gopls/internal/cache: memoize cache keys

analysisNode.cacheKey shows up as very hot in benchmarks, yet analysis
should observe the same invalidation heuristics as other package data.

Memoizing the cache keys in the snapshot reduced total CPU of
BenchmarkDiagnosePackageFiles by 10-15%.

For golang/go#53275

Change-Id: I0f60c3e78c77ac6e58b14308bb0331022babae69
Reviewed-on: https://go-review.googlesource.com/c/tools/+/622037
Auto-Submit: Robert Findley <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Alan Donovan <[email protected]>

Author: findleyr

gopls/internal/cache/analysis.go

@@ -19,6 +19,7 @@ import (
 	"go/token"
 	"go/types"
 	"log"
+	"maps"
 	urlpkg "net/url"
 	"path/filepath"
 	"reflect"
@@ -45,6 +46,7 @@ import (
 	"golang.org/x/tools/gopls/internal/util/bug"
 	"golang.org/x/tools/gopls/internal/util/frob"
 	"golang.org/x/tools/gopls/internal/util/moremaps"
+	"golang.org/x/tools/gopls/internal/util/persistent"
 	"golang.org/x/tools/internal/analysisinternal"
 	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/facts"
@@ -175,7 +177,7 @@ const AnalysisProgressTitle = "Analyzing Dependencies"
 // The analyzers list must be duplicate free; order does not matter.
 //
 // Notifications of progress may be sent to the optional reporter.
-func (s *Snapshot) Analyze(ctx context.Context, pkgs map[PackageID]*metadata.Package, analyzers []*settings.Analyzer, reporter *progress.Tracker) ([]*Diagnostic, error) {
+func (s *Snapshot) Analyze(ctx context.Context, pkgs map[PackageID]*metadata.Package, reporter *progress.Tracker) ([]*Diagnostic, error) {
 	start := time.Now() // for progress reporting
 
 	var tagStr string // sorted comma-separated list of PackageIDs
@@ -192,6 +194,7 @@ func (s *Snapshot) Analyze(ctx context.Context, pkgs map[PackageID]*metadata.Pac
 
 	// Filter and sort enabled root analyzers.
 	// A disabled analyzer may still be run if required by another.
+	analyzers := analyzers(s.Options().Staticcheck)
 	toSrc := make(map[*analysis.Analyzer]*settings.Analyzer)
 	var enabledAnalyzers []*analysis.Analyzer // enabled subset + transitive requirements
 	for _, a := range analyzers {
@@ -322,20 +325,14 @@ func (s *Snapshot) Analyze(ctx context.Context, pkgs map[PackageID]*metadata.Pac
 		roots = append(roots, root)
 	}
 
-	// Now that we have read all files,
-	// we no longer need the snapshot.
-	// (but options are needed for progress reporting)
-	options := s.Options()
-	s = nil
-
 	// Progress reporting. If supported, gopls reports progress on analysis
 	// passes that are taking a long time.
 	maybeReport := func(completed int64) {}
 
 	// Enable progress reporting if enabled by the user
 	// and we have a capable reporter.
-	if reporter != nil && reporter.SupportsWorkDoneProgress() && options.AnalysisProgressReporting {
-		var reportAfter = options.ReportAnalysisProgressAfter // tests may set this to 0
+	if reporter != nil && reporter.SupportsWorkDoneProgress() && s.Options().AnalysisProgressReporting {
+		var reportAfter = s.Options().ReportAnalysisProgressAfter // tests may set this to 0
 		const reportEvery = 1 * time.Second
 
 		ctx, cancel := context.WithCancel(ctx)
@@ -396,10 +393,34 @@ func (s *Snapshot) Analyze(ctx context.Context, pkgs map[PackageID]*metadata.Pac
 			limiter <- unit{}
 			defer func() { <-limiter }()
 
-			summary, err := an.runCached(ctx)
+			// Check to see if we already have a valid cache key. If not, compute it.
+			//
+			// The snapshot field that memoizes keys depends on whether this key is
+			// for the analysis result including all enabled analyzer, or just facty analyzers.
+			var keys *persistent.Map[PackageID, file.Hash]
+			if _, root := pkgs[an.mp.ID]; root {
+				keys = s.fullAnalysisKeys
+			} else {
+				keys = s.factyAnalysisKeys
+			}
+
+			// As keys is referenced by a snapshot field, it's guarded by s.mu.
+			s.mu.Lock()
+			key, keyFound := keys.Get(an.mp.ID)
+			s.mu.Unlock()
+
+			if !keyFound {
+				key = an.cacheKey()
+				s.mu.Lock()
+				keys.Set(an.mp.ID, key, nil)
+				s.mu.Unlock()
+			}
+
+			summary, err := an.runCached(ctx, key)
 			if err != nil {
 				return err // cancelled, or failed to produce a package
 			}
+
 			maybeReport(completed.Add(1))
 			an.summary = summary
 
@@ -488,6 +509,14 @@ func (s *Snapshot) Analyze(ctx context.Context, pkgs map[PackageID]*metadata.Pac
 	return results, nil
 }
 
+func analyzers(staticcheck bool) []*settings.Analyzer {
+	analyzers := slices.Collect(maps.Values(settings.DefaultAnalyzers))
+	if staticcheck {
+		analyzers = slices.AppendSeq(analyzers, maps.Values(settings.StaticcheckAnalyzers))
+	}
+	return analyzers
+}
+
 func (an *analysisNode) decrefPreds() {
 	if an.unfinishedPreds.Add(-1) == 0 {
 		an.summary.Actions = nil
@@ -711,21 +740,14 @@ var (
 // actions failed. It usually fails only if the package was unknown,
 // a file was missing, or the operation was cancelled.
 //
-// Postcondition: runCached must not continue to use the snapshot
-// (in background goroutines) after it has returned; see memoize.RefCounted.
-func (an *analysisNode) runCached(ctx context.Context) (*analyzeSummary, error) {
-	// At this point we have the action results (serialized
-	// packages and facts) of our immediate dependencies,
-	// and the metadata and content of this package.
-	//
-	// We now compute a hash for all our inputs, and consult a
-	// global cache of promised results. If nothing material
-	// has changed, we'll make a hit in the shared cache.
+// The provided key is the cache key for this package.
+func (an *analysisNode) runCached(ctx context.Context, key file.Hash) (*analyzeSummary, error) {
+	// At this point we have the action results (serialized packages and facts)
+	// of our immediate dependencies, and the metadata and content of this
+	// package.
 	//
-	// The hash of our inputs is based on the serialized export
-	// data and facts so that immaterial changes can be pruned
-	// without decoding.
-	key := an.cacheKey()
+	// We now consult a global cache of promised results. If nothing material has
+	// changed, we'll make a hit in the shared cache.
 
 	// Access the cache.
 	var summary *analyzeSummary

gopls/internal/cache/session.go

@@ -240,25 +240,27 @@ func (s *Session) createView(ctx context.Context, def *viewDefinition) (*View, *
 
 	s.snapshotWG.Add(1)
 	v.snapshot = &Snapshot{
-		view:             v,
-		backgroundCtx:    backgroundCtx,
-		cancel:           cancel,
-		store:            s.cache.store,
-		refcount:         1, // Snapshots are born referenced.
-		done:             s.snapshotWG.Done,
-		packages:         new(persistent.Map[PackageID, *packageHandle]),
-		meta:             new(metadata.Graph),
-		files:            newFileMap(),
-		symbolizeHandles: new(persistent.Map[protocol.DocumentURI, *memoize.Promise]),
-		shouldLoad:       new(persistent.Map[PackageID, []PackagePath]),
-		unloadableFiles:  new(persistent.Set[protocol.DocumentURI]),
-		parseModHandles:  new(persistent.Map[protocol.DocumentURI, *memoize.Promise]),
-		parseWorkHandles: new(persistent.Map[protocol.DocumentURI, *memoize.Promise]),
-		modTidyHandles:   new(persistent.Map[protocol.DocumentURI, *memoize.Promise]),
-		modVulnHandles:   new(persistent.Map[protocol.DocumentURI, *memoize.Promise]),
-		modWhyHandles:    new(persistent.Map[protocol.DocumentURI, *memoize.Promise]),
-		moduleUpgrades:   new(persistent.Map[protocol.DocumentURI, map[string]string]),
-		vulns:            new(persistent.Map[protocol.DocumentURI, *vulncheck.Result]),
+		view:              v,
+		backgroundCtx:     backgroundCtx,
+		cancel:            cancel,
+		store:             s.cache.store,
+		refcount:          1, // Snapshots are born referenced.
+		done:              s.snapshotWG.Done,
+		packages:          new(persistent.Map[PackageID, *packageHandle]),
+		fullAnalysisKeys:  new(persistent.Map[PackageID, file.Hash]),
+		factyAnalysisKeys: new(persistent.Map[PackageID, file.Hash]),
+		meta:              new(metadata.Graph),
+		files:             newFileMap(),
+		symbolizeHandles:  new(persistent.Map[protocol.DocumentURI, *memoize.Promise]),
+		shouldLoad:        new(persistent.Map[PackageID, []PackagePath]),
+		unloadableFiles:   new(persistent.Set[protocol.DocumentURI]),
+		parseModHandles:   new(persistent.Map[protocol.DocumentURI, *memoize.Promise]),
+		parseWorkHandles:  new(persistent.Map[protocol.DocumentURI, *memoize.Promise]),
+		modTidyHandles:    new(persistent.Map[protocol.DocumentURI, *memoize.Promise]),
+		modVulnHandles:    new(persistent.Map[protocol.DocumentURI, *memoize.Promise]),
+		modWhyHandles:     new(persistent.Map[protocol.DocumentURI, *memoize.Promise]),
+		moduleUpgrades:    new(persistent.Map[protocol.DocumentURI, map[string]string]),
+		vulns:             new(persistent.Map[protocol.DocumentURI, *vulncheck.Result]),
 	}
 
 	// Snapshots must observe all open files, as there are some caching

gopls/internal/cache/snapshot.go

@@ -144,6 +144,16 @@ type Snapshot struct {
 	//    be in packages, unless there is a missing import
 	packages *persistent.Map[PackageID, *packageHandle]
 
+	// fullAnalysisKeys and factyAnalysisKeys hold memoized cache keys for
+	// analysis packages. "full" refers to the cache key including all enabled
+	// analyzers, whereas "facty" is the key including only the subset of enabled
+	// analyzers that produce facts, such as is required for transitively
+	// imported packages.
+	//
+	// These keys are memoized because they can be quite expensive to compute.
+	fullAnalysisKeys  *persistent.Map[PackageID, file.Hash]
+	factyAnalysisKeys *persistent.Map[PackageID, file.Hash]
+
 	// workspacePackages contains the workspace's packages, which are loaded
 	// when the view is created. It does not contain intermediate test variants.
 	workspacePackages immutable.Map[PackageID, PackagePath]
@@ -1611,6 +1621,8 @@ func (s *Snapshot) clone(ctx, bgCtx context.Context, changed StateChange, done f
 		initialized:       s.initialized,
 		initialErr:        s.initialErr,
 		packages:          s.packages.Clone(),
+		fullAnalysisKeys:  s.fullAnalysisKeys.Clone(),
+		factyAnalysisKeys: s.factyAnalysisKeys.Clone(),
 		files:             s.files.clone(changedFiles),
 		symbolizeHandles:  cloneWithout(s.symbolizeHandles, changedFiles, nil),
 		workspacePackages: s.workspacePackages,
@@ -1888,6 +1900,12 @@ func (s *Snapshot) clone(ctx, bgCtx context.Context, changed StateChange, done f
 		// invalidation.
 		if ph, ok := result.packages.Get(id); ok {
 			needsDiagnosis = true
+
+			// Always invalidate analysis keys, as we do not implement fine-grained
+			// invalidation for analysis.
+			result.fullAnalysisKeys.Delete(id)
+			result.factyAnalysisKeys.Delete(id)
+
 			if invalidateMetadata {
 				result.packages.Delete(id)
 			} else {

gopls/internal/golang/diagnostics.go

@@ -6,14 +6,11 @@ package golang
 
 import (
 	"context"
-	"maps"
-	"slices"
 
 	"golang.org/x/tools/gopls/internal/cache"
 	"golang.org/x/tools/gopls/internal/cache/metadata"
 	"golang.org/x/tools/gopls/internal/progress"
 	"golang.org/x/tools/gopls/internal/protocol"
-	"golang.org/x/tools/gopls/internal/settings"
 	"golang.org/x/tools/gopls/internal/util/moremaps"
 )
 
@@ -35,8 +32,7 @@ func DiagnoseFile(ctx context.Context, snapshot *cache.Snapshot, uri protocol.Do
 	diags := pkgDiags[uri]
 
 	// Get analysis diagnostics.
-	analyzers := analyzers(snapshot.Options().Staticcheck)
-	pkgAnalysisDiags, err := snapshot.Analyze(ctx, map[PackageID]*metadata.Package{mp.ID: mp}, analyzers, nil)
+	pkgAnalysisDiags, err := snapshot.Analyze(ctx, map[PackageID]*metadata.Package{mp.ID: mp}, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -60,8 +56,7 @@ func Analyze(ctx context.Context, snapshot *cache.Snapshot, pkgIDs map[PackageID
 		return nil, ctx.Err()
 	}
 
-	analyzers := analyzers(snapshot.Options().Staticcheck)
-	analysisDiagnostics, err := snapshot.Analyze(ctx, pkgIDs, analyzers, tracker)
+	analysisDiagnostics, err := snapshot.Analyze(ctx, pkgIDs, tracker)
 	if err != nil {
 		return nil, err
 	}
@@ -71,14 +66,6 @@ func Analyze(ctx context.Context, snapshot *cache.Snapshot, pkgIDs map[PackageID
 // byURI is used for grouping diagnostics.
 func byURI(d *cache.Diagnostic) protocol.DocumentURI { return d.URI }
 
-func analyzers(staticcheck bool) []*settings.Analyzer {
-	analyzers := slices.Collect(maps.Values(settings.DefaultAnalyzers))
-	if staticcheck {
-		analyzers = slices.AppendSeq(analyzers, maps.Values(settings.StaticcheckAnalyzers))
-	}
-	return analyzers
-}
-
 // CombineDiagnostics combines and filters list/parse/type diagnostics from
 // tdiags with the analysis adiags, returning the resulting combined set.
 //