From 88325aa063540896255ba32a031e726c1fb0e545 Mon Sep 17 00:00:00 2001
From: Alan Donovan <adonovan@google.com>
Date: Thu, 16 Jun 2022 15:41:24 -0400
Subject: [PATCH] internal/memoize: add trace region for Handle.run

The region name includes the type of the key, such as packageHandleKey
or actionHandleKey, so we can separate these deferred computations
into their own buckets.

Change-Id: I0359127ccf47b158f353fae2bf74ba000668a40b
Reviewed-on: https://go-review.googlesource.com/c/tools/+/412817
Run-TryBot: Alan Donovan <adonovan@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>
---
 internal/memoize/memoize.go | 65 +++++++++++++++++++------------------
 1 file changed, 34 insertions(+), 31 deletions(-)

diff --git a/internal/memoize/memoize.go b/internal/memoize/memoize.go
index dec2fff683..480b87f5ce 100644
--- a/internal/memoize/memoize.go
+++ b/internal/memoize/memoize.go
@@ -14,6 +14,7 @@ import (
 	"flag"
 	"fmt"
 	"reflect"
+	"runtime/trace"
 	"sync"
 	"sync/atomic"
 
@@ -316,40 +317,42 @@ func (h *Handle) run(ctx context.Context, g *Generation, arg Arg) (interface{},
 	// Make sure that the generation isn't destroyed while we're running in it.
 	release := g.Acquire()
 	go func() {
-		defer release()
-		// 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, arg)
-		if childCtx.Err() != nil {
-			// It's possible that v was computed despite the context cancellation. In
-			// this case we should ensure that it is cleaned up.
-			if h.cleanup != nil && v != nil {
-				h.cleanup(v)
+		trace.WithRegion(childCtx, fmt.Sprintf("Handle.run %T", h.key), func() {
+			defer release()
+			// 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, arg)
+			if childCtx.Err() != nil {
+				// It's possible that v was computed despite the context cancellation. In
+				// this case we should ensure that it is cleaned up.
+				if h.cleanup != nil && v != nil {
+					h.cleanup(v)
+				}
+				return
 			}
-			return
-		}
 
-		h.mu.Lock()
-		defer h.mu.Unlock()
-		// It's theoretically possible that the handle has been cancelled out
-		// of the run that started us, and then started running again since we
-		// checked childCtx above. Even so, that should be harmless, since each
-		// run should produce the same results.
-		if h.state != stateRunning {
-			// v will never be used, so ensure that it is cleaned up.
-			if h.cleanup != nil && v != nil {
-				h.cleanup(v)
+			h.mu.Lock()
+			defer h.mu.Unlock()
+			// It's theoretically possible that the handle has been cancelled out
+			// of the run that started us, and then started running again since we
+			// checked childCtx above. Even so, that should be harmless, since each
+			// run should produce the same results.
+			if h.state != stateRunning {
+				// v will never be used, so ensure that it is cleaned up.
+				if h.cleanup != nil && v != nil {
+					h.cleanup(v)
+				}
+				return
 			}
-			return
-		}
-		// At this point v will be cleaned up whenever h is destroyed.
-		h.value = v
-		h.function = nil
-		h.state = stateCompleted
-		close(h.done)
+			// At this point v will be cleaned up whenever h is destroyed.
+			h.value = v
+			h.function = nil
+			h.state = stateCompleted
+			close(h.done)
+		})
 	}()
 
 	return h.wait(ctx)