gopls/internal/lsp: add tracing instrumentation for all Server methods

While debugging the gopls/internal/lsp tests, I noticed that I was not
seeing useful user defined tasks in the trace viewer. The reasons for
this were that (1) we were disabling all instrumentation of the event
package, and (2) we didn't have top-level tasks for each jsonrcp2
method, because the tests call into the Server directly.

Fix this by re-enabling the exporter in gopls/internal/lsp tests
(using a debug.Instance to suppress logging to stderr), and by
instrumenting all server methods with a top-level span. Some were
already instrumented, but it was inconsistent. Another advantage to
instrumenting server methods in addition to jsonrpc2 methods, is that
they can add labels specific to the request, such as the text document
URI.

Also add some instrumentation of key snapshot methods.

Change-Id: I992a6a86b175b766e6cbff8c2f2c4a5a35b5d0cf
Reviewed-on: https://go-review.googlesource.com/c/tools/+/480198
Reviewed-by: Peter Weinberger <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>
Run-TryBot: Robert Findley <[email protected]>

gopls/internal/lsp/cache/check.go

@@ -188,6 +188,9 @@ func (s *snapshot) getImportGraph(ctx context.Context) *importGraph {
 // resolveImportGraph should only be called from getImportGraph.
 func (s *snapshot) resolveImportGraph() (*importGraph, error) {
 	ctx := s.backgroundCtx
+	ctx, done := event.Start(event.Detach(ctx), "cache.resolveImportGraph")
+	defer done()
+
 	if err := s.awaitLoaded(ctx); err != nil {
 		return nil, err
 	}
@@ -334,6 +337,9 @@ type (
 //
 // Both pre and post may be called concurrently.
 func (s *snapshot) forEachPackage(ctx context.Context, ids []PackageID, pre preTypeCheck, post postTypeCheck) error {
+	ctx, done := event.Start(ctx, "cache.forEachPackage", tag.PackageCount.Of(len(ids)))
+	defer done()
+
 	if len(ids) == 0 {
 		return nil // short cut: many call sites do not handle empty ids
 	}
@@ -534,6 +540,9 @@ func (b *typeCheckBatch) handleSyntaxPackage(ctx context.Context, i int, id Pack
 // importPackage loads the given package from its export data in p.exportData
 // (which must already be populated).
 func (b *typeCheckBatch) importPackage(ctx context.Context, m *source.Metadata, data []byte) (*types.Package, error) {
+	ctx, done := event.Start(ctx, "cache.typeCheckBatch.importPackage", tag.Package.Of(string(m.ID)))
+	defer done()
+
 	impMap := b.importMap(m.ID)
 
 	var firstErr error // TODO(rfindley): unused: revisit or remove.
@@ -580,6 +589,9 @@ func (b *typeCheckBatch) importPackage(ctx context.Context, m *source.Metadata,
 // checkPackageForImport type checks, but skips function bodies and does not
 // record syntax information.
 func (b *typeCheckBatch) checkPackageForImport(ctx context.Context, ph *packageHandle) (*types.Package, error) {
+	ctx, done := event.Start(ctx, "cache.typeCheckBatch.checkPackageForImport", tag.Package.Of(string(ph.m.ID)))
+	defer done()
+
 	onError := func(e error) {
 		// Ignore errors for exporting.
 	}
@@ -627,6 +639,9 @@ func (b *typeCheckBatch) checkPackageForImport(ctx context.Context, ph *packageH
 
 // checkPackage "fully type checks" to produce a syntax package.
 func (b *typeCheckBatch) checkPackage(ctx context.Context, ph *packageHandle) (*Package, error) {
+	ctx, done := event.Start(ctx, "cache.typeCheckBatch.checkPackage", tag.Package.Of(string(ph.m.ID)))
+	defer done()
+
 	// TODO(rfindley): refactor to inline typeCheckImpl here. There is no need
 	// for so many layers to build up the package
 	// (checkPackage->typeCheckImpl->doTypeCheck).

gopls/internal/lsp/cache/load.go

@@ -103,7 +103,7 @@ func (s *snapshot) load(ctx context.Context, allowNetwork bool, scopes ...loadSc
 	}
 	sort.Strings(query) // for determinism
 
-	ctx, done := event.Start(ctx, "cache.view.load", tag.Query.Of(query))
+	ctx, done := event.Start(ctx, "cache.snapshot.load", tag.Query.Of(query))
 	defer done()
 
 	flags := source.LoadWorkspace

gopls/internal/lsp/cache/mod_tidy.go

@@ -29,6 +29,9 @@ import (
 // ModTidy returns the go.mod file that would be obtained by running
 // "go mod tidy". Concurrent requests are combined into a single command.
 func (s *snapshot) ModTidy(ctx context.Context, pm *source.ParsedModule) (*source.TidiedModule, error) {
+	ctx, done := event.Start(ctx, "cache.snapshot.ModTidy")
+	defer done()
+
 	uri := pm.URI
 	if pm.File == nil {
 		return nil, fmt.Errorf("cannot tidy unparseable go.mod file: %v", uri)

gopls/internal/lsp/cache/parse.go

@@ -36,9 +36,6 @@ func (s *snapshot) ParseGo(ctx context.Context, fh source.FileHandle, mode parse
 
 // parseGoImpl parses the Go source file whose content is provided by fh.
 func parseGoImpl(ctx context.Context, fset *token.FileSet, fh source.FileHandle, mode parser.Mode) (*source.ParsedGoFile, error) {
-	ctx, done := event.Start(ctx, "cache.parseGo", tag.File.Of(fh.URI().Filename()))
-	defer done()
-
 	ext := filepath.Ext(fh.URI().Filename())
 	if ext != ".go" && ext != "" { // files generated by cgo have no extension
 		return nil, fmt.Errorf("cannot parse non-Go file %s", fh.URI())
@@ -59,6 +56,9 @@ func parseGoImpl(ctx context.Context, fset *token.FileSet, fh source.FileHandle,
 //
 // The provided ctx is used only for logging.
 func ParseGoSrc(ctx context.Context, fset *token.FileSet, uri span.URI, src []byte, mode parser.Mode) (res *source.ParsedGoFile, fixes []fixType) {
+	ctx, done := event.Start(ctx, "cache.ParseGoSrc", tag.File.Of(uri.Filename()))
+	defer done()
+
 	file, err := parser.ParseFile(fset, uri.Filename(), src, mode)
 	var parseErr scanner.ErrorList
 	if err != nil {

gopls/internal/lsp/cache/parse_cache.go

@@ -161,7 +161,7 @@ func (c *parseCache) startParse(mode parser.Mode, fhs ...source.FileHandle) ([]*
 		}
 
 		uri := fh.URI()
-		promise := memoize.NewPromise(fmt.Sprintf("parse(%s)", uri), func(ctx context.Context, _ interface{}) interface{} {
+		promise := memoize.NewPromise("parseCache.parse", func(ctx context.Context, _ interface{}) interface{} {
 			// Allocate 2*len(content)+parsePadding to allow for re-parsing once
 			// inside of parseGoSrc without exceeding the allocated space.
 			base, nextBase := c.allocateSpace(2*len(content) + parsePadding)

gopls/internal/lsp/cache/snapshot.go

@@ -657,6 +657,9 @@ const (
 )
 
 func (s *snapshot) PackageDiagnostics(ctx context.Context, ids ...PackageID) (map[span.URI][]*source.Diagnostic, error) {
+	ctx, done := event.Start(ctx, "cache.snapshot.PackageDiagnostics")
+	defer done()
+
 	var mu sync.Mutex
 	perFile := make(map[span.URI][]*source.Diagnostic)
 	collect := func(diags []*source.Diagnostic) {
@@ -685,6 +688,9 @@ func (s *snapshot) PackageDiagnostics(ctx context.Context, ids ...PackageID) (ma
 }
 
 func (s *snapshot) References(ctx context.Context, ids ...PackageID) ([]source.XrefIndex, error) {
+	ctx, done := event.Start(ctx, "cache.snapshot.References")
+	defer done()
+
 	indexes := make([]source.XrefIndex, len(ids))
 	pre := func(i int, ph *packageHandle) bool {
 		data, err := filecache.Get(xrefsKind, ph.key)
@@ -713,6 +719,9 @@ func (index XrefIndex) Lookup(targets map[PackagePath]map[objectpath.Path]struct
 }
 
 func (s *snapshot) MethodSets(ctx context.Context, ids ...PackageID) ([]*methodsets.Index, error) {
+	ctx, done := event.Start(ctx, "cache.snapshot.MethodSets")
+	defer done()
+
 	indexes := make([]*methodsets.Index, len(ids))
 	pre := func(i int, ph *packageHandle) bool {
 		data, err := filecache.Get(methodSetsKind, ph.key)
@@ -731,6 +740,9 @@ func (s *snapshot) MethodSets(ctx context.Context, ids ...PackageID) ([]*methods
 }
 
 func (s *snapshot) MetadataForFile(ctx context.Context, uri span.URI) ([]*source.Metadata, error) {
+	ctx, done := event.Start(ctx, "cache.snapshot.MetadataForFile")
+	defer done()
+
 	s.mu.Lock()
 
 	// Start with the set of package associations derived from the last load.
@@ -1689,7 +1701,7 @@ func (ac *unappliedChanges) ReadFile(ctx context.Context, uri span.URI) (source.
 }
 
 func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileChange, forceReloadMetadata bool) (*snapshot, func()) {
-	ctx, done := event.Start(ctx, "snapshot.clone")
+	ctx, done := event.Start(ctx, "cache.snapshot.clone")
 	defer done()
 
 	reinit := false

gopls/internal/lsp/call_hierarchy.go

@@ -9,9 +9,13 @@ import (
 
 	"golang.org/x/tools/gopls/internal/lsp/protocol"
 	"golang.org/x/tools/gopls/internal/lsp/source"
+	"golang.org/x/tools/internal/event"
 )
 
 func (s *Server) prepareCallHierarchy(ctx context.Context, params *protocol.CallHierarchyPrepareParams) ([]protocol.CallHierarchyItem, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.prepareCallHierarchy")
+	defer done()
+
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.Go)
 	defer release()
 	if !ok {
@@ -22,6 +26,9 @@ func (s *Server) prepareCallHierarchy(ctx context.Context, params *protocol.Call
 }
 
 func (s *Server) incomingCalls(ctx context.Context, params *protocol.CallHierarchyIncomingCallsParams) ([]protocol.CallHierarchyIncomingCall, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.incomingCalls")
+	defer done()
+
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.Item.URI, source.Go)
 	defer release()
 	if !ok {
@@ -32,6 +39,9 @@ func (s *Server) incomingCalls(ctx context.Context, params *protocol.CallHierarc
 }
 
 func (s *Server) outgoingCalls(ctx context.Context, params *protocol.CallHierarchyOutgoingCallsParams) ([]protocol.CallHierarchyOutgoingCall, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.outgoingCalls")
+	defer done()
+
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.Item.URI, source.Go)
 	defer release()
 	if !ok {

gopls/internal/lsp/code_action.go

@@ -21,6 +21,9 @@ import (
 )
 
 func (s *Server) codeAction(ctx context.Context, params *protocol.CodeActionParams) ([]protocol.CodeAction, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.codeAction")
+	defer done()
+
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.UnknownKind)
 	defer release()
 	if !ok {

gopls/internal/lsp/code_lens.go

@@ -14,9 +14,13 @@ import (
 	"golang.org/x/tools/gopls/internal/lsp/protocol"
 	"golang.org/x/tools/gopls/internal/lsp/source"
 	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/event/tag"
 )
 
 func (s *Server) codeLens(ctx context.Context, params *protocol.CodeLensParams) ([]protocol.CodeLens, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.codeLens", tag.URI.Of(params.TextDocument.URI))
+	defer done()
+
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.UnknownKind)
 	defer release()
 	if !ok {

gopls/internal/lsp/command.go

@@ -37,6 +37,9 @@ import (
 )
 
 func (s *Server) executeCommand(ctx context.Context, params *protocol.ExecuteCommandParams) (interface{}, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.executeCommand")
+	defer done()
+
 	var found bool
 	for _, name := range s.session.Options().SupportedCommands {
 		if name == params.Command {

gopls/internal/lsp/completion.go

@@ -19,6 +19,9 @@ import (
 )
 
 func (s *Server) completion(ctx context.Context, params *protocol.CompletionParams) (*protocol.CompletionList, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.completion", tag.URI.Of(params.TextDocument.URI))
+	defer done()
+
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.UnknownKind)
 	defer release()
 	if !ok {

gopls/internal/lsp/definition.go

@@ -12,9 +12,14 @@ import (
 	"golang.org/x/tools/gopls/internal/lsp/protocol"
 	"golang.org/x/tools/gopls/internal/lsp/source"
 	"golang.org/x/tools/gopls/internal/lsp/template"
+	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/event/tag"
 )
 
 func (s *Server) definition(ctx context.Context, params *protocol.DefinitionParams) ([]protocol.Location, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.definition", tag.URI.Of(params.TextDocument.URI))
+	defer done()
+
 	// TODO(rfindley): definition requests should be multiplexed across all views.
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.UnknownKind)
 	defer release()
@@ -37,6 +42,9 @@ func (s *Server) definition(ctx context.Context, params *protocol.DefinitionPara
 }
 
 func (s *Server) typeDefinition(ctx context.Context, params *protocol.TypeDefinitionParams) ([]protocol.Location, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.typeDefinition", tag.URI.Of(params.TextDocument.URI))
+	defer done()
+
 	// TODO(rfindley): type definition requests should be multiplexed across all views.
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.Go)
 	defer release()

gopls/internal/lsp/folding_range.go

@@ -9,9 +9,14 @@ import (
 
 	"golang.org/x/tools/gopls/internal/lsp/protocol"
 	"golang.org/x/tools/gopls/internal/lsp/source"
+	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/event/tag"
 )
 
 func (s *Server) foldingRange(ctx context.Context, params *protocol.FoldingRangeParams) ([]protocol.FoldingRange, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.foldingRange", tag.URI.Of(params.TextDocument.URI))
+	defer done()
+
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.Go)
 	defer release()
 	if !ok {

gopls/internal/lsp/format.go

@@ -11,9 +11,14 @@ import (
 	"golang.org/x/tools/gopls/internal/lsp/protocol"
 	"golang.org/x/tools/gopls/internal/lsp/source"
 	"golang.org/x/tools/gopls/internal/lsp/work"
+	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/event/tag"
 )
 
 func (s *Server) formatting(ctx context.Context, params *protocol.DocumentFormattingParams) ([]protocol.TextEdit, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.formatting", tag.URI.Of(params.TextDocument.URI))
+	defer done()
+
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.UnknownKind)
 	defer release()
 	if !ok {

gopls/internal/lsp/general.go

@@ -26,6 +26,9 @@ import (
 )
 
 func (s *Server) initialize(ctx context.Context, params *protocol.ParamInitialize) (*protocol.InitializeResult, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.initialize")
+	defer done()
+
 	s.stateMu.Lock()
 	if s.state >= serverInitializing {
 		defer s.stateMu.Unlock()
@@ -192,6 +195,9 @@ See https://github.com/golang/go/issues/45732 for more information.`,
 }
 
 func (s *Server) initialized(ctx context.Context, params *protocol.InitializedParams) error {
+	ctx, done := event.Start(ctx, "lsp.Server.initialized")
+	defer done()
+
 	s.stateMu.Lock()
 	if s.state >= serverInitialized {
 		defer s.stateMu.Unlock()
@@ -585,6 +591,9 @@ func (s *Server) beginFileRequest(ctx context.Context, pURI protocol.DocumentURI
 // shutdown implements the 'shutdown' LSP handler. It releases resources
 // associated with the server and waits for all ongoing work to complete.
 func (s *Server) shutdown(ctx context.Context) error {
+	ctx, done := event.Start(ctx, "lsp.Server.shutdown")
+	defer done()
+
 	s.stateMu.Lock()
 	defer s.stateMu.Unlock()
 	if s.state < serverInitialized {
@@ -604,6 +613,9 @@ func (s *Server) shutdown(ctx context.Context) error {
 }
 
 func (s *Server) exit(ctx context.Context) error {
+	ctx, done := event.Start(ctx, "lsp.Server.exit")
+	defer done()
+
 	s.stateMu.Lock()
 	defer s.stateMu.Unlock()
 

gopls/internal/lsp/highlight.go

@@ -7,14 +7,17 @@ package lsp
 import (
 	"context"
 
-	"golang.org/x/tools/internal/event"
-	"golang.org/x/tools/internal/event/tag"
 	"golang.org/x/tools/gopls/internal/lsp/protocol"
 	"golang.org/x/tools/gopls/internal/lsp/source"
 	"golang.org/x/tools/gopls/internal/lsp/template"
+	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/event/tag"
 )
 
 func (s *Server) documentHighlight(ctx context.Context, params *protocol.DocumentHighlightParams) ([]protocol.DocumentHighlight, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.documentHighlight", tag.URI.Of(params.TextDocument.URI))
+	defer done()
+
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.Go)
 	defer release()
 	if !ok {
@@ -27,7 +30,7 @@ func (s *Server) documentHighlight(ctx context.Context, params *protocol.Documen
 
 	rngs, err := source.Highlight(ctx, snapshot, fh, params.Position)
 	if err != nil {
-		event.Error(ctx, "no highlight", err, tag.URI.Of(params.TextDocument.URI))
+		event.Error(ctx, "no highlight", err)
 	}
 	return toProtocolHighlight(rngs), nil
 }

gopls/internal/lsp/hover.go

@@ -12,9 +12,14 @@ import (
 	"golang.org/x/tools/gopls/internal/lsp/source"
 	"golang.org/x/tools/gopls/internal/lsp/template"
 	"golang.org/x/tools/gopls/internal/lsp/work"
+	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/event/tag"
 )
 
 func (s *Server) hover(ctx context.Context, params *protocol.HoverParams) (*protocol.Hover, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.hover", tag.URI.Of(params.TextDocument.URI))
+	defer done()
+
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.UnknownKind)
 	defer release()
 	if !ok {

gopls/internal/lsp/implementation.go

@@ -9,9 +9,14 @@ import (
 
 	"golang.org/x/tools/gopls/internal/lsp/protocol"
 	"golang.org/x/tools/gopls/internal/lsp/source"
+	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/event/tag"
 )
 
 func (s *Server) implementation(ctx context.Context, params *protocol.ImplementationParams) ([]protocol.Location, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.implementation", tag.URI.Of(params.TextDocument.URI))
+	defer done()
+
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.Go)
 	defer release()
 	if !ok {

gopls/internal/lsp/inlay_hint.go

@@ -10,9 +10,14 @@ import (
 	"golang.org/x/tools/gopls/internal/lsp/mod"
 	"golang.org/x/tools/gopls/internal/lsp/protocol"
 	"golang.org/x/tools/gopls/internal/lsp/source"
+	"golang.org/x/tools/internal/event"
+	"golang.org/x/tools/internal/event/tag"
 )
 
 func (s *Server) inlayHint(ctx context.Context, params *protocol.InlayHintParams) ([]protocol.InlayHint, error) {
+	ctx, done := event.Start(ctx, "lsp.Server.inlayHint", tag.URI.Of(params.TextDocument.URI))
+	defer done()
+
 	snapshot, fh, ok, release, err := s.beginFileRequest(ctx, params.TextDocument.URI, source.UnknownKind)
 	defer release()
 	if !ok {