internal/lsp: parallelize initial workspace load

The initial workspace load was happening when a view was created, in serial.
It should really just be kicked off in a separate goroutine once we create a
new view. Implementing this change required some other significant changes,
particularly the additional work being done by the WorkspacePackageIDs
method.

Some other changes had to be made while debugging. In particular, the
modification to the circular dependencies test was a consequence of
golang/go#36265.

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

internal/lsp/cache/check.go

@@ -124,7 +124,6 @@ func (s *snapshot) buildKey(ctx context.Context, id packageID, mode source.Parse
 		compiledGoFiles: compiledGoFiles,
 		mode:            mode,
 	}
-
 	// Make sure all of the depList are sorted.
 	depList := append([]packageID{}, m.deps...)
 	sort.Slice(depList, func(i, j int) bool {
@@ -137,7 +136,7 @@ func (s *snapshot) buildKey(ctx context.Context, id packageID, mode source.Parse
 	var depKeys [][]byte
 	for _, depID := range depList {
 		mode := source.ParseExported
-		if s.workspacePackages[depID] {
+		if s.isWorkspacePackage(depID) {
 			mode = source.ParseFull
 		}
 		depHandle, err := s.buildPackageHandle(ctx, depID, mode)

internal/lsp/cache/load.go

@@ -69,8 +69,6 @@ func (s *snapshot) load(ctx context.Context, scope interface{}) ([]*metadata, er
 		if err == nil {
 			err = errors.Errorf("no packages found for query %s", query)
 		}
-	}
-	if err != nil {
 		return nil, err
 	}
 	return s.updateMetadata(ctx, scope, pkgs, cfg)

internal/lsp/cache/session.go

@@ -84,6 +84,7 @@ func (s *session) createView(ctx context.Context, name string, folder span.URI,
 	}
 	v := &view{
 		session:       s,
+		initialized:   make(chan struct{}),
 		id:            strconv.FormatInt(index, 10),
 		options:       options,
 		baseCtx:       baseCtx,
@@ -116,26 +117,6 @@ func (s *session) createView(ctx context.Context, name string, folder span.URI,
 	// so we immediately add builtin.go to the list of ignored files.
 	v.buildBuiltinPackage(ctx)
 
-	// Preemptively load everything in this directory.
-	// TODO(matloob): Determine if this can be done in parallel with something else.
-	// Perhaps different calls to NewView can be run in parallel?
-	v.snapshotMu.Lock()
-	defer v.snapshotMu.Unlock() // The code after the snapshot is used isn't expensive.
-	m, err := v.snapshot.load(ctx, directoryURI(folder))
-	if err != nil {
-		// Suppress all errors.
-		log.Error(ctx, "failed to load snapshot", err, telemetry.Directory.Of(folder))
-		return v, v.snapshot, nil
-	}
-	// Prepare CheckPackageHandles for every package that's been loaded.
-	// (*snapshot).CheckPackageHandle makes the assumption that every package that's
-	// been loaded has an existing checkPackageHandle.
-	if _, err := v.snapshot.checkWorkspacePackages(ctx, m); err != nil {
-		// Suppress all errors.
-		log.Error(ctx, "failed to check snapshot", err, telemetry.Directory.Of(folder))
-		return v, v.snapshot, nil
-	}
-
 	debug.AddView(debugView{v})
 	return v, v.snapshot, nil
 }

internal/lsp/cache/snapshot.go

@@ -197,7 +197,12 @@ func (s *snapshot) shouldCheck(m []*metadata) (phs []*packageHandle, load, check
 	return phs, false, false
 }
 
-func (s *snapshot) GetReverseDependencies(id string) []string {
+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 {
+		return nil, err
+	}
+
 	ids := make(map[packageID]struct{})
 	s.transitiveReverseDependencies(packageID(id), ids)
 
@@ -208,7 +213,7 @@ func (s *snapshot) GetReverseDependencies(id string) []string {
 	for id := range ids {
 		results = append(results, string(id))
 	}
-	return results
+	return results, nil
 }
 
 // transitiveReverseDependencies populates the uris map with file URIs
@@ -217,8 +222,7 @@ func (s *snapshot) transitiveReverseDependencies(id packageID, ids map[packageID
 	if _, ok := ids[id]; ok {
 		return
 	}
-	m := s.getMetadata(id)
-	if m == nil {
+	if s.getMetadata(id) == nil {
 		return
 	}
 	ids[id] = struct{}{}
@@ -239,10 +243,26 @@ func (s *snapshot) getImportedByLocked(id packageID) []packageID {
 	if len(s.importedBy) == 0 {
 		s.rebuildImportGraph()
 	}
-
 	return s.importedBy[id]
 }
 
+func (s *snapshot) clearAndRebuildImportGraph() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	// Completely invalidate the original map.
+	s.importedBy = make(map[packageID][]packageID)
+	s.rebuildImportGraph()
+}
+
+func (s *snapshot) rebuildImportGraph() {
+	for id, m := range s.metadata {
+		for _, importID := range m.deps {
+			s.importedBy[importID] = append(s.importedBy[importID], id)
+		}
+	}
+}
+
 func (s *snapshot) addPackage(ph *packageHandle) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -256,24 +276,7 @@ func (s *snapshot) addPackage(ph *packageHandle) {
 	s.packages[ph.packageKey()] = ph
 }
 
-// checkWorkspacePackages checks the initial set of packages loaded when
-// the view is created. This is needed because
-// (*snapshot).CheckPackageHandle makes the assumption that every package that's
-// been loaded has an existing checkPackageHandle.
-func (s *snapshot) checkWorkspacePackages(ctx context.Context, m []*metadata) ([]source.PackageHandle, error) {
-	var phs []source.PackageHandle
-	for _, m := range m {
-		ph, err := s.buildPackageHandle(ctx, m.id, source.ParseFull)
-		if err != nil {
-			return nil, err
-		}
-		s.workspacePackages[m.id] = true
-		phs = append(phs, ph)
-	}
-	return phs, nil
-}
-
-func (s *snapshot) WorkspacePackageIDs(ctx context.Context) (ids []string) {
+func (s *snapshot) workspacePackageIDs() (ids []string) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -283,7 +286,10 @@ func (s *snapshot) WorkspacePackageIDs(ctx context.Context) (ids []string) {
 	return ids
 }
 
-func (s *snapshot) KnownPackages(ctx context.Context) []source.PackageHandle {
+func (s *snapshot) KnownPackages(ctx context.Context) ([]source.PackageHandle, error) {
+	if err := s.awaitInitialized(ctx); err != nil {
+		return nil, err
+	}
 	// Collect PackageHandles for all of the workspace packages first.
 	// They may need to be reloaded if their metadata has been invalidated.
 	wsPackages := make(map[packageID]bool)
@@ -324,8 +330,7 @@ func (s *snapshot) KnownPackages(ctx context.Context) []source.PackageHandle {
 		}
 		results = append(results, ph)
 	}
-
-	return results
+	return results, nil
 }
 
 func (s *snapshot) KnownImportPaths() map[string]source.Package {
@@ -450,6 +455,21 @@ func (s *snapshot) getIDs(uri span.URI) []packageID {
 	return s.ids[uri]
 }
 
+func (s *snapshot) isWorkspacePackage(id packageID) bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	_, ok := s.workspacePackages[id]
+	return ok
+}
+
+func (s *snapshot) setWorkspacePackage(id packageID) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.workspacePackages[id] = true
+}
+
 func (s *snapshot) getFileURIs() []span.URI {
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -553,7 +573,6 @@ func (s *snapshot) clone(ctx context.Context, withoutURI span.URI, withoutFileKi
 		if _, seen := transitiveIDs[id]; seen {
 			return
 		}
-
 		transitiveIDs[id] = struct{}{}
 		for _, rid := range s.getImportedByLocked(id) {
 			addRevDeps(rid)
@@ -630,20 +649,3 @@ func (s *snapshot) clone(ctx context.Context, withoutURI span.URI, withoutFileKi
 func (s *snapshot) ID() uint64 {
 	return s.id
 }
-
-func (s *snapshot) clearAndRebuildImportGraph() {
-	s.mu.Lock()
-	defer s.mu.Unlock()
-
-	// Completely invalidate the original map.
-	s.importedBy = make(map[packageID][]packageID)
-	s.rebuildImportGraph()
-}
-
-func (s *snapshot) rebuildImportGraph() {
-	for id, m := range s.metadata {
-		for _, importID := range m.deps {
-			s.importedBy[importID] = append(s.importedBy[importID], id)
-		}
-	}
-}

internal/lsp/cache/view.go

@@ -88,6 +88,12 @@ type view struct {
 	// ignoredURIs is the set of URIs of files that we ignore.
 	ignoredURIsMu sync.Mutex
 	ignoredURIs   map[span.URI]struct{}
+
+	// initialized is closed when we have attempted to load the view's workspace packages.
+	// If we failed to load initially, we don't re-try to avoid too many go/packages calls.
+	initializeOnce      sync.Once
+	initialized         chan struct{}
+	initializationError error
 }
 
 // modfiles holds the real and temporary go.mod files that are attributed to a view.
@@ -221,7 +227,6 @@ func (v *view) RunProcessEnvFunc(ctx context.Context, fn func(*imports.Options)
 			log.Print(context.Background(), "background refresh finished with err: ", tag.Of("err", nil))
 		}()
 	}
-
 	return nil
 }
 
@@ -359,6 +364,36 @@ func (v *view) getSnapshot() *snapshot {
 	return v.snapshot
 }
 
+func (v *view) WorkspacePackageIDs(ctx context.Context) ([]string, error) {
+	s := v.getSnapshot()
+
+	v.initializeOnce.Do(func() {
+		defer close(v.initialized)
+
+		// Do not cancel the call to go/packages.Load for the entire workspace.
+		meta, err := s.load(xcontext.Detach(ctx), directoryURI(v.folder))
+		if err != nil {
+			v.initializationError = err
+		}
+		for _, m := range meta {
+			s.setWorkspacePackage(m.id)
+		}
+	})
+	if v.initializationError != nil {
+		return nil, v.initializationError
+	}
+	return s.workspacePackageIDs(), nil
+}
+
+func (s *snapshot) awaitInitialized(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-s.view.initialized:
+	}
+	return s.view.initializationError
+}
+
 // invalidateContent invalidates the content of a Go file,
 // including any position and type information that depends on it.
 // It returns true if we were already tracking the given file, false otherwise.
@@ -373,6 +408,9 @@ 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
+
 	// This should be the only time we hold the view's snapshot lock for any period of time.
 	v.snapshotMu.Lock()
 	defer v.snapshotMu.Unlock()

internal/lsp/diagnostics.go

@@ -20,17 +20,22 @@ func (s *Server) diagnose(snapshot source.Snapshot, fh source.FileHandle) error
 	case source.Go:
 		go s.diagnoseFile(snapshot, fh)
 	case source.Mod:
-		go s.diagnoseSnapshot(snapshot)
+		ctx := snapshot.View().BackgroundContext()
+		go s.diagnoseSnapshot(ctx, snapshot)
 	}
 	return nil
 }
 
-func (s *Server) diagnoseSnapshot(snapshot source.Snapshot) {
-	ctx := snapshot.View().BackgroundContext()
+func (s *Server) diagnoseSnapshot(ctx context.Context, snapshot source.Snapshot) {
 	ctx, done := trace.StartSpan(ctx, "lsp:background-worker")
 	defer done()
 
-	for _, id := range snapshot.WorkspacePackageIDs(ctx) {
+	wsPackages, err := snapshot.View().WorkspacePackageIDs(ctx)
+	if err != nil {
+		log.Error(ctx, "diagnoseSnapshot: no workspace packages", err, telemetry.Directory.Of(snapshot.View().Folder))
+		return
+	}
+	for _, id := range wsPackages {
 		ph, err := snapshot.PackageHandle(ctx, id)
 		if err != nil {
 			log.Error(ctx, "diagnoseSnapshot: no PackageHandle for workspace package", err, telemetry.Package.Of(id))

internal/lsp/general.go

@@ -170,7 +170,7 @@ func (s *Server) addFolders(ctx context.Context, folders []protocol.WorkspaceFol
 			viewErrors[uri] = err
 			continue
 		}
-		go s.diagnoseSnapshot(snapshot)
+		go s.diagnoseSnapshot(ctx, snapshot)
 	}
 	if len(viewErrors) > 0 {
 		errMsg := fmt.Sprintf("Error loading workspace folders (expected %v, got %v)\n", len(folders), len(s.session.Views())-originalViews)

internal/lsp/lsp_test.go

@@ -53,7 +53,13 @@ func testLSP(t *testing.T, exporter packagestest.Exporter) {
 	options := tests.DefaultOptions()
 	session.SetOptions(options)
 	options.Env = data.Config.Env
-	if _, _, err := session.NewView(ctx, viewName, span.FileURI(data.Config.Dir), options); err != nil {
+	view, _, err := session.NewView(ctx, viewName, span.FileURI(data.Config.Dir), options)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Load the workspace packages, since this would normally happen when a view is initialized.
+	// Otherwise, tests that need to look at all workspace packages will fail.
+	if _, err := view.WorkspacePackageIDs(ctx); err != nil {
 		t.Fatal(err)
 	}
 	for filename, content := range data.Config.Overlay {
@@ -79,7 +85,6 @@ func testLSP(t *testing.T, exporter packagestest.Exporter) {
 		data: data,
 		ctx:  ctx,
 	}
-
 	tests.Run(t, r, data)
 }
 

internal/lsp/references.go

@@ -9,6 +9,7 @@ import (
 
 	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/lsp/source"
+	"golang.org/x/tools/internal/lsp/telemetry"
 	"golang.org/x/tools/internal/span"
 	"golang.org/x/tools/internal/telemetry/log"
 	"golang.org/x/tools/internal/telemetry/tag"
@@ -46,7 +47,7 @@ func (s *Server) references(ctx context.Context, params *protocol.ReferenceParam
 			if err == source.ErrNoIdentFound {
 				return nil, err
 			}
-			log.Error(ctx, "no identifier", err, tag.Of("Identifier", ident.Name))
+			log.Error(ctx, "no identifier", err, telemetry.URI.Of(uri))
 			continue
 		}
 

internal/lsp/source/diagnostics.go

@@ -92,7 +92,12 @@ func Diagnostics(ctx context.Context, snapshot Snapshot, fh FileHandle, withAnal
 		}
 	}
 	// Updates to the diagnostics for this package may need to be propagated.
-	for _, id := range snapshot.GetReverseDependencies(pkg.ID()) {
+	reverseDeps, err := snapshot.GetReverseDependencies(ctx, pkg.ID())
+	if err != nil {
+		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

internal/lsp/source/implementation.go

@@ -55,7 +55,6 @@ func Implementation(ctx context.Context, s Snapshot, f FileHandle, pp protocol.P
 var ErrNotAnInterface = errors.New("not an interface or interface method")
 
 func implementations(ctx context.Context, s Snapshot, f FileHandle, pp protocol.Position) ([]implementation, error) {
-
 	var (
 		impls []implementation
 		seen  = make(map[token.Position]bool)
@@ -97,13 +96,16 @@ func implementations(ctx context.Context, s Snapshot, f FileHandle, pp protocol.
 			allNamed []*types.Named
 			pkgs     = make(map[*types.Package]Package)
 		)
-		for _, ph := range s.KnownPackages(ctx) {
+		knownPkgs, err := s.KnownPackages(ctx)
+		if err != nil {
+			return nil, err
+		}
+		for _, ph := range knownPkgs {
 			pkg, err := ph.Check(ctx)
 			if err != nil {
 				return nil, err
 			}
 			pkgs[pkg.GetTypes()] = pkg
-
 			info := pkg.GetTypesInfo()
 			for _, obj := range info.Defs {
 				obj, ok := obj.(*types.TypeName)

internal/lsp/source/references.go

@@ -43,7 +43,11 @@ func (i *IdentifierInfo) References(ctx context.Context) ([]*ReferenceInfo, erro
 	var searchpkgs []Package
 	if i.Declaration.obj.Exported() {
 		// Only search all packages if the identifier is exported.
-		for _, id := range i.Snapshot.GetReverseDependencies(i.pkg.ID()) {
+		reverseDeps, err := i.Snapshot.GetReverseDependencies(ctx, i.pkg.ID())
+		if err != nil {
+			return nil, err
+		}
+		for _, id := range reverseDeps {
 			ph, err := i.Snapshot.PackageHandle(ctx, id)
 			if err != nil {
 				log.Error(ctx, "References: no CheckPackageHandle", err, telemetry.Package.Of(id))

internal/lsp/source/source_test.go

@@ -60,6 +60,11 @@ func testSource(t *testing.T, exporter packagestest.Exporter) {
 		data: data,
 		ctx:  ctx,
 	}
+	// Load the workspace packages, since this would normally happen when a view is initialized.
+	// Otherwise, tests that need to look at all workspace packages will fail.
+	if _, err := view.WorkspacePackageIDs(ctx); err != nil {
+		t.Fatal(err)
+	}
 	for filename, content := range data.Config.Overlay {
 		kind := source.DetectLanguage("", filename)
 		if kind != source.Go {

internal/lsp/source/view.go

@@ -46,20 +46,16 @@ type Snapshot interface {
 	// that this file belongs to.
 	PackageHandles(ctx context.Context, fh FileHandle) ([]PackageHandle, error)
 
-	// WorkspacePackageIDs returns the ids of the packages at the top-level
-	// of the snapshot's view.
-	WorkspacePackageIDs(ctx context.Context) []string
-
 	// GetActiveReverseDeps returns the active files belonging to the reverse
 	// dependencies of this file's package.
-	GetReverseDependencies(id string) []string
+	GetReverseDependencies(ctx context.Context, id string) ([]string, error)
 
 	// KnownImportPaths returns all the imported packages loaded in this snapshot,
 	// indexed by their import path.
 	KnownImportPaths() map[string]Package
 
 	// KnownPackages returns all the packages loaded in this snapshot.
-	KnownPackages(ctx context.Context) []PackageHandle
+	KnownPackages(ctx context.Context) ([]PackageHandle, error)
 }
 
 // PackageHandle represents a handle to a specific version of a package.
@@ -131,6 +127,10 @@ type View interface {
 	// the given package or one of its dependencies.
 	FindMapperInPackage(pkg Package, uri span.URI) (*protocol.ColumnMapper, 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
 }

internal/lsp/workspace.go

@@ -48,10 +48,11 @@ func (s *Server) updateConfiguration(ctx context.Context, changed interface{}) e
 		if err := s.fetchConfig(ctx, view.Name(), view.Folder(), &options); err != nil {
 			return err
 		}
-		if _, err := view.SetOptions(ctx, options); err != nil {
+		view, err := view.SetOptions(ctx, options)
+		if err != nil {
 			return err
 		}
-		go s.diagnoseSnapshot(view.Snapshot())
+		go s.diagnoseSnapshot(ctx, view.Snapshot())
 	}
 	return nil
 }