gopls/internal/lsp/cache: move upgrades and vulns onto the snapshot

To simplify the View and eliminate races, treat module upgrades and
vulnerabilities like any other state by tracking them in the snapshot.

With this change, and an additional change to track diagnostics done on
behalf of ExecuteCommand, we can make TestUpgradeCodelens properly
assertive and fix the underlying race of golang/go#58750.

Also rewrite a unit test checking vuln result expiry as a regtest.

Updates golang/go#57979
Fixes golang/go#58750

Change-Id: I4521c97f798eecd13844278a304070f661538382
Reviewed-on: https://go-review.googlesource.com/c/tools/+/540479
Reviewed-by: Alan Donovan <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: findleyr

gopls/internal/immutable/immutable.go

@@ -20,6 +20,15 @@ func MapOf[K comparable, V any](m map[K]V) Map[K, V] {
 	return Map[K, V]{m}
 }
 
+// Keys returns all keys present in the map.
+func (m Map[K, V]) Keys() []K {
+	var keys []K
+	for k := range m.m {
+		keys = append(keys, k)
+	}
+	return keys
+}
+
 // Value returns the mapped value for k.
 // It is equivalent to the commaok form of an ordinary go map, and returns
 // (zero, false) if the key is not present.

gopls/internal/lsp/cache/session.go

@@ -128,8 +128,6 @@ func (s *Session) createView(ctx context.Context, info *workspaceInformation, fo
 		initialWorkspaceLoad: make(chan struct{}),
 		initializationSema:   make(chan struct{}, 1),
 		baseCtx:              baseCtx,
-		moduleUpgrades:       map[span.URI]map[string]string{},
-		vulns:                map[span.URI]*vulncheck.Result{},
 		parseCache:           s.parseCache,
 		fs:                   s.overlayFS,
 		workspaceInformation: info,
@@ -173,6 +171,8 @@ func (s *Session) createView(ctx context.Context, info *workspaceInformation, fo
 		workspaceModFiles:    wsModFiles,
 		workspaceModFilesErr: wsModFilesErr,
 		pkgIndex:             typerefs.NewPackageIndex(),
+		moduleUpgrades:       new(persistent.Map[span.URI, map[string]string]),
+		vulns:                new(persistent.Map[span.URI, *vulncheck.Result]),
 	}
 	// Save one reference in the view.
 	v.releaseSnapshot = v.snapshot.Acquire()
@@ -501,9 +501,9 @@ func (s *Session) DidModifyFiles(ctx context.Context, changes []source.FileModif
 	}
 
 	var releases []func()
-	viewToSnapshot := map[*View]*snapshot{}
+	viewToSnapshot := make(map[*View]source.Snapshot)
 	for view, changed := range views {
-		snapshot, release := view.invalidateContent(ctx, changed)
+		snapshot, release := view.Invalidate(ctx, source.StateChange{Files: changed})
 		releases = append(releases, release)
 		viewToSnapshot[view] = snapshot
 	}

gopls/internal/lsp/cache/snapshot.go

@@ -40,6 +40,8 @@ import (
 	"golang.org/x/tools/gopls/internal/lsp/source/xrefs"
 	"golang.org/x/tools/gopls/internal/persistent"
 	"golang.org/x/tools/gopls/internal/span"
+	"golang.org/x/tools/gopls/internal/vulncheck"
+	"golang.org/x/tools/internal/constraints"
 	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/event/tag"
 	"golang.org/x/tools/internal/gocommand"
@@ -177,6 +179,13 @@ type snapshot struct {
 	// detect ignored files.
 	ignoreFilterOnce sync.Once
 	ignoreFilter     *ignoreFilter
+
+	// moduleUpgrades tracks known upgrades for module paths in each modfile.
+	// Each modfile has a map of module name to upgrade version.
+	moduleUpgrades *persistent.Map[span.URI, map[string]string]
+
+	// vulns maps each go.mod file's URI to its known vulnerabilities.
+	vulns *persistent.Map[span.URI, *vulncheck.Result]
 }
 
 var globalSnapshotID uint64
@@ -251,6 +260,8 @@ func (s *snapshot) destroy(destroyedBy string) {
 	s.modVulnHandles.Destroy()
 	s.modWhyHandles.Destroy()
 	s.unloadableFiles.Destroy()
+	s.moduleUpgrades.Destroy()
+	s.vulns.Destroy()
 }
 
 func (s *snapshot) SequenceID() uint64 {
@@ -1814,7 +1825,8 @@ func inVendor(uri span.URI) bool {
 	return found && strings.Contains(after, "/")
 }
 
-func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]source.FileHandle) (*snapshot, func()) {
+func (s *snapshot) clone(ctx, bgCtx context.Context, changed source.StateChange) (*snapshot, func()) {
+	changedFiles := changed.Files
 	ctx, done := event.Start(ctx, "cache.snapshot.clone")
 	defer done()
 
@@ -1834,20 +1846,22 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]source
 		initializedErr:       s.initializedErr,
 		packages:             s.packages.Clone(),
 		activePackages:       s.activePackages.Clone(),
-		files:                s.files.Clone(changes),
-		symbolizeHandles:     cloneWithout(s.symbolizeHandles, changes),
+		files:                s.files.Clone(changedFiles),
+		symbolizeHandles:     cloneWithout(s.symbolizeHandles, changedFiles),
 		workspacePackages:    s.workspacePackages,
 		shouldLoad:           s.shouldLoad.Clone(),      // not cloneWithout: shouldLoad is cleared on loads
 		unloadableFiles:      s.unloadableFiles.Clone(), // not cloneWithout: typing in a file doesn't necessarily make it loadable
-		parseModHandles:      cloneWithout(s.parseModHandles, changes),
-		parseWorkHandles:     cloneWithout(s.parseWorkHandles, changes),
-		modTidyHandles:       cloneWithout(s.modTidyHandles, changes),
-		modWhyHandles:        cloneWithout(s.modWhyHandles, changes),
-		modVulnHandles:       cloneWithout(s.modVulnHandles, changes),
+		parseModHandles:      cloneWithout(s.parseModHandles, changedFiles),
+		parseWorkHandles:     cloneWithout(s.parseWorkHandles, changedFiles),
+		modTidyHandles:       cloneWithout(s.modTidyHandles, changedFiles),
+		modWhyHandles:        cloneWithout(s.modWhyHandles, changedFiles),
+		modVulnHandles:       cloneWithout(s.modVulnHandles, changedFiles),
 		workspaceModFiles:    s.workspaceModFiles,
 		workspaceModFilesErr: s.workspaceModFilesErr,
 		importGraph:          s.importGraph,
 		pkgIndex:             s.pkgIndex,
+		moduleUpgrades:       cloneWith(s.moduleUpgrades, changed.ModuleUpgrades),
+		vulns:                cloneWith(s.vulns, changed.Vulns),
 	}
 
 	// Create a lease on the new snapshot.
@@ -1864,7 +1878,7 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]source
 	// vendor tree after 'go mod vendor' or 'rm -fr vendor/'.
 	//
 	// TODO(rfindley): revisit the location of this check.
-	for uri := range changes {
+	for uri := range changedFiles {
 		if inVendor(uri) && s.initializedErr != nil ||
 			strings.HasSuffix(string(uri), "/vendor/modules.txt") {
 			reinit = true
@@ -1877,7 +1891,7 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]source
 	// where a file is added on disk; we don't want to read the newly added file
 	// into the old snapshot, as that will break our change detection below.
 	oldFiles := make(map[span.URI]source.FileHandle)
-	for uri := range changes {
+	for uri := range changedFiles {
 		if fh, ok := s.files.Get(uri); ok {
 			oldFiles[uri] = fh
 		}
@@ -1898,7 +1912,7 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]source
 	}
 
 	if workURI, _ := s.view.GOWORK(); workURI != "" {
-		if newFH, ok := changes[workURI]; ok {
+		if newFH, ok := changedFiles[workURI]; ok {
 			result.workspaceModFiles, result.workspaceModFilesErr = computeWorkspaceModFiles(ctx, s.view.gomod, workURI, s.view.effectiveGO111MODULE(), result)
 			if changedOnDisk(oldFiles[workURI], newFH) {
 				reinit = true
@@ -1907,14 +1921,14 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]source
 	}
 
 	// Reinitialize if any workspace mod file has changed on disk.
-	for uri, newFH := range changes {
+	for uri, newFH := range changedFiles {
 		if _, ok := result.workspaceModFiles[uri]; ok && changedOnDisk(oldFiles[uri], newFH) {
 			reinit = true
 		}
 	}
 
 	// Finally, process sumfile changes that may affect loading.
-	for uri, newFH := range changes {
+	for uri, newFH := range changedFiles {
 		if !changedOnDisk(oldFiles[uri], newFH) {
 			continue // like with go.mod files, we only reinit when things change on disk
 		}
@@ -1955,7 +1969,7 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]source
 	anyFileOpenedOrClosed := false // opened files affect workspace packages
 	anyFileAdded := false          // adding a file can resolve missing dependencies
 
-	for uri, newFH := range changes {
+	for uri, newFH := range changedFiles {
 		// The original FileHandle for this URI is cached on the snapshot.
 		oldFH, _ := oldFiles[uri] // may be nil
 		_, oldOpen := oldFH.(*Overlay)
@@ -2160,14 +2174,24 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]source
 	return result, release
 }
 
-func cloneWithout[V any](m *persistent.Map[span.URI, V], changes map[span.URI]source.FileHandle) *persistent.Map[span.URI, V] {
+// cloneWithout clones m then deletes from it the keys of changes.
+func cloneWithout[K constraints.Ordered, V1, V2 any](m *persistent.Map[K, V1], changes map[K]V2) *persistent.Map[K, V1] {
 	m2 := m.Clone()
 	for k := range changes {
 		m2.Delete(k)
 	}
 	return m2
 }
 
+// cloneWith clones m then inserts the changes into it.
+func cloneWith[K constraints.Ordered, V any](m *persistent.Map[K, V], changes map[K]V) *persistent.Map[K, V] {
+	m2 := m.Clone()
+	for k, v := range changes {
+		m2.Set(k, v, nil)
+	}
+	return m2
+}
+
 // deleteMostRelevantModFile deletes the mod file most likely to be the mod
 // file for the changed URI, if it exists.
 //

gopls/internal/lsp/cache/view.go

@@ -65,15 +65,6 @@ type View struct {
 
 	importsState *importsState
 
-	// moduleUpgrades tracks known upgrades for module paths in each modfile.
-	// Each modfile has a map of module name to upgrade version.
-	moduleUpgradesMu sync.Mutex
-	moduleUpgrades   map[span.URI]map[string]string
-
-	// vulns maps each go.mod file's URI to its known vulnerabilities.
-	vulnsMu sync.Mutex
-	vulns   map[span.URI]*vulncheck.Result
-
 	// parseCache holds an LRU cache of recently parsed files.
 	parseCache *parseCache
 
@@ -864,14 +855,13 @@ func (s *snapshot) loadWorkspace(ctx context.Context, firstAttempt bool) (loadEr
 	return loadErr
 }
 
-// invalidateContent invalidates the content of a Go file,
-// including any position and type information that depends on it.
+// Invalidate processes the provided state change, invalidating any derived
+// results that depend on the changed state.
 //
-// invalidateContent returns a non-nil snapshot for the new content, along with
-// a callback which the caller must invoke to release that snapshot.
-//
-// newOptions may be nil, in which case options remain unchanged.
-func (v *View) invalidateContent(ctx context.Context, changes map[span.URI]source.FileHandle) (*snapshot, func()) {
+// The resulting snapshot is non-nil, representing the outcome of the state
+// change. The second result is a function that must be called to release the
+// snapshot when the snapshot is no longer needed.
+func (v *View) Invalidate(ctx context.Context, changed source.StateChange) (source.Snapshot, func()) {
 	// Detach the context so that content invalidation cannot be canceled.
 	ctx = xcontext.Detach(ctx)
 
@@ -893,7 +883,7 @@ func (v *View) invalidateContent(ctx context.Context, changes map[span.URI]sourc
 	prevSnapshot.AwaitInitialized(ctx)
 
 	// Save one lease of the cloned snapshot in the view.
-	v.snapshot, v.releaseSnapshot = prevSnapshot.clone(ctx, v.baseCtx, changes)
+	v.snapshot, v.releaseSnapshot = prevSnapshot.clone(ctx, v.baseCtx, changed)
 
 	prevReleaseSnapshot()
 	v.destroy(prevSnapshot, "View.invalidateContent")
@@ -1037,75 +1027,44 @@ func (v *View) IsGoPrivatePath(target string) bool {
 	return globsMatchPath(v.goprivate, target)
 }
 
-func (v *View) ModuleUpgrades(modfile span.URI) map[string]string {
-	v.moduleUpgradesMu.Lock()
-	defer v.moduleUpgradesMu.Unlock()
-
+func (s *snapshot) ModuleUpgrades(modfile span.URI) map[string]string {
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	upgrades := map[string]string{}
-	for mod, ver := range v.moduleUpgrades[modfile] {
+	orig, _ := s.moduleUpgrades.Get(modfile)
+	for mod, ver := range orig {
 		upgrades[mod] = ver
 	}
 	return upgrades
 }
 
-func (v *View) RegisterModuleUpgrades(modfile span.URI, upgrades map[string]string) {
-	// Return early if there are no upgrades.
-	if len(upgrades) == 0 {
-		return
-	}
-
-	v.moduleUpgradesMu.Lock()
-	defer v.moduleUpgradesMu.Unlock()
-
-	m := v.moduleUpgrades[modfile]
-	if m == nil {
-		m = make(map[string]string)
-		v.moduleUpgrades[modfile] = m
-	}
-	for mod, ver := range upgrades {
-		m[mod] = ver
-	}
-}
-
-func (v *View) ClearModuleUpgrades(modfile span.URI) {
-	v.moduleUpgradesMu.Lock()
-	defer v.moduleUpgradesMu.Unlock()
-
-	delete(v.moduleUpgrades, modfile)
-}
-
-const maxGovulncheckResultAge = 1 * time.Hour // Invalidate results older than this limit.
-var timeNow = time.Now                        // for testing
+// MaxGovulncheckResultsAge defines the maximum vulnerability age considered
+// valid by gopls.
+//
+// Mutable for testing.
+var MaxGovulncheckResultAge = 1 * time.Hour
 
-func (v *View) Vulnerabilities(modfiles ...span.URI) map[span.URI]*vulncheck.Result {
+// TODO(rfindley): move to snapshot.go
+func (s *snapshot) Vulnerabilities(modfiles ...span.URI) map[span.URI]*vulncheck.Result {
 	m := make(map[span.URI]*vulncheck.Result)
-	now := timeNow()
-	v.vulnsMu.Lock()
-	defer v.vulnsMu.Unlock()
+	now := time.Now()
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
 
 	if len(modfiles) == 0 { // empty means all modfiles
-		for modfile := range v.vulns {
-			modfiles = append(modfiles, modfile)
-		}
+		modfiles = s.vulns.Keys()
 	}
 	for _, modfile := range modfiles {
-		vuln := v.vulns[modfile]
-		if vuln != nil && now.Sub(vuln.AsOf) > maxGovulncheckResultAge {
-			v.vulns[modfile] = nil // same as SetVulnerabilities(modfile, nil)
+		vuln, _ := s.vulns.Get(modfile)
+		if vuln != nil && now.Sub(vuln.AsOf) > MaxGovulncheckResultAge {
 			vuln = nil
 		}
 		m[modfile] = vuln
 	}
 	return m
 }
 
-func (v *View) SetVulnerabilities(modfile span.URI, vulns *vulncheck.Result) {
-	v.vulnsMu.Lock()
-	defer v.vulnsMu.Unlock()
-
-	v.vulns[modfile] = vulns
-}
-
 func (v *View) GoVersion() int {
 	return v.workspaceInformation.goversion
 }