[gopls-release-branch.0.5] 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>
(cherry picked from commit c43c25c95a5031135f3c557359cce41997c8b2bc)
Reviewed-on: https://go-review.googlesource.com/c/tools/+/258286
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
Reviewed-by: Danish Dua <danishdua@google.com>
Trust: Rebecca Stambler <rstambler@golang.org>
diff --git a/internal/lsp/source/completion/completion.go b/internal/lsp/source/completion/completion.go
index 4b40607..5a3bbd5 100644
--- a/internal/lsp/source/completion/completion.go
+++ b/internal/lsp/source/completion/completion.go
@@ -179,6 +179,11 @@
 	// 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 @@
 
 	// 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.View().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 @@
 		}
 	}
 
-	return c.snapshot.View().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 @@
 	}
 
 	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
 
 	var mu sync.Mutex
 	add := func(pkgExport imports.PackageExport) {
@@ -1153,9 +1169,12 @@
 			cancel()
 		}
 	}
-	return c.snapshot.View().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 @@
 	}
 
 	ctx, cancel := context.WithCancel(ctx)
-	defer cancel()
 
 	var mu sync.Mutex
 	add := func(pkg imports.ImportFix) {
@@ -1454,9 +1472,11 @@
 		})
 		count++
 	}
-	return c.snapshot.View().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.
diff --git a/internal/lsp/source/completion/deep_completion.go b/internal/lsp/source/completion/deep_completion.go
index f2d2026..ec485f2 100644
--- a/internal/lsp/source/completion/deep_completion.go
+++ b/internal/lsp/source/completion/deep_completion.go
@@ -22,8 +22,8 @@
 	// 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 @@
 		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
+				}
 			}
 		}