internal/imports: optimize scan implementations

In scan implementations, stop after cancellation, and swallow the
context's error for convenience.

In the module implementation specifically, try to avoid scanning if the
cache is enough to satisfy the user. When we do have to scan, prioritize
module dependencies before the whole cache.

Change-Id: I23dc98df016f9fca4f31c7ded3d11bc257c29b94
Reviewed-on: https://go-review.googlesource.com/c/tools/+/212857
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

@@ -663,11 +663,7 @@ func getAllCandidates(ctx context.Context, wrapped func(ImportFix), prefix strin
 			return false
 		},
 	}
-	err := getCandidatePkgs(ctx, callback, filename, env)
+	return getCandidatePkgs(ctx, callback, filename, env)
-	if err != nil {
-		return err
-	}
-	return nil
 }
 
 // A PackageExport is a package and its exports.
@@ -701,11 +697,7 @@ func getPackageExports(ctx context.Context, wrapped func(PackageExport), complet
 			})
 		},
 	}
-	err := getCandidatePkgs(ctx, callback, filename, env)
+	return getCandidatePkgs(ctx, callback, filename, env)
-	if err != nil {
-		return err
-	}
-	return nil
 }
 
 // ProcessEnv contains environment variables and settings that affect the use of
@@ -1146,6 +1138,10 @@ func (r *gopathResolver) scan(ctx context.Context, callback *scanCallback, exclu
 	roots := filterRoots(gopathwalk.SrcDirsRoots(r.env.buildContext()), exclude)
 	gopathwalk.Walk(roots, add, gopathwalk.Options{Debug: r.env.Debug, ModulesEnabled: false})
 	for _, dir := range r.cache.Keys() {
+		if ctx.Err() != nil {
+			return nil
+		}
+
 		info, ok := r.cache.Load(dir)
 		if !ok {
 			continue

internal/imports/mod.go

@@ -22,9 +22,11 @@ import (
 // ModuleResolver implements resolver for modules using the go command as little
 // as feasible.
 type ModuleResolver struct {
-	env            *ProcessEnv
+	env             *ProcessEnv
-	moduleCacheDir string
+	moduleCacheDir  string
-	dummyVendorMod *ModuleJSON // If vendoring is enabled, the pseudo-module that represents the /vendor directory.
+	dummyVendorMod  *ModuleJSON // If vendoring is enabled, the pseudo-module that represents the /vendor directory.
+	roots           []gopathwalk.Root
+	walkedRootIndex int
 
 	Initialized   bool
 	Main          *ModuleJSON
@@ -85,6 +87,37 @@ func (r *ModuleResolver) init() error {
 		return count(j) < count(i) // descending order
 	})
 
+	r.roots = []gopathwalk.Root{
+		{filepath.Join(r.env.GOROOT, "/src"), gopathwalk.RootGOROOT},
+	}
+	if r.Main != nil {
+		r.roots = append(r.roots, gopathwalk.Root{r.Main.Dir, gopathwalk.RootCurrentModule})
+	}
+	if vendorEnabled {
+		r.roots = append(r.roots, gopathwalk.Root{r.dummyVendorMod.Dir, gopathwalk.RootOther})
+	} else {
+		addDep := func(mod *ModuleJSON) {
+			if mod.Replace == nil {
+				// This is redundant with the cache, but we'll skip it cheaply enough.
+				r.roots = append(r.roots, gopathwalk.Root{mod.Dir, gopathwalk.RootModuleCache})
+			} else {
+				r.roots = append(r.roots, gopathwalk.Root{mod.Dir, gopathwalk.RootOther})
+			}
+		}
+		// Walk dependent modules before scanning the full mod cache, direct deps first.
+		for _, mod := range r.ModsByModPath {
+			if !mod.Indirect {
+				addDep(mod)
+			}
+		}
+		for _, mod := range r.ModsByModPath {
+			if mod.Indirect {
+				addDep(mod)
+			}
+		}
+		r.roots = append(r.roots, gopathwalk.Root{r.moduleCacheDir, gopathwalk.RootModuleCache})
+	}
+
 	if r.moduleCacheCache == nil {
 		r.moduleCacheCache = &dirInfoCache{
 			dirs: map[string]*directoryPackageInfo{},
@@ -126,6 +159,7 @@ func (r *ModuleResolver) initAllMods() error {
 }
 
 func (r *ModuleResolver) ClearForNewScan() {
+	r.walkedRootIndex = 0
 	r.otherCache = &dirInfoCache{
 		dirs: map[string]*directoryPackageInfo{},
 	}
@@ -338,29 +372,49 @@ func (r *ModuleResolver) scan(ctx context.Context, callback *scanCallback, exclu
 		return err
 	}
 
-	// Walk GOROOT, GOPATH/pkg/mod, and the main module.
+	processDir := func(info directoryPackageInfo) {
-	roots := []gopathwalk.Root{
+		// Skip this directory if we were not able to get the package information successfully.
-		{filepath.Join(r.env.GOROOT, "/src"), gopathwalk.RootGOROOT},
+		if scanned, err := info.reachedStatus(directoryScanned); !scanned || err != nil {
-	}
+			return
-	if r.Main != nil {
+		}
-		roots = append(roots, gopathwalk.Root{r.Main.Dir, gopathwalk.RootCurrentModule})
+
-	}
+		pkg, err := r.canonicalize(info)
-	if r.dummyVendorMod != nil {
+		if err != nil {
-		roots = append(roots, gopathwalk.Root{r.dummyVendorMod.Dir, gopathwalk.RootOther})
+			return
-	} else {
+		}
-		roots = append(roots, gopathwalk.Root{r.moduleCacheDir, gopathwalk.RootModuleCache})
+
-		// Walk replace targets, just in case they're not in any of the above.
+		if callback.dirFound(pkg) {
-		for _, mod := range r.ModsByModPath {
+			var err error
-			if mod.Replace != nil {
+			pkg.packageName, err = r.cachePackageName(info)
-				roots = append(roots, gopathwalk.Root{mod.Dir, gopathwalk.RootOther})
+			if err != nil {
+				return
 			}
 		}
+
+		if callback.packageNameLoaded(pkg) {
+			_, exports, err := r.loadExports(ctx, pkg)
+			if err != nil {
+				return
+			}
+			callback.exportsLoaded(pkg, exports)
+		}
 	}
 
-	roots = filterRoots(roots, exclude)
+	// Everything we already had is in the cache. Process it now, in hopes we
+	// we don't need anything new.
+	for _, dir := range r.cacheKeys() {
+		if ctx.Err() != nil {
+			return nil
+		}
+		info, ok := r.cacheLoad(dir)
+		if !ok {
+			continue
+		}
+		processDir(info)
+	}
 
-	// We assume cached directories have not changed. We can skip them and their
+	// We assume cached directories are fully cached, including all their
-	// children.
+	// children, and have not changed. We can skip them.
 	skip := func(root gopathwalk.Root, dir string) bool {
 		info, ok := r.cacheLoad(dir)
 		if !ok {
@@ -373,45 +427,31 @@ func (r *ModuleResolver) scan(ctx context.Context, callback *scanCallback, exclu
 		return packageScanned
 	}
 
-	// Add anything new to the cache. We'll process everything in it below.
+	// Add anything new to the cache, and process it if we're still looking.
 	add := func(root gopathwalk.Root, dir string) {
-		r.cacheStore(r.scanDirForPackage(root, dir))
+		info := r.scanDirForPackage(root, dir)
+		r.cacheStore(info)
+		if ctx.Err() == nil {
+			processDir(info)
+		}
 	}
 
-	gopathwalk.WalkSkip(roots, add, skip, gopathwalk.Options{Debug: r.env.Debug, ModulesEnabled: true})
+	// We can't cancel walks, because we need them to finish to have a usable
-
+	// cache. We can do them one by one and stop in between.
-	// Everything we already had, and everything new, is now in the cache.
+	// TODO(heschi): Run asynchronously and detach on cancellation? Would risk
-	for _, dir := range r.cacheKeys() {
+	// racy callbacks.
-		info, ok := r.cacheLoad(dir)
+rootLoop:
-		if !ok {
+	for ; r.walkedRootIndex < len(r.roots); r.walkedRootIndex++ {
-			continue
+		root := r.roots[r.walkedRootIndex]
-		}
+		for _, rt := range exclude {
-
+			if root.Type == rt {
-		// Skip this directory if we were not able to get the package information successfully.
+				continue rootLoop
-		if scanned, err := info.reachedStatus(directoryScanned); !scanned || err != nil {
-			continue
-		}
-
-		pkg, err := r.canonicalize(info)
-		if err != nil {
-			continue
-		}
-
-		if callback.dirFound(pkg) {
-			var err error
-			pkg.packageName, err = r.cachePackageName(info)
-			if err != nil {
-				continue
 			}
 		}
-
+		if ctx.Err() != nil {
-		if callback.packageNameLoaded(pkg) {
+			return nil
-			_, exports, err := r.loadExports(ctx, pkg)
-			if err != nil {
-				continue
-			}
-			callback.exportsLoaded(pkg, exports)
 		}
+		gopathwalk.WalkSkip([]gopathwalk.Root{root}, add, skip, gopathwalk.Options{Debug: r.env.Debug, ModulesEnabled: true})
 	}
 	return nil
 }

internal/lsp/source/completion.go

@@ -663,7 +663,7 @@ func (c *completer) selector(sel *ast.SelectorExpr) error {
 		}
 		if err := c.snapshot.View().RunProcessEnvFunc(ctx, func(opts *imports.Options) error {
 			return imports.GetPackageExports(ctx, add, id.Name, c.filename, opts)
-		}); err != nil && err != context.Canceled {
+		}); err != nil {
 			return err
 		}
 	}
@@ -886,7 +886,7 @@ func (c *completer) lexical() error {
 		}
 		if err := c.snapshot.View().RunProcessEnvFunc(ctx, func(opts *imports.Options) error {
 			return imports.GetAllCandidates(ctx, add, prefix, c.filename, opts)
-		}); err != nil && err != context.Canceled {
+		}); err != nil {
 			return err
 		}
 	}