internal/lsp/source: run deep completions before unimported completions

Unimported completions are expensive and can use up a large portion of
completion budget just to find initial deep search candidates. This
change moves these expensive operations which search through the module
cache to after normal deep completions so we search through more useful
candidates first.

Fixes golang/go#41434
Fixes golang/go#41665

Change-Id: I6f3963f8c65c1a97833a35738d2e96420de2f6ee
Reviewed-on: https://go-review.googlesource.com/c/tools/+/257974
Run-TryBot: Danish Dua <danishdua@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>
Trust: Danish Dua <danishdua@google.com>

Author: dandua98

internal/lsp/source/completion/completion.go

@@ -179,6 +179,11 @@ type completer struct {
 	// items is the list of completion items returned.
 	items []CompletionItem
 
+	// completionCallbacks is a list of callbacks to collect completions that
+	// require expensive operations. This includes operations where we search
+	// through the entire module cache.
+	completionCallbacks []func(opts *imports.Options) error
+
 	// surrounding describes the identifier surrounding the position.
 	surrounding *Selection
 
@@ -516,6 +521,17 @@ func Completion(ctx context.Context, snapshot source.Snapshot, fh source.FileHan
 
 	// Deep search collected candidates and their members for more candidates.
 	c.deepSearch(ctx)
+	c.deepState.searchQueue = nil
+
+	for _, callback := range c.completionCallbacks {
+		if err := c.snapshot.RunProcessEnvFunc(ctx, callback); err != nil {
+			return nil, nil, err
+		}
+	}
+
+	// Search candidates populated by expensive operations like
+	// unimportedMembers etc. for more completion items.
+	c.deepSearch(ctx)
 
 	// Statement candidates offer an entire statement in certain contexts, as
 	// opposed to a single object. Add statement candidates last because they
@@ -799,9 +815,10 @@ func (c *completer) populateImportCompletions(ctx context.Context, searchImport
 		}
 	}
 
-	return c.snapshot.RunProcessEnvFunc(ctx, func(opts *imports.Options) error {
+	c.completionCallbacks = append(c.completionCallbacks, func(opts *imports.Options) error {
 		return imports.GetImportPaths(ctx, searchImports, prefix, c.filename, c.pkg.GetTypes().Name(), opts.Env)
 	})
+	return nil
 }
 
 // populateCommentCompletions yields completions for comments preceding or in declarations.
@@ -1126,7 +1143,6 @@ func (c *completer) unimportedMembers(ctx context.Context, id *ast.Ident) error
 	}
 
 	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
 
 	var mu sync.Mutex
 	add := func(pkgExport imports.PackageExport) {
@@ -1153,9 +1169,12 @@ func (c *completer) unimportedMembers(ctx context.Context, id *ast.Ident) error
 			cancel()
 		}
 	}
-	return c.snapshot.RunProcessEnvFunc(ctx, func(opts *imports.Options) error {
+
+	c.completionCallbacks = append(c.completionCallbacks, func(opts *imports.Options) error {
+		defer cancel()
 		return imports.GetPackageExports(ctx, add, id.Name, c.filename, c.pkg.GetTypes().Name(), opts.Env)
 	})
+	return nil
 }
 
 // unimportedScore returns a score for an unimported package that is generally
@@ -1422,7 +1441,6 @@ func (c *completer) unimportedPackages(ctx context.Context, seen map[string]stru
 	}
 
 	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
 
 	var mu sync.Mutex
 	add := func(pkg imports.ImportFix) {
@@ -1454,9 +1472,11 @@ func (c *completer) unimportedPackages(ctx context.Context, seen map[string]stru
 		})
 		count++
 	}
-	return c.snapshot.RunProcessEnvFunc(ctx, func(opts *imports.Options) error {
+	c.completionCallbacks = append(c.completionCallbacks, func(opts *imports.Options) error {
+		defer cancel()
 		return imports.GetAllCandidates(ctx, add, prefix, c.filename, c.pkg.GetTypes().Name(), opts.Env)
 	})
+	return nil
 }
 
 // alreadyImports reports whether f has an import with the specified path.

internal/lsp/source/completion/deep_completion.go

@@ -22,8 +22,8 @@ type deepCompletionState struct {
 	// enabled indicates wether deep completion is permitted.
 	enabled bool
 
-	// queueClosed is used to disable adding new items to search queue once
-	// we're running out of our time budget.
+	// queueClosed is used to disable adding new sub-fields to search queue
+	// once we're running out of our time budget.
 	queueClosed bool
 
 	// searchQueue holds the current breadth first search queue.
@@ -158,14 +158,16 @@ outer:
 		c.deepState.candidateCount++
 		if c.opts.budget > 0 && c.deepState.candidateCount%100 == 0 {
 			spent := float64(time.Since(c.startTime)) / float64(c.opts.budget)
-			if spent > 1.0 {
+			select {
+			case <-ctx.Done():
 				return
-			}
-			// If we are almost out of budgeted time, no further elements
-			// should be added to the queue. This ensures remaining time is
-			// used for processing current queue.
-			if !c.deepState.queueClosed && spent >= 0.85 {
-				c.deepState.queueClosed = true
+			default:
+				// If we are almost out of budgeted time, no further elements
+				// should be added to the queue. This ensures remaining time is
+				// used for processing current queue.
+				if !c.deepState.queueClosed && spent >= 0.85 {
+					c.deepState.queueClosed = true
+				}
 			}
 		}