internal/lsp/cache: invalidate packages in setMetadata

When setting metadata, we must invalidate any packages that have been
computed with old metadata.

Also make setMetadata atomic, by locking around it in Load. This is just
writing memory after a Load, so should be fast and infrequent. It is
critical that our various maps related to metadata are coherent, so it
makes sense to err on the side of coarser locking, IMO.

Fix a race in TestUpgradeCodeLens.

Change-Id: Iee8003b7e52b9f21681bdba08a009824f4ac6268
Reviewed-on: https://go-review.googlesource.com/c/tools/+/331809
Trust: Robert Findley <rfindley@google.com>
Run-TryBot: Robert Findley <rfindley@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/gopls/internal/regtest/codelens/codelens_test.go b/gopls/internal/regtest/codelens/codelens_test.go
index a2a1ae4..d89b8e0 100644
--- a/gopls/internal/regtest/codelens/codelens_test.go
+++ b/gopls/internal/regtest/codelens/codelens_test.go
@@ -166,11 +166,16 @@
 				if vendoring {
 					env.RunGoCommand("mod", "vendor")
 				}
+				env.Await(env.DoneWithChangeWatchedFiles())
 				env.OpenFile("go.mod")
 				env.ExecuteCodeLensCommand("go.mod", command.CheckUpgrades)
 				d := &protocol.PublishDiagnosticsParams{}
-				env.Await(OnceMet(env.DiagnosticAtRegexpWithMessage("go.mod", `require`, "can be upgraded"),
-					ReadDiagnostics("go.mod", d)))
+				env.Await(
+					OnceMet(
+						env.DiagnosticAtRegexpWithMessage("go.mod", `require`, "can be upgraded"),
+						ReadDiagnostics("go.mod", d),
+					),
+				)
 				env.ApplyQuickFixes("go.mod", d.Diagnostics)
 				env.Await(env.DoneWithChangeWatchedFiles())
 				if got := env.Editor.BufferText("go.mod"); got != wantGoMod {
diff --git a/internal/lsp/cache/load.go b/internal/lsp/cache/load.go
index 5a69875..1baa3e5 100644
--- a/internal/lsp/cache/load.go
+++ b/internal/lsp/cache/load.go
@@ -194,7 +194,9 @@
 			continue
 		}
 		// Set the metadata for this package.
-		m, err := s.setMetadata(ctx, packagePath(pkg.PkgPath), pkg, cfg, map[packageID]struct{}{})
+		s.mu.Lock()
+		m, err := s.setMetadataLocked(ctx, packagePath(pkg.PkgPath), pkg, cfg, map[packageID]struct{}{})
+		s.mu.Unlock()
 		if err != nil {
 			return err
 		}
@@ -388,10 +390,10 @@
 	return span.URIFromPath(v.(*workspaceDirData).dir), nil
 }
 
-// setMetadata extracts metadata from pkg and records it in s. It
+// setMetadataLocked extracts metadata from pkg and records it in s. It
 // recurses through pkg.Imports to ensure that metadata exists for all
 // dependencies.
-func (s *snapshot) setMetadata(ctx context.Context, pkgPath packagePath, pkg *packages.Package, cfg *packages.Config, seen map[packageID]struct{}) (*metadata, error) {
+func (s *snapshot) setMetadataLocked(ctx context.Context, pkgPath packagePath, pkg *packages.Package, cfg *packages.Config, seen map[packageID]struct{}) (*metadata, error) {
 	id := packageID(pkg.ID)
 	if _, ok := seen[id]; ok {
 		return nil, errors.Errorf("import cycle detected: %q", id)
@@ -429,7 +431,7 @@
 		m.goFiles = append(m.goFiles, uri)
 		uris[uri] = struct{}{}
 	}
-	s.updateIDForURIs(id, uris)
+	s.updateIDForURIsLocked(id, uris)
 
 	// TODO(rstambler): is this still necessary?
 	copied := map[packageID]struct{}{
@@ -452,16 +454,14 @@
 			m.missingDeps[importPkgPath] = struct{}{}
 			continue
 		}
-		if s.noValidMetadataForID(importID) {
-			if _, err := s.setMetadata(ctx, importPkgPath, importPkg, cfg, copied); err != nil {
+		if s.noValidMetadataForIDLocked(importID) {
+			if _, err := s.setMetadataLocked(ctx, importPkgPath, importPkg, cfg, copied); err != nil {
 				event.Error(ctx, "error in dependency", err)
 			}
 		}
 	}
 
 	// Add the metadata to the cache.
-	s.mu.Lock()
-	defer s.mu.Unlock()
 
 	// If we've already set the metadata for this snapshot, reuse it.
 	if original, ok := s.metadata[m.id]; ok && original.valid {
@@ -473,6 +473,11 @@
 			metadata: m,
 			valid:    true,
 		}
+		// Invalidate any packages we may have associated with this metadata.
+		for _, mode := range []source.ParseMode{source.ParseHeader, source.ParseExported, source.ParseFull} {
+			key := packageKey{mode, m.id}
+			delete(s.packages, key)
+		}
 	}
 
 	// Set the workspace packages. If any of the package's files belong to the
diff --git a/internal/lsp/cache/snapshot.go b/internal/lsp/cache/snapshot.go
index 207be3a..21ca384 100644
--- a/internal/lsp/cache/snapshot.go
+++ b/internal/lsp/cache/snapshot.go
@@ -1114,19 +1114,22 @@
 // noValidMetadataForID reports whether there is no valid metadata for the
 // given ID.
 func (s *snapshot) noValidMetadataForID(id packageID) bool {
-	m := s.getMetadata(id)
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.noValidMetadataForIDLocked(id)
+}
+
+func (s *snapshot) noValidMetadataForIDLocked(id packageID) bool {
+	m := s.metadata[id]
 	return m == nil || !m.valid
 }
 
-// updateIDForURIs adds the given ID to the set of known IDs for the given URI.
+// updateIDForURIsLocked adds the given ID to the set of known IDs for the given URI.
 // Any existing invalid IDs are removed from the set of known IDs. IDs that are
 // not "command-line-arguments" are preferred, so if a new ID comes in for a
 // URI that previously only had "command-line-arguments", the new ID will
 // replace the "command-line-arguments" ID.
-func (s *snapshot) updateIDForURIs(id packageID, uris map[span.URI]struct{}) {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
+func (s *snapshot) updateIDForURIsLocked(id packageID, uris map[span.URI]struct{}) {
 	for uri := range uris {
 		// Collect the new set of IDs, preserving any valid existing IDs.
 		newIDs := []packageID{id}