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>

internal/lsp/cache/check.go

@@ -16,6 +16,7 @@ import (
 	"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 @@ type importer struct {
 }
 
 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 @@ func (imp *importer) getPkg(id packageID) (*pkg, error) {
 
 		// 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 @@ func (imp *importer) getPkg(id packageID) (*pkg, error) {
 	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 @@ func (imp *importer) typeCheck(id packageID) (*pkg, error) {
 	)
 	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 @@ func (imp *importer) typeCheck(id packageID) (*pkg, error) {
 		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 @@ func (imp *importer) typeCheck(id packageID) (*pkg, error) {
 		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 @@ func (imp *importer) typeCheck(id packageID) (*pkg, error) {
 	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 @@ func (imp *importer) cachePackage(ctx context.Context, pkg *pkg, meta *metadata,
 	// 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
 		}

internal/lsp/cache/load.go

@@ -43,10 +43,10 @@ func (v *view) loadParseTypecheck(ctx context.Context, f *goFile) ([]packages.Er
 		}
 		// 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
 		}

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) }

internal/lsp/protocol/protocol.go

@@ -8,6 +8,7 @@ import (
 	"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 defaultMessageBufferSize = 20
 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) {

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 @@ func (nullSpan) SetStatus(status Status)
 
 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{}
 	}

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) }

internal/memoize/memoize.go

@@ -46,6 +46,7 @@ type Handle struct {
 // 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 @@ func (s *Store) Bind(key interface{}, function Function) *Handle {
 		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 @@ func (e *entry) get(ctx context.Context, f Function) (interface{}, bool) {
 			// 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(detatchContext(ctx))
-			value = f(ctx)
 
 			// The function has completed. Update the value in the entry.
 			e.mu.Lock()