internal/imports: skip scanned module cache dirs

Packages found in the module cache do not change. When we encounter a
directory we have already processed in the module cache, skip that
directory and add the packages that have already been computed.

Change-Id: Ib1bf0bf22727110b8073b415b145034acceb6787
Reviewed-on: https://go-review.googlesource.com/c/tools/+/186921
Run-TryBot: Suzy Mueller <suzmue@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>

internal/gopathwalk/walk.go

@@ -59,12 +59,24 @@ func SrcDirsRoots(ctx *build.Context) []Root {
 // paths of the containing source directory and the package directory.
 // add will be called concurrently.
 func Walk(roots []Root, add func(root Root, dir string), opts Options) {
+	WalkSkip(roots, add, func(Root, string) bool { return false }, opts)
+}
+
+// WalkSkip walks Go source directories ($GOROOT, $GOPATH, etc) to find packages.
+// For each package found, add will be called (concurrently) with the absolute
+// paths of the containing source directory and the package directory.
+// For each directory that will be scanned, skip will be called (concurrently)
+// with the absolute paths of the containing source directory and the directory.
+// If skip returns false on a directory it will be processed.
+// add will be called concurrently.
+// skip will be called concurrently.
+func WalkSkip(roots []Root, add func(root Root, dir string), skip func(root Root, dir string) bool, opts Options) {
 	for _, root := range roots {
-		walkDir(root, add, opts)
+		walkDir(root, add, skip, opts)
 	}
 }
 
-func walkDir(root Root, add func(Root, string), opts Options) {
+func walkDir(root Root, add func(Root, string), skip func(root Root, dir string) bool, opts Options) {
 	if _, err := os.Stat(root.Path); os.IsNotExist(err) {
 		if opts.Debug {
 			log.Printf("skipping nonexistent directory: %v", root.Path)
@@ -77,6 +89,7 @@ func walkDir(root Root, add func(Root, string), opts Options) {
 	w := &walker{
 		root: root,
 		add:  add,
+		skip: skip,
 		opts: opts,
 	}
 	w.init()
@@ -93,6 +106,7 @@ func walkDir(root Root, add func(Root, string), opts Options) {
 type walker struct {
 	root Root                    // The source directory to scan.
 	add  func(Root, string)      // The callback that will be invoked for every possible Go package dir.
+	skip func(Root, string) bool // The callback that will be invoked for every dir. dir is skipped if it returns true.
 	opts Options                 // Options passed to Walk by the user.
 
 	ignoredDirs []os.FileInfo // The ignored directories, loaded from .goimportsignore files.
@@ -151,12 +165,16 @@ func (w *walker) getIgnoredDirs(path string) []string {
 	return ignoredDirs
 }
 
-func (w *walker) shouldSkipDir(fi os.FileInfo) bool {
+func (w *walker) shouldSkipDir(fi os.FileInfo, dir string) bool {
 	for _, ignoredDir := range w.ignoredDirs {
 		if os.SameFile(fi, ignoredDir) {
 			return true
 		}
 	}
+	if w.skip != nil {
+		// Check with the user specified callback.
+		return w.skip(w.root, dir)
+	}
 	return false
 }
 
@@ -184,7 +202,7 @@ func (w *walker) walk(path string, typ os.FileMode) error {
 			return filepath.SkipDir
 		}
 		fi, err := os.Lstat(path)
-		if err == nil && w.shouldSkipDir(fi) {
+		if err == nil && w.shouldSkipDir(fi, path) {
 			return filepath.SkipDir
 		}
 		return nil
@@ -224,7 +242,7 @@ func (w *walker) shouldTraverse(dir string, fi os.FileInfo) bool {
 	if !ts.IsDir() {
 		return false
 	}
-	if w.shouldSkipDir(ts) {
+	if w.shouldSkipDir(ts, dir) {
 		return false
 	}
 	// Check for symlink loops by statting each directory component

internal/gopathwalk/walk_test.go

@@ -11,6 +11,7 @@ import (
 	"reflect"
 	"runtime"
 	"strings"
+	"sync"
 	"testing"
 )
 
@@ -107,14 +108,52 @@ func TestSkip(t *testing.T) {
 	}
 
 	var found []string
-	walkDir(Root{filepath.Join(dir, "src"), RootGOPATH}, func(root Root, dir string) {
+	var mu sync.Mutex
+	walkDir(Root{filepath.Join(dir, "src"), RootGOPATH},
+		func(root Root, dir string) {
+			mu.Lock()
+			defer mu.Unlock()
 			found = append(found, dir[len(root.Path)+1:])
+		}, func(root Root, dir string) bool {
+			return false
 		}, Options{ModulesEnabled: false, Debug: true})
 	if want := []string{"shouldfind"}; !reflect.DeepEqual(found, want) {
 		t.Errorf("expected to find only %v, got %v", want, found)
 	}
 }
 
+// TestSkipFunction tests that scan successfully skips directories from user callback.
+func TestSkipFunction(t *testing.T) {
+	dir, err := ioutil.TempDir("", "goimports-")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(dir)
+
+	if err := mapToDir(dir, map[string]string{
+		"ignoreme/f.go":           "package ignoreme",    // ignored by skip
+		"ignoreme/subignore/f.go": "package subignore",   // also ignored by skip
+		"shouldfind/f.go":         "package shouldfind;", // not ignored
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	var found []string
+	var mu sync.Mutex
+	walkDir(Root{filepath.Join(dir, "src"), RootGOPATH},
+		func(root Root, dir string) {
+			mu.Lock()
+			defer mu.Unlock()
+			found = append(found, dir[len(root.Path)+1:])
+		}, func(root Root, dir string) bool {
+			return strings.HasSuffix(dir, "ignoreme")
+		},
+		Options{ModulesEnabled: false})
+	if want := []string{"shouldfind"}; !reflect.DeepEqual(found, want) {
+		t.Errorf("expected to find only %v, got %v", want, found)
+	}
+}
+
 func mapToDir(destDir string, files map[string]string) error {
 	for path, contents := range files {
 		file := filepath.Join(destDir, "src", path)

internal/imports/mod.go

@@ -31,10 +31,6 @@ type ModuleResolver struct {
 	ModsByModPath []*ModuleJSON // All modules, ordered by # of path components in module Path...
 	ModsByDir     []*ModuleJSON // ...or Dir.
 
-	// ModCachePkgs contains the canonicalized importPath and directory of packages
-	// in the module cache. Keyed by absolute directory.
-	ModCachePkgs map[string]*pkg
-
 	// moduleCacheInfo stores information about the module cache.
 	moduleCacheInfo *moduleCacheInfo
 }
@@ -97,7 +93,6 @@ func (r *ModuleResolver) init() error {
 		return count(j) < count(i) // descending order
 	})
 
-	r.ModCachePkgs = make(map[string]*pkg)
 	if r.moduleCacheInfo == nil {
 		r.moduleCacheInfo = &moduleCacheInfo{
 			modCacheDirInfo: make(map[string]*directoryPackageInfo),
@@ -268,46 +263,89 @@ func (r *ModuleResolver) scan(_ references) ([]*pkg, error) {
 	dupCheck := make(map[string]bool)
 	var mu sync.Mutex
 
-	gopathwalk.Walk(roots, func(root gopathwalk.Root, dir string) {
+	// Packages in the module cache are immutable. If we have
+	// already seen this package on a previous scan of the module
+	// cache, return that result.
+	skip := func(root gopathwalk.Root, dir string) bool {
 		mu.Lock()
 		defer mu.Unlock()
+		// If we have already processed this directory on this walk, skip it.
+		if _, dup := dupCheck[dir]; dup {
+			return true
+		}
 
+		// If we have saved this directory information, skip it.
+		info, ok := r.moduleCacheInfo.Load(dir)
+		if !ok {
+			return false
+		}
+		// This directory can be skipped as long as we have already scanned it.
+		// Packages with errors will continue to have errors, so there is no need
+		// to rescan them.
+		packageScanned, _ := info.reachedStatus(directoryScanned)
+		return packageScanned
+	}
+
+	add := func(root gopathwalk.Root, dir string) {
+		mu.Lock()
+		defer mu.Unlock()
 		if _, dup := dupCheck[dir]; dup {
 			return
 		}
 
-		dupCheck[dir] = true
-
-		absDir := dir
-		// Packages in the module cache are immutable. If we have
-		// already seen this package on a previous scan of the module
-		// cache, return that result.
-		if p, ok := r.ModCachePkgs[absDir]; ok {
-			result = append(result, p)
-			return
-		}
-
-		info, ok := r.moduleCacheInfo.Load(dir)
-		if !ok {
-			var err error
-			info, err = r.scanDirForPackage(root, dir)
+		info, err := r.scanDirForPackage(root, dir)
 		if err != nil {
 			return
 		}
 		if root.Type == gopathwalk.RootModuleCache {
+			// Save this package information in the cache and return.
+			// Packages from the module cache are added after Walk.
 			r.moduleCacheInfo.Store(dir, info)
-			}
+			return
 		}
 
-		if info.status < directoryScanned ||
-			(info.status == directoryScanned && info.err != nil) {
+		// Skip this package if there was an error loading package info.
+		if info.err != nil {
 			return
 		}
 
 		// The rest of this function canonicalizes the packages using the results
 		// of initializing the resolver from 'go list -m'.
-		importPath := info.nonCanonicalImportPath
+		res, err := r.canonicalize(info.nonCanonicalImportPath, info.dir, info.needsReplace)
+		if err != nil {
+			return
+		}
 
+		result = append(result, res)
+	}
+
+	gopathwalk.WalkSkip(roots, add, skip, gopathwalk.Options{Debug: r.env.Debug, ModulesEnabled: true})
+
+	// Add the packages from the modules in the mod cache that were skipped.
+	for _, dir := range r.moduleCacheInfo.Keys() {
+		info, ok := r.moduleCacheInfo.Load(dir)
+		if !ok {
+			continue
+		}
+
+		// Skip this directory if we were not able to get the package information successfully.
+		if scanned, err := info.reachedStatus(directoryScanned); !scanned || err != nil {
+			continue
+		}
+
+		res, err := r.canonicalize(info.nonCanonicalImportPath, info.dir, info.needsReplace)
+		if err != nil {
+			continue
+		}
+		result = append(result, res)
+	}
+
+	return result, nil
+}
+
+// canonicalize gets the result of canonicalizing the packages using the results
+// of initializing the resolver from 'go list -m'.
+func (r *ModuleResolver) canonicalize(importPath, dir string, needsReplace bool) (res *pkg, err error) {
 	// Check if the directory is underneath a module that's in scope.
 	if mod := r.findModuleByDir(dir); mod != nil {
 		// It is. If dir is the target of a replace directive,
@@ -318,9 +356,8 @@ func (r *ModuleResolver) scan(_ references) ([]*pkg, error) {
 			dirInMod := dir[len(mod.Dir)+len("/"):]
 			importPath = path.Join(mod.Path, filepath.ToSlash(dirInMod))
 		}
-		} else if info.needsReplace {
-			// This package needed a replace target we don't have.
-			return
+	} else if needsReplace {
+		return nil, fmt.Errorf("needed this package to be in scope: %s", dir)
 	}
 
 	// We may have discovered a package that has a different version
@@ -328,23 +365,10 @@ func (r *ModuleResolver) scan(_ references) ([]*pkg, error) {
 	if _, canonicalDir := r.findPackage(importPath); canonicalDir != "" {
 		dir = canonicalDir
 	}
-
-		res := &pkg{
+	return &pkg{
 		importPathShort: VendorlessPath(importPath),
 		dir:             dir,
-		}
-
-		if root.Type == gopathwalk.RootModuleCache {
-			// Save the results of processing this directory.
-			// This needs to be invalidated when the results of
-			// 'go list -m' would change, as the directory and
-			// importPath in this map depend on those results.
-			r.ModCachePkgs[absDir] = res
-		}
-
-		result = append(result, res)
-	}, gopathwalk.Options{Debug: r.env.Debug, ModulesEnabled: true})
-	return result, nil
+	}, nil
 }
 
 func (r *ModuleResolver) loadExports(ctx context.Context, expectPackage string, pkg *pkg) (map[string]bool, error) {

internal/imports/mod_cache.go

@@ -109,3 +109,13 @@ func (d *moduleCacheInfo) Load(dir string) (directoryPackageInfo, bool) {
 	}
 	return *info, true
 }
+
+// Keys returns the keys currently present in d.
+func (d *moduleCacheInfo) Keys() (keys []string) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	for key := range d.modCacheDirInfo {
+		keys = append(keys, key)
+	}
+	return keys
+}

internal/imports/mod_cache_test.go

@@ -2,6 +2,7 @@ package imports
 
 import (
 	"fmt"
+	"sort"
 	"testing"
 )
 
@@ -50,3 +51,73 @@ func TestDirectoryPackageInfoReachedStatus(t *testing.T) {
 		}
 	}
 }
+
+func TestModCacheInfo(t *testing.T) {
+	m := &moduleCacheInfo{
+		modCacheDirInfo: make(map[string]*directoryPackageInfo),
+	}
+
+	dirInfo := []struct {
+		dir  string
+		info directoryPackageInfo
+	}{
+		{
+			dir: "mypackage",
+			info: directoryPackageInfo{
+				status:                 directoryScanned,
+				dir:                    "mypackage",
+				nonCanonicalImportPath: "example.com/mypackage",
+				needsReplace:           false,
+			},
+		},
+		{
+			dir: "bad package",
+			info: directoryPackageInfo{
+				status: directoryScanned,
+				err:    fmt.Errorf("bad package"),
+			},
+		},
+		{
+			dir: "mypackage/other",
+			info: directoryPackageInfo{
+				dir:                    "mypackage/other",
+				nonCanonicalImportPath: "example.com/mypackage/other",
+				needsReplace:           false,
+			},
+		},
+	}
+
+	for _, d := range dirInfo {
+		m.Store(d.dir, d.info)
+	}
+
+	for _, d := range dirInfo {
+		val, ok := m.Load(d.dir)
+		if !ok {
+			t.Errorf("directory not loaded: %s", d.dir)
+		}
+
+		if val != d.info {
+			t.Errorf("expected: %v, got: %v", d.info, val)
+		}
+	}
+
+	var wantKeys []string
+	for _, d := range dirInfo {
+		wantKeys = append(wantKeys, d.dir)
+	}
+	sort.Strings(wantKeys)
+
+	gotKeys := m.Keys()
+	sort.Strings(gotKeys)
+
+	if len(gotKeys) != len(wantKeys) {
+		t.Errorf("different length of keys. expected: %d, got: %d", len(wantKeys), len(gotKeys))
+	}
+
+	for i, want := range wantKeys {
+		if want != gotKeys[i] {
+			t.Errorf("%d: expected %s, got %s", i, want, gotKeys[i])
+		}
+	}
+}

internal/imports/mod_test.go

@@ -137,6 +137,25 @@ import _ "example.com"
 
 }
 
+// Tests that scanning the module cache > 1 time is able to find the same module
+// in the module cache.
+func TestModMultipleScansWithSubdirs(t *testing.T) {
+	mt := setup(t, `
+-- go.mod --
+module x
+
+require rsc.io/quote v1.5.2
+
+-- x.go --
+package x
+import _ "rsc.io/quote"
+`, "")
+	defer mt.cleanup()
+
+	mt.assertScanFinds("rsc.io/quote/buggy", "buggy")
+	mt.assertScanFinds("rsc.io/quote/buggy", "buggy")
+}
+
 // Tests that scanning the module cache > 1 after changing a package in module cache to make it unimportable
 // is able to find the same module.
 func TestModCacheEditModFile(t *testing.T) {
@@ -151,6 +170,9 @@ import _ "rsc.io/quote"
 `, "")
 	defer mt.cleanup()
 	found := mt.assertScanFinds("rsc.io/quote", "quote")
+	if found == nil {
+		t.Fatal("rsc.io/quote not found in initial scan.")
+	}
 
 	// Update the go.mod file of example.com so that it changes its module path (not allowed).
 	if err := os.Chmod(filepath.Join(found.dir, "go.mod"), 0644); err != nil {
@@ -176,7 +198,6 @@ import _ "rsc.io/quote"
 	mt.resolver.Main = nil
 	mt.resolver.ModsByModPath = nil
 	mt.resolver.ModsByDir = nil
-	mt.resolver.ModCachePkgs = nil
 	mt.assertScanFinds("rsc.io/quote", "quote")
 
 }