internal/lsp/cache: clone the metadata graph when clearing ShouldLoad

The Metadata type was mutated in exactly one place: when setting the
ShouldLoad bit after package loading completes. Address this by cloning
the metadata graph when clearing ShouldLoad. After this change, metadata
graphs and the data within them are immutable.

This also fixes a range-variable capture bug in load.go: previously we
were deferring a call to clearShouldLoad for the range variable scope.
After this change, we properly clear the ShouldLoad bit for all scopes.

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

internal/lsp/cache/load.go

@@ -37,6 +37,15 @@ func (s *snapshot) load(ctx context.Context, allowNetwork bool, scopes ...interf
 	var query []string
 	var containsDir bool // for logging
 
+	// Unless the context was canceled, set "shouldLoad" to false for all
+	// of the metadata we attempted to load.
+	defer func() {
+		if errors.Is(err, context.Canceled) {
+			return
+		}
+		s.clearShouldLoad(scopes...)
+	}()
+
 	// Keep track of module query -> module path so that we can later correlate query
 	// errors with errors.
 	moduleQueries := make(map[string]string)
@@ -44,14 +53,6 @@ func (s *snapshot) load(ctx context.Context, allowNetwork bool, scopes ...interf
 		if !s.shouldLoad(scope) {
 			continue
 		}
-		// Unless the context was canceled, set "shouldLoad" to false for all
-		// of the metadata we attempted to load.
-		defer func() {
-			if errors.Is(err, context.Canceled) {
-				return
-			}
-			s.clearShouldLoad(scope)
-		}()
 		switch scope := scope.(type) {
 		case PackagePath:
 			if source.IsCommandLineArguments(string(scope)) {
@@ -196,7 +197,7 @@ func (s *snapshot) load(ctx context.Context, allowNetwork bool, scopes ...interf
 		}
 		// TODO: once metadata is immutable, we shouldn't have to lock here.
 		s.mu.Lock()
-		err := s.setMetadataLocked(ctx, PackagePath(pkg.PkgPath), pkg, cfg, query, updates, nil)
+		err := s.computeMetadataUpdates(ctx, PackagePath(pkg.PkgPath), pkg, cfg, query, updates, nil)
 		s.mu.Unlock()
 		if err != nil {
 			return err
@@ -438,10 +439,10 @@ func getWorkspaceDir(ctx context.Context, h *memoize.Handle, g *memoize.Generati
 	return span.URIFromPath(v.(*workspaceDirData).dir), nil
 }
 
-// setMetadataLocked extracts metadata from pkg and records it in s. It
-// recurs through pkg.Imports to ensure that metadata exists for all
-// dependencies.
-func (s *snapshot) setMetadataLocked(ctx context.Context, pkgPath PackagePath, pkg *packages.Package, cfg *packages.Config, query []string, updates map[PackageID]*KnownMetadata, path []PackageID) error {
+// computeMetadataUpdates populates the updates map with metadata updates to
+// apply, based on the given pkg. It recurs through pkg.Imports to ensure that
+// metadata exists for all dependencies.
+func (s *snapshot) computeMetadataUpdates(ctx context.Context, pkgPath PackagePath, pkg *packages.Package, cfg *packages.Config, query []string, updates map[PackageID]*KnownMetadata, path []PackageID) error {
 	id := PackageID(pkg.ID)
 	if new := updates[id]; new != nil {
 		return nil
@@ -532,7 +533,7 @@ func (s *snapshot) setMetadataLocked(ctx context.Context, pkgPath PackagePath, p
 			continue
 		}
 		if s.noValidMetadataForIDLocked(importID) {
-			if err := s.setMetadataLocked(ctx, importPkgPath, importPkg, cfg, query, updates, append(path, id)); err != nil {
+			if err := s.computeMetadataUpdates(ctx, importPkgPath, importPkg, cfg, query, updates, append(path, id)); err != nil {
 				event.Error(ctx, "error in dependency", err)
 			}
 		}

internal/lsp/cache/snapshot.go

@@ -1191,10 +1191,12 @@ func (s *snapshot) shouldLoad(scope interface{}) bool {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
+	g := s.meta
+
 	switch scope := scope.(type) {
 	case PackagePath:
 		var meta *KnownMetadata
-		for _, m := range s.meta.metadata {
+		for _, m := range g.metadata {
 			if m.PkgPath != scope {
 				continue
 			}
@@ -1206,12 +1208,12 @@ func (s *snapshot) shouldLoad(scope interface{}) bool {
 		return false
 	case fileURI:
 		uri := span.URI(scope)
-		ids := s.meta.ids[uri]
+		ids := g.ids[uri]
 		if len(ids) == 0 {
 			return true
 		}
 		for _, id := range ids {
-			m, ok := s.meta.metadata[id]
+			m, ok := g.metadata[id]
 			if !ok || m.ShouldLoad {
 				return true
 			}
@@ -1222,34 +1224,40 @@ func (s *snapshot) shouldLoad(scope interface{}) bool {
 	}
 }
 
-func (s *snapshot) clearShouldLoad(scope interface{}) {
+func (s *snapshot) clearShouldLoad(scopes ...interface{}) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	switch scope := scope.(type) {
-	case PackagePath:
-		var meta *KnownMetadata
-		for _, m := range s.meta.metadata {
-			if m.PkgPath == scope {
-				meta = m
+	g := s.meta
+
+	var updates map[PackageID]*KnownMetadata
+	markLoaded := func(m *KnownMetadata) {
+		if updates == nil {
+			updates = make(map[PackageID]*KnownMetadata)
+		}
+		next := *m
+		next.ShouldLoad = false
+		updates[next.ID] = &next
+	}
+	for _, scope := range scopes {
+		switch scope := scope.(type) {
+		case PackagePath:
+			for _, m := range g.metadata {
+				if m.PkgPath == scope {
+					markLoaded(m)
+				}
 			}
-		}
-		if meta == nil {
-			return
-		}
-		meta.ShouldLoad = false
-	case fileURI:
-		uri := span.URI(scope)
-		ids := s.meta.ids[uri]
-		if len(ids) == 0 {
-			return
-		}
-		for _, id := range ids {
-			if m, ok := s.meta.metadata[id]; ok {
-				m.ShouldLoad = false
+		case fileURI:
+			uri := span.URI(scope)
+			ids := g.ids[uri]
+			for _, id := range ids {
+				if m, ok := g.metadata[id]; ok {
+					markLoaded(m)
+				}
 			}
 		}
 	}
+	s.meta = g.Clone(updates)
 }
 
 // noValidMetadataForURILocked reports whether there is any valid metadata for