internal/lsp: connect memoize actions to their callers

This adds the ability to tie a background context to the context that created it
in traces, and also cleans up and annotates the context used in type checking.
This gives us detailed connected traces of all the type checking and parsing
logic.

Change-Id: I32721220a50ecb9b4404a4e9354343389d7a5219
Reviewed-on: https://go-review.googlesource.com/c/tools/+/183757
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/internal/lsp/cache/check.go b/internal/lsp/cache/check.go
index 1ea134f..1a0d754 100644
--- a/internal/lsp/cache/check.go
+++ b/internal/lsp/cache/check.go
@@ -16,6 +16,7 @@
 	"golang.org/x/tools/go/analysis"
 	"golang.org/x/tools/go/packages"
 	"golang.org/x/tools/internal/lsp/source"
+	"golang.org/x/tools/internal/lsp/telemetry/trace"
 	"golang.org/x/tools/internal/span"
 )
 
@@ -34,18 +35,19 @@
 }
 
 func (imp *importer) Import(pkgPath string) (*types.Package, error) {
+	ctx := imp.ctx
 	id, ok := imp.view.mcache.ids[packagePath(pkgPath)]
 	if !ok {
 		return nil, fmt.Errorf("no known ID for %s", pkgPath)
 	}
-	pkg, err := imp.getPkg(id)
+	pkg, err := imp.getPkg(ctx, id)
 	if err != nil {
 		return nil, err
 	}
 	return pkg.types, nil
 }
 
-func (imp *importer) getPkg(id packageID) (*pkg, error) {
+func (imp *importer) getPkg(ctx context.Context, id packageID) (*pkg, error) {
 	if _, ok := imp.seen[id]; ok {
 		return nil, fmt.Errorf("circular import detected")
 	}
@@ -65,20 +67,20 @@
 
 		// This goroutine becomes responsible for populating
 		// the entry and broadcasting its readiness.
-		e.pkg, e.err = imp.typeCheck(id)
+		e.pkg, e.err = imp.typeCheck(ctx, id)
 		close(e.ready)
 	}
 
 	if e.err != nil {
 		// If the import had been previously canceled, and that error cached, try again.
-		if e.err == context.Canceled && imp.ctx.Err() == nil {
+		if e.err == context.Canceled && ctx.Err() == nil {
 			imp.view.pcache.mu.Lock()
 			// Clear out canceled cache entry if it is still there.
 			if imp.view.pcache.packages[id] == e {
 				delete(imp.view.pcache.packages, id)
 			}
 			imp.view.pcache.mu.Unlock()
-			return imp.getPkg(id)
+			return imp.getPkg(ctx, id)
 		}
 		return nil, e.err
 	}
@@ -86,7 +88,9 @@
 	return e.pkg, nil
 }
 
-func (imp *importer) typeCheck(id packageID) (*pkg, error) {
+func (imp *importer) typeCheck(ctx context.Context, id packageID) (*pkg, error) {
+	ctx, ts := trace.StartSpan(ctx, "cache.importer.typeCheck")
+	defer ts.End()
 	meta, ok := imp.view.mcache.packages[id]
 	if !ok {
 		return nil, fmt.Errorf("no metadata for %v", id)
@@ -119,11 +123,11 @@
 	)
 	for _, filename := range meta.files {
 		uri := span.FileURI(filename)
-		f, err := imp.view.getFile(imp.ctx, uri)
+		f, err := imp.view.getFile(ctx, uri)
 		if err != nil {
 			continue
 		}
-		ph := imp.view.session.cache.ParseGoHandle(f.Handle(imp.ctx), mode)
+		ph := imp.view.session.cache.ParseGoHandle(f.Handle(ctx), mode)
 		phs = append(phs, ph)
 		files = append(files, &astFile{
 			uri:       ph.File().Identity().URI,
@@ -136,7 +140,7 @@
 		go func(i int, ph source.ParseGoHandle) {
 			defer wg.Done()
 
-			files[i].file, files[i].err = ph.Parse(imp.ctx)
+			files[i].file, files[i].err = ph.Parse(ctx)
 		}(i, ph)
 	}
 	wg.Wait()
@@ -175,7 +179,7 @@
 		IgnoreFuncBodies: mode == source.ParseExported,
 		Importer: &importer{
 			view:          imp.view,
-			ctx:           imp.ctx,
+			ctx:           ctx,
 			fset:          imp.fset,
 			topLevelPkgID: imp.topLevelPkgID,
 			seen:          seen,
@@ -187,7 +191,7 @@
 	check.Files(pkg.GetSyntax())
 
 	// Add every file in this package to our cache.
-	if err := imp.cachePackage(imp.ctx, pkg, meta, mode); err != nil {
+	if err := imp.cachePackage(ctx, pkg, meta, mode); err != nil {
 		return nil, err
 	}
 
@@ -213,7 +217,7 @@
 	// We lock the package cache, but we shouldn't get any inconsistencies
 	// because we are still holding the lock on the view.
 	for importPath := range meta.children {
-		importPkg, err := imp.getPkg(importPath)
+		importPkg, err := imp.getPkg(ctx, importPath)
 		if err != nil {
 			continue
 		}
diff --git a/internal/lsp/cache/load.go b/internal/lsp/cache/load.go
index 2b6661e..84e9f7e 100644
--- a/internal/lsp/cache/load.go
+++ b/internal/lsp/cache/load.go
@@ -43,10 +43,10 @@
 		}
 		// Start prefetching direct imports.
 		for importID := range m.children {
-			go imp.getPkg(importID)
+			go imp.getPkg(ctx, importID)
 		}
 		// Type-check package.
-		pkg, err := imp.getPkg(imp.topLevelPkgID)
+		pkg, err := imp.getPkg(ctx, imp.topLevelPkgID)
 		if err != nil {
 			return nil, err
 		}
diff --git a/internal/lsp/protocol/detatch.go b/internal/lsp/protocol/detatch.go
new file mode 100644
index 0000000..9e34b09
--- /dev/null
+++ b/internal/lsp/protocol/detatch.go
@@ -0,0 +1,21 @@
+// Copyright 2019 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package protocol
+
+import (
+	"context"
+	"time"
+)
+
+// detatch returns a context that keeps all the values of its parent context
+// but detatches from the cancellation and error handling.
+func detatchContext(ctx context.Context) context.Context { return detatchedContext{ctx} }
+
+type detatchedContext struct{ parent context.Context }
+
+func (v detatchedContext) Deadline() (time.Time, bool)       { return time.Time{}, false }
+func (v detatchedContext) Done() <-chan struct{}             { return nil }
+func (v detatchedContext) Err() error                        { return nil }
+func (v detatchedContext) Value(key interface{}) interface{} { return v.parent.Value(key) }
diff --git a/internal/lsp/protocol/protocol.go b/internal/lsp/protocol/protocol.go
index a956fe4..2802c3f 100644
--- a/internal/lsp/protocol/protocol.go
+++ b/internal/lsp/protocol/protocol.go
@@ -8,6 +8,7 @@
 	"context"
 
 	"golang.org/x/tools/internal/jsonrpc2"
+	"golang.org/x/tools/internal/lsp/telemetry/trace"
 	"golang.org/x/tools/internal/lsp/xlog"
 )
 
@@ -15,7 +16,10 @@
 const defaultRejectIfOverloaded = false
 
 func canceller(ctx context.Context, conn *jsonrpc2.Conn, id jsonrpc2.ID) {
-	conn.Notify(context.Background(), "$/cancelRequest", &CancelParams{ID: id})
+	ctx = detatchContext(ctx)
+	ctx, span := trace.StartSpan(ctx, "protocol.canceller")
+	defer span.End()
+	conn.Notify(ctx, "$/cancelRequest", &CancelParams{ID: id})
 }
 
 func NewClient(stream jsonrpc2.Stream, client Client) (*jsonrpc2.Conn, Server, xlog.Logger) {
diff --git a/internal/lsp/telemetry/trace/trace.go b/internal/lsp/telemetry/trace/trace.go
index 96760b5..9fc62df 100644
--- a/internal/lsp/telemetry/trace/trace.go
+++ b/internal/lsp/telemetry/trace/trace.go
@@ -5,11 +5,12 @@
 // Package tag adds support for telemetry tracins.
 package trace
 
-import "context"
+import (
+	"context"
+)
 
 type Span interface {
 	AddAttributes(attributes ...Attribute)
-
 	AddMessageReceiveEvent(messageID, uncompressedByteSize, compressedByteSize int64)
 	AddMessageSendEvent(messageID, uncompressedByteSize, compressedByteSize int64)
 	Annotate(attributes []Attribute, str string)
@@ -41,6 +42,7 @@
 
 var (
 	FromContext = func(ctx context.Context) Span { return nullSpan{} }
+	NewContext  = func(ctx context.Context, span Span) context.Context { return ctx }
 	StartSpan   = func(ctx context.Context, name string, options ...interface{}) (context.Context, Span) {
 		return ctx, nullSpan{}
 	}
diff --git a/internal/memoize/detatch.go b/internal/memoize/detatch.go
new file mode 100644
index 0000000..6112de6
--- /dev/null
+++ b/internal/memoize/detatch.go
@@ -0,0 +1,21 @@
+// Copyright 2019 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package memoize
+
+import (
+	"context"
+	"time"
+)
+
+// detatch returns a context that keeps all the values of its parent context
+// but detatches from the cancellation and error handling.
+func detatchContext(ctx context.Context) context.Context { return detatchedContext{ctx} }
+
+type detatchedContext struct{ parent context.Context }
+
+func (v detatchedContext) Deadline() (time.Time, bool)       { return time.Time{}, false }
+func (v detatchedContext) Done() <-chan struct{}             { return nil }
+func (v detatchedContext) Err() error                        { return nil }
+func (v detatchedContext) Value(key interface{}) interface{} { return v.parent.Value(key) }
diff --git a/internal/memoize/memoize.go b/internal/memoize/memoize.go
index 6f9c621..9e0cb4c 100644
--- a/internal/memoize/memoize.go
+++ b/internal/memoize/memoize.go
@@ -46,6 +46,7 @@
 // there is only one instance of the result live at any given time.
 type entry struct {
 	noCopy
+	key interface{}
 	// mu contols access to the typ and ptr fields
 	mu sync.Mutex
 	// the calculated value, as stored in an interface{}
@@ -93,7 +94,7 @@
 		if s.entries == nil {
 			s.entries = make(map[interface{}]*entry)
 		}
-		e = &entry{}
+		e = &entry{key: key}
 		s.entries[key] = e
 	}
 	return &Handle{
@@ -179,8 +180,7 @@
 			// Use the background context to avoid canceling the function.
 			// The function cannot be canceled even if the context is canceled
 			// because multiple goroutines may depend on it.
-			ctx := context.Background()
-			value = f(ctx)
+			value = f(detatchContext(ctx))
 
 			// The function has completed. Update the value in the entry.
 			e.mu.Lock()