internal/imports: redesign scan API

We have multiple use cases for scanning: goimports, import completion,
and unimported completions. All three need slightly different features,
and the latter have very different performance considerations. Scanning
everything all at once and returning it was not good enough for them.

Instead, design the API as a series of callbacks for each
directory/package: first we discover its existence, then we load its
package name, then we load its exports. At each step the caller can
choose whether to proceed with the package. Import completion can stop
before loading exports, goimports can apply its directory name
heuristics, and in the future we'll be able to stop the scan short once
we've found all the results we want for completions.

I don't intend any significant changes here but there may be some little
ones around the edges.

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

internal/imports/fix.go

@@ -82,6 +82,7 @@ type ImportFix struct {
 	IdentName string
 	// FixType is the type of fix this is (AddImport, DeleteImport, SetImportName).
 	FixType   ImportFixType
+	relevance int // see pkg
 }
 
 // An ImportInfo represents a single import statement.
@@ -585,61 +586,59 @@ func getFixes(fset *token.FileSet, f *ast.File, filename string, env *ProcessEnv
 }
 
 // getCandidatePkgs returns the list of pkgs that are accessible from filename,
-// optionall filtered to only packages named pkgName.
+// filtered to those that match pkgnameFilter.
-func getCandidatePkgs(pkgName, filename string, env *ProcessEnv) ([]*pkg, error) {
+func getCandidatePkgs(ctx context.Context, wrappedCallback *scanCallback, filename string, env *ProcessEnv) error {
 	// TODO(heschi): filter out current package. (Don't forget x_test can import x.)
 
-	var result []*pkg
 	// Start off with the standard library.
-	for importPath := range stdlib {
+	for importPath, exports := range stdlib {
-		if pkgName != "" && path.Base(importPath) != pkgName {
+		p := &pkg{
-			continue
-		}
-		result = append(result, &pkg{
 			dir:             filepath.Join(env.GOROOT, "src", importPath),
 			importPathShort: importPath,
 			packageName:     path.Base(importPath),
 			relevance:       0,
-		})
+		}
+		if wrappedCallback.packageNameLoaded(p) {
+			wrappedCallback.exportsLoaded(p, exports)
+		}
 	}
 
 	// Exclude goroot results -- getting them is relatively expensive, not cached,
 	// and generally redundant with the in-memory version.
 	exclude := []gopathwalk.RootType{gopathwalk.RootGOROOT}
-	// Only the go/packages resolver uses the first argument, and nobody uses that resolver.
-	scannedPkgs, err := env.GetResolver().scan(nil, true, exclude)
-	if err != nil {
-		return nil, err
-	}
 
+	var mu sync.Mutex
 	dupCheck := map[string]struct{}{}
-	for _, pkg := range scannedPkgs {
+
-		if pkgName != "" && pkg.packageName != pkgName {
+	scanFilter := &scanCallback{
-			continue
+		dirFound: func(pkg *pkg) bool {
-		}
+			return canUse(filename, pkg.dir) && wrappedCallback.dirFound(pkg)
-		if !canUse(filename, pkg.dir) {
+		},
-			continue
+		packageNameLoaded: func(pkg *pkg) bool {
-		}
+			mu.Lock()
+			defer mu.Unlock()
 			if _, ok := dupCheck[pkg.importPathShort]; ok {
-			continue
+				return false
 			}
 			dupCheck[pkg.importPathShort] = struct{}{}
-		result = append(result, pkg)
+			return wrappedCallback.packageNameLoaded(pkg)
+		},
+		exportsLoaded: func(pkg *pkg, exports []string) {
+			wrappedCallback.exportsLoaded(pkg, exports)
+		},
+	}
+	return env.GetResolver().scan(ctx, scanFilter, exclude)
 }
 
-	// Sort first by relevance, then by package name, with import path as a tiebreaker.
+// Compare first by relevance, then by package name, with import path as a tiebreaker.
-	sort.Slice(result, func(i, j int) bool {
+func compareFix(fi, fj *ImportFix) bool {
-		pi, pj := result[i], result[j]
+	if fi.relevance != fj.relevance {
-		if pi.relevance != pj.relevance {
+		return fi.relevance < fj.relevance
-			return pi.relevance < pj.relevance
 	}
-		if pi.packageName != pj.packageName {
+	if fi.IdentName != fj.IdentName {
-			return pi.packageName < pj.packageName
+		return fi.IdentName < fj.IdentName
 	}
-		return pi.importPathShort < pj.importPathShort
+	return fi.StmtInfo.ImportPath < fj.StmtInfo.ImportPath
-	})
-
-	return result, nil
 }
 
 func candidateImportName(pkg *pkg) string {
@@ -650,23 +649,39 @@ func candidateImportName(pkg *pkg) string {
 }
 
 // getAllCandidates gets all of the candidates to be imported, regardless of if they are needed.
-func getAllCandidates(filename string, env *ProcessEnv) ([]ImportFix, error) {
+func getAllCandidates(ctx context.Context, prefix string, filename string, env *ProcessEnv) ([]ImportFix, error) {
-	pkgs, err := getCandidatePkgs("", filename, env)
+	var mu sync.Mutex
-	if err != nil {
+	var results []ImportFix
-		return nil, err
+	filter := &scanCallback{
-	}
+		dirFound: func(pkg *pkg) bool {
-	result := make([]ImportFix, 0, len(pkgs))
+			// TODO(heschi): apply dir match heuristics like pkgIsCandidate
-	for _, pkg := range pkgs {
+			return true
-		result = append(result, ImportFix{
+		},
+		packageNameLoaded: func(pkg *pkg) bool {
+			if strings.HasPrefix(pkg.packageName, prefix) {
+				mu.Lock()
+				defer mu.Unlock()
+				results = append(results, ImportFix{
 					StmtInfo: ImportInfo{
 						ImportPath: pkg.importPathShort,
 						Name:       candidateImportName(pkg),
 					},
 					IdentName: pkg.packageName,
 					FixType:   AddImport,
+					relevance: pkg.relevance,
 				})
 			}
-	return result, nil
+			return false
+		},
+	}
+	err := getCandidatePkgs(ctx, filter, filename, env)
+	if err != nil {
+		return nil, err
+	}
+	sort.Slice(results, func(i, j int) bool {
+		return compareFix(&results[i], &results[j])
+	})
+	return results, nil
 }
 
 // A PackageExport is a package and its exports.
@@ -675,41 +690,42 @@ type PackageExport struct {
 	Exports []string
 }
 
-func getPackageExports(completePackage, filename string, env *ProcessEnv) ([]PackageExport, error) {
+func getPackageExports(ctx context.Context, completePackage, filename string, env *ProcessEnv) ([]PackageExport, error) {
-	pkgs, err := getCandidatePkgs(completePackage, filename, env)
+	var mu sync.Mutex
-	if err != nil {
+	var results []PackageExport
-		return nil, err
+	callback := &scanCallback{
-	}
+		dirFound: func(pkg *pkg) bool {
-
+			// TODO(heschi): apply dir match heuristics like pkgIsCandidate
-	results := make([]PackageExport, 0, len(pkgs))
+			return true
-	for _, pkg := range pkgs {
+		},
-		fix := &ImportFix{
+		packageNameLoaded: func(pkg *pkg) bool {
+			return pkg.packageName == completePackage
+		},
+		exportsLoaded: func(pkg *pkg, exports []string) {
+			mu.Lock()
+			defer mu.Unlock()
+			sort.Strings(exports)
+			results = append(results, PackageExport{
+				Fix: &ImportFix{
 					StmtInfo: ImportInfo{
 						ImportPath: pkg.importPathShort,
 						Name:       candidateImportName(pkg),
 					},
 					IdentName: pkg.packageName,
 					FixType:   AddImport,
-		}
+					relevance: pkg.relevance,
-		var exports []string
+				},
-		if e, ok := stdlib[pkg.importPathShort]; ok {
-			exports = e
-		} else {
-			exports, err = loadExportsForPackage(context.Background(), env, completePackage, pkg)
-			if err != nil {
-				if env.Debug {
-					env.Logf("while completing %q, error loading exports from %q: %v", completePackage, pkg.importPathShort, err)
-				}
-				continue
-			}
-		}
-		sort.Strings(exports)
-		results = append(results, PackageExport{
-			Fix:     fix,
 				Exports: exports,
 			})
+		},
 	}
-
+	err := getCandidatePkgs(ctx, callback, filename, env)
+	if err != nil {
+		return nil, err
+	}
+	sort.Slice(results, func(i, j int) bool {
+		return compareFix(results[i].Fix, results[j].Fix)
+	})
 	return results, nil
 }
 
@@ -839,10 +855,8 @@ func addStdlibCandidates(pass *pass, refs references) {
 type Resolver interface {
 	// loadPackageNames loads the package names in importPaths.
 	loadPackageNames(importPaths []string, srcDir string) (map[string]string, error)
-	// scan finds (at least) the packages satisfying refs. If loadNames is true,
+	// scan works with callback to search for packages. See scanCallback for details.
-	// package names will be set on the results, and dirs whose package name
+	scan(ctx context.Context, callback *scanCallback, exclude []gopathwalk.RootType) error
-	// could not be determined will be excluded.
-	scan(refs references, loadNames bool, exclude []gopathwalk.RootType) ([]*pkg, error)
 	// loadExports returns the set of exported symbols in the package at dir.
 	// loadExports may be called concurrently.
 	loadExports(ctx context.Context, pkg *pkg) (string, []string, error)
@@ -850,8 +864,47 @@ type Resolver interface {
 	ClearForNewScan()
 }
 
+// A scanCallback controls a call to scan and receives its results.
+// In general, minor errors will be silently discarded; a user should not
+// expect to receive a full series of calls for everything.
+type scanCallback struct {
+	// dirFound is called when a directory is found that is possibly a Go package.
+	// pkg will be populated with everything except packageName.
+	// If it returns true, the package's name will be loaded.
+	dirFound func(pkg *pkg) bool
+	// packageNameLoaded is called when a package is found and its name is loaded.
+	// If it returns true, the package's exports will be loaded.
+	packageNameLoaded func(pkg *pkg) bool
+	// exportsLoaded is called when a package's exports have been loaded.
+	exportsLoaded func(pkg *pkg, exports []string)
+}
+
 func addExternalCandidates(pass *pass, refs references, filename string) error {
-	dirScan, err := pass.env.GetResolver().scan(refs, false, nil)
+	var mu sync.Mutex
+	found := make(map[string][]pkgDistance)
+	callback := &scanCallback{
+		dirFound: func(pkg *pkg) bool {
+			return pkgIsCandidate(filename, refs, pkg)
+		},
+		packageNameLoaded: func(pkg *pkg) bool {
+			if _, want := refs[pkg.packageName]; !want {
+				return false
+			}
+			if pkg.dir == pass.srcDir && pass.f.Name.Name == pkg.packageName {
+				// The candidate is in the same directory and has the
+				// same package name. Don't try to import ourselves.
+				return false
+			}
+			if !canUse(filename, pkg.dir) {
+				return false
+			}
+			mu.Lock()
+			defer mu.Unlock()
+			found[pkg.packageName] = append(found[pkg.packageName], pkgDistance{pkg, distance(pass.srcDir, pkg.dir)})
+			return false // We'll do our own loading after we sort.
+		},
+	}
+	err := pass.env.GetResolver().scan(context.Background(), callback, nil)
 	if err != nil {
 		return err
 	}
@@ -878,7 +931,7 @@ func addExternalCandidates(pass *pass, refs references, filename string) error {
 		go func(pkgName string, symbols map[string]bool) {
 			defer wg.Done()
 
-			found, err := findImport(ctx, pass, dirScan, pkgName, symbols, filename)
+			found, err := findImport(ctx, pass, found[pkgName], pkgName, symbols, filename)
 
 			if err != nil {
 				firstErrOnce.Do(func() {
@@ -1093,7 +1146,7 @@ func distance(basepath, targetpath string) int {
 	return strings.Count(p, string(filepath.Separator)) + 1
 }
 
-func (r *gopathResolver) scan(_ references, loadNames bool, exclude []gopathwalk.RootType) ([]*pkg, error) {
+func (r *gopathResolver) scan(ctx context.Context, callback *scanCallback, exclude []gopathwalk.RootType) error {
 	r.init()
 	add := func(root gopathwalk.Root, dir string) {
 		// We assume cached directories have not changed. We can skip them and their
@@ -1113,32 +1166,36 @@ func (r *gopathResolver) scan(_ references, loadNames bool, exclude []gopathwalk
 	}
 	roots := filterRoots(gopathwalk.SrcDirsRoots(r.env.buildContext()), exclude)
 	gopathwalk.Walk(roots, add, gopathwalk.Options{Debug: r.env.Debug, ModulesEnabled: false})
-	var result []*pkg
 	for _, dir := range r.cache.Keys() {
 		info, ok := r.cache.Load(dir)
 		if !ok {
 			continue
 		}
-		if loadNames {
-			var err error
-			info, err = r.cache.CachePackageName(info)
-			if err != nil {
-				continue
-			}
-		}
 
 		p := &pkg{
 			importPathShort: info.nonCanonicalImportPath,
 			dir:             dir,
 			relevance:       1,
-			packageName:     info.packageName,
 		}
 		if info.rootType == gopathwalk.RootGOROOT {
 			p.relevance = 0
 		}
-		result = append(result, p)
+
+		if callback.dirFound(p) {
+			var err error
+			p.packageName, err = r.cache.CachePackageName(info)
+			if err != nil {
+				continue
 			}
-	return result, nil
+		}
+
+		if callback.packageNameLoaded(p) {
+			if _, exports, err := r.loadExports(ctx, p); err == nil {
+				callback.exportsLoaded(p, exports)
+			}
+		}
+	}
+	return nil
 }
 
 func filterRoots(roots []gopathwalk.Root, exclude []gopathwalk.RootType) []gopathwalk.Root {
@@ -1238,29 +1295,7 @@ func loadExportsFromFiles(ctx context.Context, env *ProcessEnv, dir string) (str
 
 // findImport searches for a package with the given symbols.
 // If no package is found, findImport returns ("", false, nil)
-func findImport(ctx context.Context, pass *pass, dirScan []*pkg, pkgName string, symbols map[string]bool, filename string) (*pkg, error) {
+func findImport(ctx context.Context, pass *pass, candidates []pkgDistance, pkgName string, symbols map[string]bool, filename string) (*pkg, error) {
-	pkgDir, err := filepath.Abs(filename)
-	if err != nil {
-		return nil, err
-	}
-	pkgDir = filepath.Dir(pkgDir)
-
-	// Find candidate packages, looking only at their directory names first.
-	var candidates []pkgDistance
-	for _, pkg := range dirScan {
-		if pkg.dir == pkgDir && pass.f.Name.Name == pkgName {
-			// The candidate is in the same directory and has the
-			// same package name. Don't try to import ourselves.
-			continue
-		}
-		if pkgIsCandidate(filename, pkgName, pkg) {
-			candidates = append(candidates, pkgDistance{
-				pkg:      pkg,
-				distance: distance(pkgDir, pkg.dir),
-			})
-		}
-	}
-
 	// Sort the candidates by their import package length,
 	// assuming that shorter package names are better than long
 	// ones.  Note that this sorts by the de-vendored name, so
@@ -1273,7 +1308,6 @@ func findImport(ctx context.Context, pass *pass, dirScan []*pkg, pkgName string,
 	}
 
 	// Collect exports for packages with matching names.
-
 	rescv := make([]chan *pkg, len(candidates))
 	for i := range candidates {
 		rescv[i] = make(chan *pkg, 1)
@@ -1368,7 +1402,7 @@ func loadExportsForPackage(ctx context.Context, env *ProcessEnv, expectPkg strin
 // filename is the file being formatted.
 // pkgIdent is the package being searched for, like "client" (if
 // searching for "client.New")
-func pkgIsCandidate(filename, pkgIdent string, pkg *pkg) bool {
+func pkgIsCandidate(filename string, refs references, pkg *pkg) bool {
 	// Check "internal" and "vendor" visibility:
 	if !canUse(filename, pkg.dir) {
 		return false
@@ -1386,6 +1420,7 @@ func pkgIsCandidate(filename, pkgIdent string, pkg *pkg) bool {
 	// "bar", which is strongly discouraged
 	// anyway. There's no reason goimports needs
 	// to be slow just to accommodate that.
+	for pkgIdent := range refs {
 		lastTwo := lastTwoComponents(pkg.importPathShort)
 		if strings.Contains(lastTwo, pkgIdent) {
 			return true
@@ -1396,7 +1431,7 @@ func pkgIsCandidate(filename, pkgIdent string, pkg *pkg) bool {
 				return true
 			}
 		}
-
+	}
 	return false
 }
 

internal/imports/fix_test.go

@@ -5,6 +5,7 @@
 package imports
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"go/build"
@@ -2310,7 +2311,8 @@ func TestPkgIsCandidate(t *testing.T) {
 	}
 	for i, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got := pkgIsCandidate(tt.filename, tt.pkgIdent, tt.pkg)
+			refs := references{tt.pkgIdent: nil}
+			got := pkgIsCandidate(tt.filename, refs, tt.pkg)
 			if got != tt.want {
 				t.Errorf("test %d. pkgIsCandidate(%q, %q, %+v) = %v; want %v",
 					i, tt.filename, tt.pkgIdent, *tt.pkg, got, tt.want)
@@ -2513,7 +2515,7 @@ func TestGetCandidates(t *testing.T) {
 			},
 		},
 	}.test(t, func(t *goimportTest) {
-		candidates, err := getAllCandidates("x.go", t.env)
+		candidates, err := getAllCandidates(context.Background(), "", "x.go", t.env)
 		if err != nil {
 			t.Fatalf("GetAllCandidates() = %v", err)
 		}
@@ -2549,7 +2551,7 @@ func TestGetPackageCompletions(t *testing.T) {
 			},
 		},
 	}.test(t, func(t *goimportTest) {
-		candidates, err := getPackageExports("rand", "x.go", t.env)
+		candidates, err := getPackageExports(context.Background(), "rand", "x.go", t.env)
 		if err != nil {
 			t.Fatalf("getPackageCompletions() = %v", err)
 		}

internal/imports/imports.go

@@ -11,6 +11,7 @@ package imports
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"fmt"
 	"go/ast"
 	"go/build"
@@ -115,23 +116,23 @@ func ApplyFixes(fixes []*ImportFix, filename string, src []byte, opt *Options, e
 	return formatFile(fileSet, file, src, nil, opt)
 }
 
-// GetAllCandidates gets all of the standard library candidate packages to import in
+// GetAllCandidates gets all of the packages starting with prefix that can be
-// sorted order on import path.
+// imported by filename, sorted by import path.
-func GetAllCandidates(filename string, opt *Options) (pkgs []ImportFix, err error) {
+func GetAllCandidates(ctx context.Context, prefix string, filename string, opt *Options) (pkgs []ImportFix, err error) {
 	_, opt, err = initialize(filename, nil, opt)
 	if err != nil {
 		return nil, err
 	}
-	return getAllCandidates(filename, opt.Env)
+	return getAllCandidates(ctx, prefix, filename, opt.Env)
 }
 
 // GetPackageExports returns all known packages with name pkg and their exports.
-func GetPackageExports(pkg, filename string, opt *Options) (exports []PackageExport, err error) {
+func GetPackageExports(ctx context.Context, pkg, filename string, opt *Options) (exports []PackageExport, err error) {
 	_, opt, err = initialize(filename, nil, opt)
 	if err != nil {
 		return nil, err
 	}
-	return getPackageExports(pkg, filename, opt.Env)
+	return getPackageExports(ctx, pkg, filename, opt.Env)
 }
 
 // initialize sets the values for opt and src.

internal/imports/mod.go

@@ -13,7 +13,6 @@ import (
 	"sort"
 	"strconv"
 	"strings"
-	"sync"
 
 	"golang.org/x/tools/internal/gopathwalk"
 	"golang.org/x/tools/internal/module"
@@ -212,7 +211,7 @@ func (r *ModuleResolver) cacheKeys() []string {
 }
 
 // cachePackageName caches the package name for a dir already in the cache.
-func (r *ModuleResolver) cachePackageName(info directoryPackageInfo) (directoryPackageInfo, error) {
+func (r *ModuleResolver) cachePackageName(info directoryPackageInfo) (string, error) {
 	if info.rootType == gopathwalk.RootModuleCache {
 		return r.moduleCacheCache.CachePackageName(info)
 	}
@@ -334,9 +333,9 @@ func (r *ModuleResolver) loadPackageNames(importPaths []string, srcDir string) (
 	return names, nil
 }
 
-func (r *ModuleResolver) scan(_ references, loadNames bool, exclude []gopathwalk.RootType) ([]*pkg, error) {
+func (r *ModuleResolver) scan(ctx context.Context, callback *scanCallback, exclude []gopathwalk.RootType) error {
 	if err := r.init(); err != nil {
-		return nil, err
+		return err
 	}
 
 	// Walk GOROOT, GOPATH/pkg/mod, and the main module.
@@ -360,15 +359,9 @@ func (r *ModuleResolver) scan(_ references, loadNames bool, exclude []gopathwalk
 
 	roots = filterRoots(roots, exclude)
 
-	var result []*pkg
-	var mu sync.Mutex
-
 	// We assume cached directories have not changed. We can skip them and their
 	// children.
 	skip := func(root gopathwalk.Root, dir string) bool {
-		mu.Lock()
-		defer mu.Unlock()
-
 		info, ok := r.cacheLoad(dir)
 		if !ok {
 			return false
@@ -382,9 +375,6 @@ func (r *ModuleResolver) scan(_ references, loadNames bool, exclude []gopathwalk
 
 	// Add anything new to the cache. We'll process everything in it below.
 	add := func(root gopathwalk.Root, dir string) {
-		mu.Lock()
-		defer mu.Unlock()
-
 		r.cacheStore(r.scanDirForPackage(root, dir))
 	}
 
@@ -402,22 +392,28 @@ func (r *ModuleResolver) scan(_ references, loadNames bool, exclude []gopathwalk
 			continue
 		}
 
-		// If we want package names, make sure the cache has them.
+		pkg, err := r.canonicalize(info)
-		if loadNames {
+		if err != nil {
-			var err error
+			continue
-			if info, err = r.cachePackageName(info); err != nil {
+		}
-				continue
+
-			}
+		if callback.dirFound(pkg) {
-		}
+			var err error
-
+			pkg.packageName, err = r.cachePackageName(info)
-		res, err := r.canonicalize(info)
 			if err != nil {
 				continue
 			}
-		result = append(result, res)
 		}
 
-	return result, nil
+		if callback.packageNameLoaded(pkg) {
+			_, exports, err := r.loadExports(ctx, pkg)
+			if err != nil {
+				continue
+			}
+			callback.exportsLoaded(pkg, exports)
+		}
+	}
+	return nil
 }
 
 // canonicalize gets the result of canonicalizing the packages using the results

internal/imports/mod_cache.go

@@ -129,17 +129,17 @@ func (d *dirInfoCache) Keys() (keys []string) {
 	return keys
 }
 
-func (d *dirInfoCache) CachePackageName(info directoryPackageInfo) (directoryPackageInfo, error) {
+func (d *dirInfoCache) CachePackageName(info directoryPackageInfo) (string, error) {
 	if loaded, err := info.reachedStatus(nameLoaded); loaded {
-		return info, err
+		return info.packageName, err
 	}
 	if scanned, err := info.reachedStatus(directoryScanned); !scanned || err != nil {
-		return info, fmt.Errorf("cannot read package name, scan error: %v", err)
+		return "", fmt.Errorf("cannot read package name, scan error: %v", err)
 	}
 	info.packageName, info.err = packageDirToName(info.dir)
 	info.status = nameLoaded
 	d.Store(info.dir, info)
-	return info, info.err
+	return info.packageName, info.err
 }
 
 func (d *dirInfoCache) CacheExports(ctx context.Context, env *ProcessEnv, info directoryPackageInfo) (string, []string, error) {
@@ -150,7 +150,7 @@ func (d *dirInfoCache) CacheExports(ctx context.Context, env *ProcessEnv, info d
 		return "", nil, err
 	}
 	info.packageName, info.exports, info.err = loadExportsFromFiles(ctx, env, info.dir)
-	if info.err == context.Canceled {
+	if info.err == context.Canceled || info.err == context.DeadlineExceeded {
 		return info.packageName, info.exports, info.err
 	}
 	// The cache structure wants things to proceed linearly. We can skip a

internal/imports/mod_test.go

@@ -4,6 +4,7 @@ package imports
 
 import (
 	"archive/zip"
+	"context"
 	"fmt"
 	"go/build"
 	"io/ioutil"
@@ -89,7 +90,7 @@ package z
 
 	mt.assertFound("y", "y")
 
-	scan, err := mt.resolver.scan(nil, false, nil)
+	scan, err := scanToSlice(mt.resolver, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -572,7 +573,7 @@ func (t *modTest) assertFound(importPath, pkgName string) (string, *pkg) {
 
 func (t *modTest) assertScanFinds(importPath, pkgName string) *pkg {
 	t.Helper()
-	scan, err := t.resolver.scan(nil, true, nil)
+	scan, err := scanToSlice(t.resolver, nil)
 	if err != nil {
 		t.Errorf("scan failed: %v", err)
 	}
@@ -585,6 +586,26 @@ func (t *modTest) assertScanFinds(importPath, pkgName string) *pkg {
 	return nil
 }
 
+func scanToSlice(resolver Resolver, exclude []gopathwalk.RootType) ([]*pkg, error) {
+	var mu sync.Mutex
+	var result []*pkg
+	filter := &scanCallback{
+		dirFound: func(pkg *pkg) bool {
+			return true
+		},
+		packageNameLoaded: func(pkg *pkg) bool {
+			mu.Lock()
+			defer mu.Unlock()
+			result = append(result, pkg)
+			return true
+		},
+		exportsLoaded: func(pkg *pkg, exports []string) {
+		},
+	}
+	err := resolver.scan(context.Background(), filter, exclude)
+	return result, err
+}
+
 // assertModuleFoundInDir is the same as assertFound, but also checks that the
 // package was found in an active module whose Dir matches dirRE.
 func (t *modTest) assertModuleFoundInDir(importPath, pkgName, dirRE string) {
@@ -829,7 +850,7 @@ func TestInvalidModCache(t *testing.T) {
 		WorkingDir:  dir,
 	}
 	resolver := &ModuleResolver{env: env}
-	resolver.scan(nil, true, nil)
+	scanToSlice(resolver, nil)
 }
 
 func TestGetCandidatesRanking(t *testing.T) {
@@ -864,7 +885,7 @@ import _ "rsc.io/quote"
 		// Out of scope modules
 		{"quote", "rsc.io/quote/v2"},
 	}
-	candidates, err := getAllCandidates("foo.go", mt.env)
+	candidates, err := getAllCandidates(context.Background(), "", "foo.go", mt.env)
 	if err != nil {
 		t.Fatalf("getAllCandidates() = %v", err)
 	}
@@ -889,10 +910,10 @@ func BenchmarkScanModCache(b *testing.B) {
 		Logf:   log.Printf,
 	}
 	exclude := []gopathwalk.RootType{gopathwalk.RootGOROOT}
-	env.GetResolver().scan(nil, true, exclude)
+	scanToSlice(env.GetResolver(), exclude)
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		env.GetResolver().scan(nil, true, exclude)
+		scanToSlice(env.GetResolver(), exclude)
 		env.GetResolver().(*ModuleResolver).ClearForNewScan()
 	}
 }

internal/lsp/cache/view.go

@@ -206,9 +206,9 @@ func (v *view) RunProcessEnvFunc(ctx context.Context, fn func(*imports.Options)
 
 			log.Print(context.Background(), "background imports cache refresh starting")
 			v.processEnv.GetResolver().ClearForNewScan()
-			_, err := imports.GetAllCandidates("", opts)
+			// TODO(heschi): prime the cache
 			v.cacheRefreshTime = time.Now()
-			log.Print(context.Background(), "background refresh finished with err: ", tag.Of("err", err))
+			log.Print(context.Background(), "background refresh finished with err: ", tag.Of("err", nil))
 		}()
 	}
 

internal/lsp/source/completion.go

@@ -828,8 +828,14 @@ func (c *completer) lexical() error {
 	}
 
 	if c.opts.Unimported && len(c.items) < unimportedTarget {
+		ctx, cancel := context.WithDeadline(c.ctx, c.startTime.Add(c.opts.Budget))
+		defer cancel()
 		// Suggest packages that have not been imported yet.
-		pkgs, err := CandidateImports(c.ctx, c.snapshot.View(), c.filename)
+		prefix := ""
+		if c.surrounding != nil {
+			prefix = c.surrounding.Prefix()
+		}
+		pkgs, err := CandidateImports(ctx, prefix, c.snapshot.View(), c.filename)
 		if err != nil {
 			return err
 		}

internal/lsp/source/format.go

@@ -313,7 +313,7 @@ func trimToFirstNonImport(fset *token.FileSet, f *ast.File, src []byte, err erro
 }
 
 // CandidateImports returns every import that could be added to filename.
-func CandidateImports(ctx context.Context, view View, filename string) ([]imports.ImportFix, error) {
+func CandidateImports(ctx context.Context, prefix string, view View, filename string) ([]imports.ImportFix, error) {
 	ctx, done := trace.StartSpan(ctx, "source.CandidateImports")
 	defer done()
 
@@ -330,7 +330,7 @@ func CandidateImports(ctx context.Context, view View, filename string) ([]import
 	var imps []imports.ImportFix
 	importFn := func(opts *imports.Options) error {
 		var err error
-		imps, err = imports.GetAllCandidates(filename, opts)
+		imps, err = imports.GetAllCandidates(ctx, prefix, filename, opts)
 		return err
 	}
 	err := view.RunProcessEnvFunc(ctx, importFn, options)
@@ -356,7 +356,7 @@ func PackageExports(ctx context.Context, view View, pkg, filename string) ([]imp
 	var pkgs []imports.PackageExport
 	importFn := func(opts *imports.Options) error {
 		var err error
-		pkgs, err = imports.GetPackageExports(pkg, filename, opts)
+		pkgs, err = imports.GetPackageExports(ctx, pkg, filename, opts)
 		return err
 	}
 	err := view.RunProcessEnvFunc(ctx, importFn, options)