internal/memoize: fix race on read of handle.function

Late into CL 206879 I started nulling out a handle's function when the
handle finished running. That invalidated a previous assumption that the
field was immutable. Fix the assumption, and since the case of having
multiple computations in flight is at least a little bit possible, try
harder to avoid duplicate work.

Fixes golang/go#35995.

Change-Id: Ib5e3640f931f95e35748f28f5a82cf75585b305d
Reviewed-on: https://go-review.googlesource.com/c/tools/+/210077
Run-TryBot: Heschi Kreinick <heschi@google.com>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
diff --git a/internal/memoize/memoize.go b/internal/memoize/memoize.go
index 8660111..232692c 100644
--- a/internal/memoize/memoize.go
+++ b/internal/memoize/memoize.go
@@ -193,8 +193,14 @@
 	h.cancel = cancel
 	h.state = stateRunning
 	h.done = make(chan struct{})
+	function := h.function // Read under the lock
 	go func() {
-		v := h.function(childCtx)
+		// Just in case the function does something expensive without checking
+		// the context, double-check we're still alive.
+		if childCtx.Err() != nil {
+			return
+		}
+		v := function(childCtx)
 		if childCtx.Err() != nil {
 			return
 		}