internal/lsp: fix concurrency issues in view initialization

Address issues from CL 254940.

Change-Id: If4839a8694979fd1951b3fad77abb10860b42b7c
Reviewed-on: https://go-review.googlesource.com/c/tools/+/255349
Trust: Rebecca Stambler <rstambler@golang.org>
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>

internal/lsp/cache/session.go

@@ -239,12 +239,13 @@ func (s *Session) createView(ctx context.Context, name string, folder span.URI,
 	// Initialize the view without blocking.
 	initCtx, initCancel := context.WithCancel(xcontext.Detach(ctx))
 	v.initCancelFirstAttempt = initCancel
+	snapshot := v.snapshot
+	release := snapshot.generation.Acquire(initCtx)
 	go func() {
-		release := v.snapshot.generation.Acquire(initCtx)
-		v.initialize(initCtx, v.snapshot, true)
+		v.initialize(initCtx, snapshot, true)
 		release()
 	}()
-	return v, v.snapshot, v.snapshot.generation.Acquire(ctx), nil
+	return v, snapshot, snapshot.generation.Acquire(ctx), nil
 }
 
 // findWorkspaceModules walks the view's root folder, looking for go.mod files.

internal/lsp/cache/snapshot.go

@@ -526,7 +526,7 @@ func (s *snapshot) KnownPackages(ctx context.Context) ([]source.Package, error)
 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.
-	s.view.AwaitInitialized(ctx)
+	s.AwaitInitialized(ctx)
 
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -705,7 +705,7 @@ func (s *snapshot) IsSaved(uri span.URI) bool {
 
 func (s *snapshot) awaitLoaded(ctx context.Context) error {
 	// Do not return results until the snapshot's view has been initialized.
-	s.view.AwaitInitialized(ctx)
+	s.AwaitInitialized(ctx)
 
 	if err := s.reloadWorkspace(ctx); err != nil {
 		return err
@@ -724,6 +724,17 @@ func (s *snapshot) awaitLoaded(ctx context.Context) error {
 	return nil
 }
 
+func (s *snapshot) AwaitInitialized(ctx context.Context) {
+	select {
+	case <-ctx.Done():
+		return
+	case <-s.view.initialized:
+	}
+	// We typically prefer to run something as intensive as the IWL without
+	// blocking. I'm not sure if there is a way to do that here.
+	s.view.initialize(ctx, s, false)
+}
+
 // reloadWorkspace reloads the metadata for all invalidated workspace packages.
 func (s *snapshot) reloadWorkspace(ctx context.Context) error {
 	// If the view's build configuration is invalid, we cannot reload by package path.
@@ -1202,7 +1213,7 @@ func (s *snapshot) findWorkspaceDirectories(ctx context.Context, modFH source.Fi
 }
 
 func (s *snapshot) BuiltinPackage(ctx context.Context) (*source.BuiltinPackage, error) {
-	s.view.AwaitInitialized(ctx)
+	s.AwaitInitialized(ctx)
 
 	if s.builtin == nil {
 		return nil, errors.Errorf("no builtin package for view %s", s.view.name)

internal/lsp/cache/view.go

@@ -677,15 +677,9 @@ func checkIgnored(suffix string) bool {
 }
 
 func (v *View) Snapshot(ctx context.Context) (source.Snapshot, func()) {
-	s := v.getSnapshot()
-	return s, s.generation.Acquire(ctx)
-}
-
-func (v *View) getSnapshot() *snapshot {
 	v.snapshotMu.Lock()
 	defer v.snapshotMu.Unlock()
-
-	return v.snapshot
+	return v.snapshot, v.snapshot.generation.Acquire(ctx)
 }
 
 func (v *View) initialize(ctx context.Context, s *snapshot, firstAttempt bool) {
@@ -743,18 +737,6 @@ func (v *View) initialize(ctx context.Context, s *snapshot, firstAttempt bool) {
 	})
 }
 
-// AwaitInitialized waits until a view is initialized
-func (v *View) AwaitInitialized(ctx context.Context) {
-	select {
-	case <-ctx.Done():
-		return
-	case <-v.initialized:
-	}
-	// We typically prefer to run something as intensive as the IWL without
-	// blocking. I'm not sure if there is a way to do that here.
-	v.initialize(ctx, v.getSnapshot(), false)
-}
-
 // 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.
@@ -767,7 +749,7 @@ func (v *View) invalidateContent(ctx context.Context, uris map[span.URI]source.V
 	v.cancelBackground()
 
 	// Do not clone a snapshot until its view has finished initializing.
-	v.AwaitInitialized(ctx)
+	v.snapshot.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/general.go

@@ -205,8 +205,11 @@ func (s *Server) addFolders(ctx context.Context, folders []protocol.WorkspaceFol
 			work.end(fmt.Sprintf("Error loading packages: %s", err))
 			continue
 		}
+		var swg sync.WaitGroup
+		swg.Add(1)
 		go func() {
-			view.AwaitInitialized(ctx)
+			defer swg.Done()
+			snapshot.AwaitInitialized(ctx)
 			work.end("Finished loading packages.")
 		}()
 
@@ -226,6 +229,7 @@ func (s *Server) addFolders(ctx context.Context, folders []protocol.WorkspaceFol
 		wg.Add(1)
 		go func() {
 			s.diagnoseDetached(snapshot)
+			swg.Wait()
 			release()
 			wg.Done()
 		}()

internal/lsp/source/view.go

@@ -40,6 +40,9 @@ type Snapshot interface {
 	// if it is not already part of the snapshot.
 	GetFile(ctx context.Context, uri span.URI) (VersionedFileHandle, error)
 
+	// AwaitInitialized waits until the snapshot's view is initialized.
+	AwaitInitialized(ctx context.Context)
+
 	// IsOpen returns whether the editor currently has a file open.
 	IsOpen(uri span.URI) bool
 
@@ -142,9 +145,6 @@ type View interface {
 	// Shutdown closes this view, and detaches it from its session.
 	Shutdown(ctx context.Context)
 
-	// AwaitInitialized waits until a view is initialized
-	AwaitInitialized(ctx context.Context)
-
 	// WriteEnv writes the view-specific environment to the io.Writer.
 	WriteEnv(ctx context.Context, w io.Writer) error