internal/lsp/cache: use GetHandle not Bind for actions

This change uses a persistent.Map for actions, just like packages.
Actions are now reference counted rather than generational.

Also:
- note optimization opportunities.
- minor cleanups.

Change-Id: Ibbac8848a3beb3fe19056a7b160d2185155e7021
Reviewed-on: https://go-review.googlesource.com/c/tools/+/415504
gopls-CI: kokoro <noreply+kokoro@google.com>
Auto-Submit: Alan Donovan <adonovan@google.com>
Run-TryBot: Alan Donovan <adonovan@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>

4 files changed, 75 insertions, 36 deletions

diff --git a/internal/lsp/cache/analysis.go b/internal/lsp/cache/analysis.go
index db2ca2a8b..4b437858e 100644
--- a/internal/lsp/cache/analysis.go
+++ b/internal/lsp/cache/analysis.go
@@ -85,12 +85,21 @@ type packageFactKey struct {
 }
 
 func (s *snapshot) actionHandle(ctx context.Context, id PackageID, a *analysis.Analyzer) (*actionHandle, error) {
+	// TODO(adonovan): opt: this block of code sequentially loads a package
+	// (and all its dependencies), then sequentially creates action handles
+	// for the direct dependencies (whose packages have by then been loaded
+	// as a consequence of ph.check) which does a sequential recursion
+	// down the action graph. Only once all that work is complete do we
+	// put a handle in the cache. As with buildPackageHandle, this does
+	// not exploit the natural parallelism in the problem, and the naive
+	// use of concurrency would lead to an exponential amount of duplicated
+	// work. We should instead use an atomically updated future cache
+	// and a parallel graph traversal.
 	ph, err := s.buildPackageHandle(ctx, id, source.ParseFull)
 	if err != nil {
 		return nil, err
 	}
-	act := s.getActionHandle(id, ph.mode, a)
-	if act != nil {
+	if act := s.getActionHandle(id, ph.mode, a); act != nil {
 		return act, nil
 	}
 	if len(ph.key) == 0 {
@@ -100,12 +109,9 @@ func (s *snapshot) actionHandle(ctx context.Context, id PackageID, a *analysis.A
 	if err != nil {
 		return nil, err
 	}
-	act = &actionHandle{
-		analyzer: a,
-		pkg:      pkg,
-	}
+
+	// Add a dependency on each required analyzer.
 	var deps []*actionHandle
-	// Add a dependency on each required analyzers.
 	for _, req := range a.Requires {
 		reqActionHandle, err := s.actionHandle(ctx, id, req)
 		if err != nil {
@@ -131,7 +137,7 @@ func (s *snapshot) actionHandle(ctx context.Context, id PackageID, a *analysis.A
 		}
 	}
 
-	h := s.generation.Bind(buildActionKey(a, ph), func(ctx context.Context, arg memoize.Arg) interface{} {
+	handle, release := s.generation.GetHandle(buildActionKey(a, ph), func(ctx context.Context, arg memoize.Arg) interface{} {
 		snapshot := arg.(*snapshot)
 		// Analyze dependencies first.
 		results, err := execAll(ctx, snapshot, deps)
@@ -142,9 +148,13 @@ func (s *snapshot) actionHandle(ctx context.Context, id PackageID, a *analysis.A
 		}
 		return runAnalysis(ctx, snapshot, a, pkg, results)
 	})
-	act.handle = h
 
-	act = s.addActionHandle(act)
+	act := &actionHandle{
+		analyzer: a,
+		pkg:      pkg,
+		handle:   handle,
+	}
+	act = s.addActionHandle(act, release)
 	return act, nil
 }
 
diff --git a/internal/lsp/cache/maps.go b/internal/lsp/cache/maps.go
index af0417771..f8e03057c 100644
--- a/internal/lsp/cache/maps.go
+++ b/internal/lsp/cache/maps.go
@@ -197,16 +197,18 @@ type packagesMap struct {
 func newPackagesMap() packagesMap {
 	return packagesMap{
 		impl: persistent.NewMap(func(a, b interface{}) bool {
-			left := a.(packageKey)
-			right := b.(packageKey)
-			if left.mode != right.mode {
-				return left.mode < right.mode
-			}
-			return left.id < right.id
+			return packageKeyLess(a.(packageKey), b.(packageKey))
 		}),
 	}
 }
 
+func packageKeyLess(x, y packageKey) bool {
+	if x.mode != y.mode {
+		return x.mode < y.mode
+	}
+	return x.id < y.id
+}
+
 func (m packagesMap) Clone() packagesMap {
 	return packagesMap{
 		impl: m.impl.Clone(),
@@ -285,3 +287,13 @@ func (s knownDirsSet) Insert(key span.URI) {
 func (s knownDirsSet) Remove(key span.URI) {
 	s.impl.Delete(key)
 }
+
+// actionKeyLessInterface is the less-than relation for actionKey
+// values wrapped in an interface.
+func actionKeyLessInterface(a, b interface{}) bool {
+	x, y := a.(actionKey), b.(actionKey)
+	if x.analyzer.Name != y.analyzer.Name {
+		return x.analyzer.Name < y.analyzer.Name
+	}
+	return packageKeyLess(x.pkg, y.pkg)
+}
diff --git a/internal/lsp/cache/session.go b/internal/lsp/cache/session.go
index aca46dd2a..98d3c2504 100644
--- a/internal/lsp/cache/session.go
+++ b/internal/lsp/cache/session.go
@@ -16,6 +16,7 @@ import (
 	"golang.org/x/tools/internal/imports"
 	"golang.org/x/tools/internal/lsp/progress"
 	"golang.org/x/tools/internal/lsp/source"
+	"golang.org/x/tools/internal/persistent"
 	"golang.org/x/tools/internal/span"
 	"golang.org/x/tools/internal/xcontext"
 )
@@ -238,7 +239,7 @@ func (s *Session) createView(ctx context.Context, name string, folder span.URI,
 		goFiles:              newGoFilesMap(),
 		parseKeysByURI:       newParseKeysByURIMap(),
 		symbols:              make(map[span.URI]*symbolHandle),
-		actions:              make(map[actionKey]*actionHandle),
+		actions:              persistent.NewMap(actionKeyLessInterface),
 		workspacePackages:    make(map[PackageID]PackagePath),
 		unloadableFiles:      make(map[span.URI]struct{}),
 		parseModHandles:      make(map[span.URI]*parseModHandle),
diff --git a/internal/lsp/cache/snapshot.go b/internal/lsp/cache/snapshot.go
index 36dcafeac..93316653a 100644
--- a/internal/lsp/cache/snapshot.go
+++ b/internal/lsp/cache/snapshot.go
@@ -36,6 +36,7 @@ import (
 	"golang.org/x/tools/internal/lsp/source"
 	"golang.org/x/tools/internal/memoize"
 	"golang.org/x/tools/internal/packagesinternal"
+	"golang.org/x/tools/internal/persistent"
 	"golang.org/x/tools/internal/span"
 	"golang.org/x/tools/internal/typesinternal"
 )
@@ -87,7 +88,7 @@ type snapshot struct {
 	// TODO(rfindley): consider merging this with files to reduce burden on clone.
 	symbols map[span.URI]*symbolHandle
 
-	// packages maps a packageKey to a set of packageHandles to which that file belongs.
+	// packages maps a packageKey to a *packageHandle.
 	// It may be invalidated when a file's content changes.
 	packages packagesMap
 
@@ -95,8 +96,9 @@ type snapshot struct {
 	// It may be invalidated when metadata changes or a new file is opened or closed.
 	isActivePackageCache isActivePackageCacheMap
 
-	// actions maps an actionkey to its actionHandle.
-	actions map[actionKey]*actionHandle
+	// actions maps an actionKey to the handle for the future
+	// result of execution an analysis pass on a package.
+	actions *persistent.Map // from actionKey to *actionHandle
 
 	// workspacePackages contains the workspace's packages, which are loaded
 	// when the view is created.
@@ -149,6 +151,7 @@ func (s *snapshot) Destroy(destroyedBy string) {
 	s.generation.Destroy(destroyedBy)
 	s.packages.Destroy()
 	s.isActivePackageCache.Destroy()
+	s.actions.Destroy()
 	s.files.Destroy()
 	s.goFiles.Destroy()
 	s.parseKeysByURI.Destroy()
@@ -1177,9 +1180,6 @@ func (s *snapshot) addSymbolHandle(uri span.URI, sh *symbolHandle) *symbolHandle
 }
 
 func (s *snapshot) getActionHandle(id PackageID, m source.ParseMode, a *analysis.Analyzer) *actionHandle {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
 	key := actionKey{
 		pkg: packageKey{
 			id:   id,
@@ -1187,13 +1187,18 @@ func (s *snapshot) getActionHandle(id PackageID, m source.ParseMode, a *analysis
 		},
 		analyzer: a,
 	}
-	return s.actions[key]
-}
 
-func (s *snapshot) addActionHandle(ah *actionHandle) *actionHandle {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	ah, ok := s.actions.Get(key)
+	if !ok {
+		return nil
+	}
+	return ah.(*actionHandle)
+}
+
+func (s *snapshot) addActionHandle(ah *actionHandle, release func()) *actionHandle {
 	key := actionKey{
 		analyzer: ah.analyzer,
 		pkg: packageKey{
@@ -1201,10 +1206,17 @@ func (s *snapshot) addActionHandle(ah *actionHandle) *actionHandle {
 			mode: ah.pkg.mode,
 		},
 	}
-	if ah, ok := s.actions[key]; ok {
-		return ah
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	// If another thread since cached a different handle,
+	// return it instead of overriding it.
+	if result, ok := s.actions.Get(key); ok {
+		release()
+		return result.(*actionHandle)
 	}
-	s.actions[key] = ah
+	s.actions.Set(key, ah, func(_, _ interface{}) { release() })
 	return ah
 }
 
@@ -1716,7 +1728,7 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 		initializedErr:       s.initializedErr,
 		packages:             s.packages.Clone(),
 		isActivePackageCache: s.isActivePackageCache.Clone(),
-		actions:              make(map[actionKey]*actionHandle, len(s.actions)),
+		actions:              s.actions.Clone(),
 		files:                s.files.Clone(),
 		goFiles:              s.goFiles.Clone(),
 		parseKeysByURI:       s.parseKeysByURI.Clone(),
@@ -1920,13 +1932,17 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 		}
 	}
 
-	// Copy the package analysis information.
-	for k, v := range s.actions {
-		if _, ok := idsToInvalidate[k.pkg.id]; ok {
-			continue
+	// Copy actions.
+	// TODO(adonovan): opt: avoid iteration over s.actions.
+	var actionsToDelete []actionKey
+	s.actions.Range(func(k, _ interface{}) {
+		key := k.(actionKey)
+		if _, ok := idsToInvalidate[key.pkg.id]; ok {
+			actionsToDelete = append(actionsToDelete, key)
 		}
-		newGen.Inherit(v.handle)
-		result.actions[k] = v
+	})
+	for _, key := range actionsToDelete {
+		result.actions.Delete(key)
 	}
 
 	// If the workspace mode has changed, we must delete all metadata, as it