internal/lsp: reload workspace package metadata on demand

This change removes functions from the snapshot that return package IDs.
We prefer PackageHandles, since getting PackageHandles in a granular
fashion is not effective and causes us to spawn many `go list`
processes. By only ever returning PackageHandles, we can batch metadata
reloads for workspace packages. This enables us to add a check to
confirm that the snapshot is in a good state before returning important
data, like reverse dependencies and workspace package handles.

Change-Id: Icffc8d8e0449864f207c15aa211e84cb158c163f
Reviewed-on: https://go-review.googlesource.com/c/tools/+/214383
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>

internal/lsp/cache/load.go

@@ -38,10 +38,15 @@ type metadata struct {
 func (s *snapshot) load(ctx context.Context, scope interface{}) ([]*metadata, error) {
 	var query string
 	switch scope := scope.(type) {
+	case []packagePath:
+		for i, p := range scope {
+			if i != 0 {
+				query += " "
+			}
+			query += string(p)
+		}
 	case packagePath:
 		query = string(scope)
-	case packageID:
-		query = string(scope)
 	case fileURI:
 		query = fmt.Sprintf("file=%s", span.URI(scope).Filename())
 	case directoryURI:

internal/lsp/cache/session.go

@@ -293,7 +293,7 @@ func (s *session) DidModifyFile(ctx context.Context, c source.FileModification)
 			// If the file was already known in the snapshot,
 			// then use the already known file kind. Otherwise,
 			// detect the file kind. This should only be needed for file creates.
-			if fh := view.getSnapshot().findFileHandle(ctx, f); fh != nil {
+			if fh := view.getSnapshot().findFileHandle(f); fh != nil {
 				kind = fh.Identity().Kind
 			} else {
 				kind = source.DetectLanguage("", c.URI.Filename())

internal/lsp/cache/snapshot.go

@@ -121,21 +121,20 @@ func (s *snapshot) PackageHandles(ctx context.Context, fh source.FileHandle) ([]
 	return results, nil
 }
 
-func (s *snapshot) PackageHandle(ctx context.Context, pkgID string) (source.PackageHandle, error) {
-	id := packageID(pkgID)
-	var meta []*metadata
-	if m := s.getMetadata(id); m != nil {
-		meta = append(meta, m)
+func (s *snapshot) packageHandle(ctx context.Context, id packageID) (*packageHandle, error) {
+	m := s.getMetadata(id)
+
+	// Don't reload metadata in this function.
+	// Callers of this function must reload metadata themselves.
+	if m == nil {
+		return nil, errors.Errorf("%s has no metadata", id)
 	}
-	// We might need to reload the package. If it is a workspace package,
-	// it may be a test variant and therefore have a different scope.
-	var scope interface{} = id
-	if path, ok := s.isWorkspacePackage(id); ok {
-		scope = path
+	phs, load, check := s.shouldCheck([]*metadata{m})
+	if load {
+		return nil, errors.Errorf("%s needs loading", id)
 	}
-	phs, err := s.packageHandles(ctx, scope, meta)
-	if err != nil {
-		return nil, err
+	if check {
+		return s.buildPackageHandle(ctx, m.id, source.ParseFull)
 	}
 	var result *packageHandle
 	for _, ph := range phs {
@@ -239,21 +238,23 @@ func (s *snapshot) shouldCheck(m []*metadata) (phs []*packageHandle, load, check
 	return phs, false, false
 }
 
-func (s *snapshot) GetReverseDependencies(ctx context.Context, id string) ([]string, error) {
-	// Do not return results until the view has been initialized.
-	if err := s.awaitInitialized(ctx); err != nil {
+func (s *snapshot) GetReverseDependencies(ctx context.Context, id string) ([]source.PackageHandle, error) {
+	if err := s.awaitLoaded(ctx); err != nil {
 		return nil, err
 	}
-
 	ids := make(map[packageID]struct{})
 	s.transitiveReverseDependencies(packageID(id), ids)
 
 	// Make sure to delete the original package ID from the map.
 	delete(ids, packageID(id))
 
-	var results []string
+	var results []source.PackageHandle
 	for id := range ids {
-		results = append(results, string(id))
+		ph, err := s.packageHandle(ctx, id)
+		if err != nil {
+			return nil, err
+		}
+		results = append(results, ph)
 	}
 	return results, nil
 }
@@ -318,18 +319,33 @@ func (s *snapshot) addPackage(ph *packageHandle) {
 	s.packages[ph.packageKey()] = ph
 }
 
-func (s *snapshot) workspacePackageIDs() (ids []string) {
+func (s *snapshot) workspacePackageIDs() (ids []packageID) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
 	for id := range s.workspacePackages {
-		ids = append(ids, string(id))
+		ids = append(ids, id)
 	}
 	return ids
 }
 
+func (s *snapshot) WorkspacePackages(ctx context.Context) ([]source.PackageHandle, error) {
+	if err := s.awaitLoaded(ctx); err != nil {
+		return nil, err
+	}
+	var results []source.PackageHandle
+	for _, pkgID := range s.workspacePackageIDs() {
+		ph, err := s.packageHandle(ctx, pkgID)
+		if err != nil {
+			return nil, err
+		}
+		results = append(results, ph)
+	}
+	return results, nil
+}
+
 func (s *snapshot) KnownPackages(ctx context.Context) ([]source.PackageHandle, error) {
-	if err := s.awaitInitialized(ctx); err != nil {
+	if err := s.awaitLoaded(ctx); err != nil {
 		return nil, err
 	}
 	// Collect PackageHandles for all of the workspace packages first.
@@ -343,7 +359,7 @@ func (s *snapshot) KnownPackages(ctx context.Context) ([]source.PackageHandle, e
 
 	var results []source.PackageHandle
 	for pkgID := range wsPackages {
-		ph, err := s.PackageHandle(ctx, string(pkgID))
+		ph, err := s.packageHandle(ctx, pkgID)
 		if err != nil {
 			return nil, err
 		}
@@ -374,7 +390,13 @@ func (s *snapshot) KnownPackages(ctx context.Context) ([]source.PackageHandle, e
 	return results, nil
 }
 
-func (s *snapshot) KnownImportPaths() map[string]source.Package {
+func (s *snapshot) CachedImportPaths(ctx context.Context) (map[string]source.Package, error) {
+	// Don't reload workspace package metadata.
+	// This function is meant to only return currently cached information.
+	if err := s.view.awaitInitialized(ctx); err != nil {
+		return nil, err
+	}
+
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -395,7 +417,7 @@ func (s *snapshot) KnownImportPaths() map[string]source.Package {
 			}
 		}
 	}
-	return results
+	return results, nil
 }
 
 func (s *snapshot) getPackage(id packageID, m source.ParseMode) *packageHandle {
@@ -542,13 +564,40 @@ func (s *snapshot) getFileHandle(f *fileBase) source.FileHandle {
 	return s.files[f.URI()]
 }
 
-func (s *snapshot) findFileHandle(ctx context.Context, f *fileBase) source.FileHandle {
+func (s *snapshot) findFileHandle(f *fileBase) source.FileHandle {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
 	return s.files[f.URI()]
 }
 
+func (s *snapshot) awaitLoaded(ctx context.Context) error {
+	// Do not return results until the snapshot's view has been initialized.
+	if err := s.view.awaitInitialized(ctx); err != nil {
+		return err
+	}
+	// Make sure that the workspace is in a valid state.
+	return s.reloadWorkspace(ctx)
+}
+
+// reloadWorkspace reloads the metadata for all invalidated workspace packages.
+func (s *snapshot) reloadWorkspace(ctx context.Context) error {
+	s.mu.Lock()
+	var scope []packagePath
+	for id, pkgPath := range s.workspacePackages {
+		if s.metadata[id] == nil {
+			scope = append(scope, pkgPath)
+		}
+	}
+	s.mu.Unlock()
+
+	if len(scope) == 0 {
+		return nil
+	}
+	_, err := s.load(ctx, scope)
+	return err
+}
+
 func (s *snapshot) clone(ctx context.Context, withoutURI span.URI) *snapshot {
 	s.mu.Lock()
 	defer s.mu.Unlock()

internal/lsp/cache/view.go

@@ -416,6 +416,7 @@ func (v *view) Shutdown(ctx context.Context) {
 }
 
 func (v *view) shutdown(context.Context) {
+	// TODO: Cancel the view's initialization.
 	v.mu.Lock()
 	defer v.mu.Unlock()
 	if v.cancel != nil {
@@ -467,15 +468,6 @@ func (v *view) getSnapshot() *snapshot {
 	return v.snapshot
 }
 
-func (v *view) WorkspacePackageIDs(ctx context.Context) ([]string, error) {
-	s := v.getSnapshot()
-
-	if err := s.awaitInitialized(ctx); err != nil {
-		return nil, err
-	}
-	return s.workspacePackageIDs(), nil
-}
-
 func (v *view) initialize(ctx context.Context, s *snapshot) {
 	v.initializeOnce.Do(func() {
 		defer close(v.initialized)
@@ -493,13 +485,13 @@ func (v *view) initialize(ctx context.Context, s *snapshot) {
 	})
 }
 
-func (s *snapshot) awaitInitialized(ctx context.Context) error {
+func (v *view) awaitInitialized(ctx context.Context) error {
 	select {
 	case <-ctx.Done():
 		return ctx.Err()
-	case <-s.view.initialized:
+	case <-v.initialized:
 	}
-	return s.view.initializationError
+	return v.initializationError
 }
 
 // invalidateContent invalidates the content of a Go file,
@@ -521,8 +513,8 @@ func (v *view) invalidateContent(ctx context.Context, uri span.URI, kind source.
 		v.cancelBackground()
 	}
 
-	// Do not clone a snapshot until the workspace load has been completed.
-	<-v.initialized
+	// Do not clone a snapshot until its view has finished initializing.
+	_ = v.awaitInitialized(ctx)
 
 	// This should be the only time we hold the view's snapshot lock for any period of time.
 	v.snapshotMu.Lock()

internal/lsp/diagnostics.go

@@ -20,25 +20,20 @@ func (s *Server) diagnoseSnapshot(ctx context.Context, snapshot source.Snapshot)
 	ctx, done := trace.StartSpan(ctx, "lsp:background-worker")
 	defer done()
 
-	wsPackages, err := snapshot.View().WorkspacePackageIDs(ctx)
+	wsPackages, err := snapshot.WorkspacePackages(ctx)
 	if err != nil {
 		log.Error(ctx, "diagnoseSnapshot: no workspace packages", err, telemetry.Directory.Of(snapshot.View().Folder))
 		return
 	}
-	for _, id := range wsPackages {
-		go func(id string) {
-			ph, err := snapshot.PackageHandle(ctx, id)
-			if err != nil {
-				log.Error(ctx, "diagnoseSnapshot: no PackageHandle for package", err, telemetry.Package.Of(id))
-				return
-			}
+	for _, ph := range wsPackages {
+		go func(ph source.PackageHandle) {
 			reports, _, err := source.PackageDiagnostics(ctx, snapshot, ph, false, snapshot.View().Options().DisabledAnalyses)
 			if err != nil {
 				log.Error(ctx, "diagnoseSnapshot: no diagnostics", err, telemetry.Package.Of(ph.ID()))
 				return
 			}
 			s.publishReports(ctx, reports, false)
-		}(id)
+		}(ph)
 	}
 	// Run diagnostics on the go.mod file.
 	s.diagnoseModfile(snapshot)

internal/lsp/source/completion.go

@@ -649,7 +649,10 @@ func (c *completer) selector(sel *ast.SelectorExpr) error {
 
 func (c *completer) unimportedMembers(id *ast.Ident) error {
 	// Try loaded packages first. They're relevant, fast, and fully typed.
-	known := c.snapshot.KnownImportPaths()
+	known, err := c.snapshot.CachedImportPaths(c.ctx)
+	if err != nil {
+		return err
+	}
 	for path, pkg := range known {
 		if pkg.GetTypes().Name() != id.Name {
 			continue

internal/lsp/source/diagnostics.go

@@ -120,11 +120,7 @@ func PackageDiagnostics(ctx context.Context, snapshot Snapshot, ph PackageHandle
 		log.Error(ctx, "no reverse dependencies", err)
 		return reports, warningMsg, nil
 	}
-	for _, id := range reverseDeps {
-		ph, err := snapshot.PackageHandle(ctx, id)
-		if err != nil {
-			return nil, warningMsg, err
-		}
+	for _, ph := range reverseDeps {
 		pkg, err := ph.Check(ctx)
 		if err != nil {
 			return nil, warningMsg, err

internal/lsp/source/references.go

@@ -45,11 +45,7 @@ func (i *IdentifierInfo) References(ctx context.Context) ([]*ReferenceInfo, erro
 		if err != nil {
 			return nil, err
 		}
-		for _, id := range reverseDeps {
-			ph, err := i.Snapshot.PackageHandle(ctx, id)
-			if err != nil {
-				return nil, err
-			}
+		for _, ph := range reverseDeps {
 			pkg, err := ph.Check(ctx)
 			if err != nil {
 				return nil, err

internal/lsp/source/view.go

@@ -38,23 +38,24 @@ type Snapshot interface {
 	// ModFiles returns the FileHandles of the go.mod files attached to the view associated with this snapshot.
 	ModFiles(ctx context.Context) (FileHandle, FileHandle, error)
 
-	// PackageHandle returns the CheckPackageHandle for the given package ID.
-	PackageHandle(ctx context.Context, id string) (PackageHandle, error)
-
 	// PackageHandles returns the CheckPackageHandles for the packages
 	// that this file belongs to.
 	PackageHandles(ctx context.Context, fh FileHandle) ([]PackageHandle, error)
 
 	// GetActiveReverseDeps returns the active files belonging to the reverse
 	// dependencies of this file's package.
-	GetReverseDependencies(ctx context.Context, id string) ([]string, error)
+	GetReverseDependencies(ctx context.Context, id string) ([]PackageHandle, error)
 
-	// KnownImportPaths returns all the imported packages loaded in this snapshot,
+	// CachedImportPaths returns all the imported packages loaded in this snapshot,
 	// indexed by their import path.
-	KnownImportPaths() map[string]Package
+	CachedImportPaths(ctx context.Context) (map[string]Package, error)
 
 	// KnownPackages returns all the packages loaded in this snapshot.
 	KnownPackages(ctx context.Context) ([]PackageHandle, error)
+
+	// WorkspacePackages returns the PackageHandles for the snapshot's
+	// top-level packages.
+	WorkspacePackages(ctx context.Context) ([]PackageHandle, error)
 }
 
 // PackageHandle represents a handle to a specific version of a package.
@@ -118,10 +119,6 @@ type View interface {
 	// original one will be.
 	SetOptions(context.Context, Options) (View, error)
 
-	// WorkspacePackageIDs returns the ids of the packages at the top-level
-	// of the snapshot's view.
-	WorkspacePackageIDs(ctx context.Context) ([]string, error)
-
 	// Snapshot returns the current snapshot for the view.
 	Snapshot() Snapshot
 }