cmd/goimports, imports: optimize directory scanning and other things

This brings goimports from 160ms to 100ms on my laptop, and under 50ms
on my Linux machine.

Using cmd/trace, I noticed that filepath.Walk is inherently slow.
See https://golang.org/issue/16399 for details.

Instead, this CL introduces a new (private) filepath.Walk
implementation, optimized for speed and avoiding unnecessary work.

In addition to avoid an Lstat per file, it also reads directories
concurrently. The old goimports code did that too, but now that logic
is removed from goimports and the code is simplified.

This also adds some profiling command line flags to goimports that I
found useful.

Updates golang/go#16367 (goimports is slow)
Updates golang/go#16399 (filepath.Walk is slow)

Change-Id: I708d570cbaad3fa9ad75a12054f5a932ee159b84
Reviewed-on: https://go-review.googlesource.com/25001
Reviewed-by: Andrew Gerrand <adg@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

cmd/goimports/goimports.go

@@ -5,6 +5,7 @@
 package main
 
 import (
+	"bufio"
 	"bytes"
 	"flag"
 	"fmt"
@@ -16,6 +17,8 @@ import (
 	"os/exec"
 	"path/filepath"
 	"runtime"
+	"runtime/pprof"
+	"runtime/trace"
 	"strings"
 
 	"golang.org/x/tools/imports"
@@ -29,6 +32,11 @@ var (
 	srcdir  = flag.String("srcdir", "", "choose imports as if source code is from `dir`")
 	verbose = flag.Bool("v", false, "verbose logging")
 
+	cpuProfile     = flag.String("cpuprofile", "", "CPU profile output")
+	memProfile     = flag.String("memprofile", "", "memory profile output")
+	memProfileRate = flag.Int("memrate", 0, "if > 0, sets runtime.MemProfileRate")
+	traceProfile   = flag.String("trace", "", "trace profile output")
+
 	options = &imports.Options{
 		TabWidth:  8,
 		TabIndent: true,
@@ -152,10 +160,50 @@ var parseFlags = func() []string {
 	return flag.Args()
 }
 
+func bufferedFileWriter(dest string) (w io.Writer, close func()) {
+	f, err := os.Create(dest)
+	if err != nil {
+		log.Fatal(err)
+	}
+	bw := bufio.NewWriter(f)
+	return bw, func() {
+		if err := bw.Flush(); err != nil {
+			log.Fatalf("error flushing %v: %v", dest, err)
+		}
+		if err := f.Close(); err != nil {
+			log.Fatal(err)
+		}
+	}
+}
+
 func gofmtMain() {
 	flag.Usage = usage
 	paths := parseFlags()
 
+	if *cpuProfile != "" {
+		bw, flush := bufferedFileWriter(*cpuProfile)
+		pprof.StartCPUProfile(bw)
+		defer flush()
+		defer pprof.StopCPUProfile()
+	}
+	if *traceProfile != "" {
+		bw, flush := bufferedFileWriter(*traceProfile)
+		trace.Start(bw)
+		defer flush()
+		defer trace.Stop()
+	}
+	if *memProfileRate > 0 {
+		runtime.MemProfileRate = *memProfileRate
+		bw, flush := bufferedFileWriter(*memProfile)
+		defer func() {
+			runtime.GC() // materialize all statistics
+			if err := pprof.WriteHeapProfile(bw); err != nil {
+				log.Fatal(err)
+			}
+			flush()
+		}()
+	}
+
 	if *verbose {
 		log.SetFlags(log.LstdFlags | log.Lmicroseconds)
 		imports.Debug = true

imports/fastwalk.go

@@ -0,0 +1,172 @@
+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// A faster implementation of filepath.Walk.
+//
+// filepath.Walk's design necessarily calls os.Lstat on each file,
+// even if the caller needs less info. And goimports only need to know
+// the type of each file. The kernel interface provides the type in
+// the Readdir call but the standard library ignored it.
+// fastwalk_unix.go contains a fork of the syscall routines.
+//
+// See golang.org/issue/16399
+
+package imports
+
+import (
+	"errors"
+	"os"
+	"path/filepath"
+	"runtime"
+)
+
+// traverseLink is a sentinel error for fastWalk, similar to filepath.SkipDir.
+var traverseLink = errors.New("traverse symlink, assuming target is a directory")
+
+// fastWalk walks the file tree rooted at root, calling walkFn for
+// each file or directory in the tree, including root.
+//
+// If fastWalk returns filepath.SkipDir, the directory is skipped.
+//
+// Unlike filepath.Walk:
+//   * file stat calls must be done by the user.
+//     The only provided metadata is the file type, which does not include
+//     any permission bits.
+//   * multiple goroutines stat the filesystem concurrently. The provided
+//     walkFn must be safe for concurrent use.
+//   * fastWalk can follow symlinks if walkFn returns the traverseLink
+//     sentinel error. It is the walkFn's responsibility to prevent
+//     fastWalk from going into symlink cycles.
+func fastWalk(root string, walkFn func(path string, typ os.FileMode) error) error {
+	// TODO(bradfitz): make numWorkers configurable? We used a
+	// minimum of 4 to give the kernel more info about multiple
+	// things we want, in hopes its I/O scheduling can take
+	// advantage of that. Hopefully most are in cache. Maybe 4 is
+	// even too low of a minimum. Profile more.
+	numWorkers := 4
+	if n := runtime.NumCPU(); n > numWorkers {
+		numWorkers = n
+	}
+	w := &walker{
+		fn:       walkFn,
+		enqueuec: make(chan walkItem, numWorkers), // buffered for performance
+		workc:    make(chan walkItem, numWorkers), // buffered for performance
+		donec:    make(chan struct{}),
+
+		// buffered for correctness & not leaking goroutines:
+		resc: make(chan error, numWorkers),
+	}
+	defer close(w.donec)
+	// TODO(bradfitz): start the workers as needed? maybe not worth it.
+	for i := 0; i < numWorkers; i++ {
+		go w.doWork()
+	}
+	todo := []walkItem{{dir: root}}
+	out := 0
+	for {
+		workc := w.workc
+		var workItem walkItem
+		if len(todo) == 0 {
+			workc = nil
+		} else {
+			workItem = todo[len(todo)-1]
+		}
+		select {
+		case workc <- workItem:
+			todo = todo[:len(todo)-1]
+			out++
+		case it := <-w.enqueuec:
+			todo = append(todo, it)
+		case err := <-w.resc:
+			out--
+			if err != nil {
+				return err
+			}
+			if out == 0 && len(todo) == 0 {
+				// It's safe to quit here, as long as the buffered
+				// enqueue channel isn't also readable, which might
+				// happen if the worker sends both another unit of
+				// work and its result before the other select was
+				// scheduled and both w.resc and w.enqueuec were
+				// readable.
+				select {
+				case it := <-w.enqueuec:
+					todo = append(todo, it)
+				default:
+					return nil
+				}
+			}
+		}
+	}
+}
+
+// doWork reads directories as instructed (via workc) and runs the
+// user's callback function.
+func (w *walker) doWork() {
+	for {
+		select {
+		case <-w.donec:
+			return
+		case it := <-w.workc:
+			w.resc <- w.walk(it.dir, !it.callbackDone)
+		}
+	}
+}
+
+type walker struct {
+	fn func(path string, typ os.FileMode) error
+
+	donec    chan struct{} // closed on fastWalk's return
+	workc    chan walkItem // to workers
+	enqueuec chan walkItem // from workers
+	resc     chan error    // from workers
+}
+
+type walkItem struct {
+	dir          string
+	callbackDone bool // callback already called; don't do it again
+}
+
+func (w *walker) enqueue(it walkItem) {
+	select {
+	case w.enqueuec <- it:
+	case <-w.donec:
+	}
+}
+
+func (w *walker) onDirEnt(dirName, baseName string, typ os.FileMode) error {
+	joined := dirName + string(os.PathSeparator) + baseName
+	if typ == os.ModeDir {
+		w.enqueue(walkItem{dir: joined})
+		return nil
+	}
+
+	err := w.fn(joined, typ)
+	if typ == os.ModeSymlink {
+		if err == traverseLink {
+			// Set callbackDone so we don't call it twice for both the
+			// symlink-as-symlink and the symlink-as-directory later:
+			w.enqueue(walkItem{dir: joined, callbackDone: true})
+			return nil
+		}
+		if err == filepath.SkipDir {
+			// Permit SkipDir on symlinks too.
+			return nil
+		}
+	}
+	return err
+}
+func (w *walker) walk(root string, runUserCallback bool) error {
+	if runUserCallback {
+		err := w.fn(root, os.ModeDir)
+		if err == filepath.SkipDir {
+			return nil
+		}
+		if err != nil {
+			return err
+		}
+	}
+
+	return readDir(root, w.onDirEnt)
+}

imports/fastwalk_dirent_fileno.go

@@ -0,0 +1,13 @@
+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// +build freebsd openbsd netbsd
+
+package imports
+
+import "syscall"
+
+func direntInode(dirent *syscall.Dirent) uint64 {
+	return uint64(dirent.Fileno)
+}

imports/fastwalk_dirent_ino.go

@@ -0,0 +1,13 @@
+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// +build linux darwin
+
+package imports
+
+import "syscall"
+
+func direntInode(dirent *syscall.Dirent) uint64 {
+	return uint64(dirent.Ino)
+}

imports/fastwalk_portable.go

@@ -0,0 +1,29 @@
+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// +build !linux,!darwin,!freebsd,!openbsd,!netbsd
+
+package imports
+
+import (
+	"io/ioutil"
+	"os"
+)
+
+// readDir calls fn for each directory entry in dirName.
+// It does not descend into directories or follow symlinks.
+// If fn returns a non-nil error, readDir returns with that error
+// immediately.
+func readDir(dirName string, fn func(dirName, entName string, typ os.FileMode) error) error {
+	fis, err := ioutil.ReadDir(dirName)
+	if err != nil {
+		return err
+	}
+	for _, fi := range fis {
+		if err := fn(dirName, fi.Name(), fi.Mode()&os.ModeType); err != nil {
+			return err
+		}
+	}
+	return nil
+}

imports/fastwalk_test.go

@@ -0,0 +1,171 @@
+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package imports
+
+import (
+	"bytes"
+	"flag"
+	"fmt"
+	"os"
+	"path/filepath"
+	"reflect"
+	"runtime"
+	"sort"
+	"strings"
+	"sync"
+	"testing"
+)
+
+func formatFileModes(m map[string]os.FileMode) string {
+	var keys []string
+	for k := range m {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+	var buf bytes.Buffer
+	for _, k := range keys {
+		fmt.Fprintf(&buf, "%-20s: %v\n", k, m[k])
+	}
+	return buf.String()
+}
+
+func testFastWalk(t *testing.T, files map[string]string, callback func(path string, typ os.FileMode) error, want map[string]os.FileMode) {
+	testConfig{
+		gopathFiles: files,
+	}.test(t, func(t *goimportTest) {
+		got := map[string]os.FileMode{}
+		var mu sync.Mutex
+		if err := fastWalk(t.gopath, func(path string, typ os.FileMode) error {
+			mu.Lock()
+			defer mu.Unlock()
+			if !strings.HasPrefix(path, t.gopath) {
+				t.Fatal("bogus prefix on %q, expect %q", path, t.gopath)
+			}
+			key := filepath.ToSlash(strings.TrimPrefix(path, t.gopath))
+			if old, dup := got[key]; dup {
+				t.Fatalf("callback called twice for key %q: %v -> %v", key, old, typ)
+			}
+			got[key] = typ
+			return callback(path, typ)
+		}); err != nil {
+			t.Fatalf("callback returned: %v", err)
+		}
+		if !reflect.DeepEqual(got, want) {
+			t.Errorf("walk mismatch.\n got:\n%v\nwant:\n%v", formatFileModes(got), formatFileModes(want))
+		}
+	})
+}
+
+func TestFastWalk_Basic(t *testing.T) {
+	testFastWalk(t, map[string]string{
+		"foo/foo.go":   "one",
+		"bar/bar.go":   "two",
+		"skip/skip.go": "skip",
+	},
+		func(path string, typ os.FileMode) error {
+			return nil
+		},
+		map[string]os.FileMode{
+			"":                  os.ModeDir,
+			"/src":              os.ModeDir,
+			"/src/bar":          os.ModeDir,
+			"/src/bar/bar.go":   0,
+			"/src/foo":          os.ModeDir,
+			"/src/foo/foo.go":   0,
+			"/src/skip":         os.ModeDir,
+			"/src/skip/skip.go": 0,
+		})
+}
+
+func TestFastWalk_Symlink(t *testing.T) {
+	switch runtime.GOOS {
+	case "windows", "plan9":
+		t.Skipf("skipping on %s", runtime.GOOS)
+	}
+	testFastWalk(t, map[string]string{
+		"foo/foo.go": "one",
+		"bar/bar.go": "LINK:../foo.go",
+		"symdir":     "LINK:foo",
+	},
+		func(path string, typ os.FileMode) error {
+			return nil
+		},
+		map[string]os.FileMode{
+			"":                os.ModeDir,
+			"/src":            os.ModeDir,
+			"/src/bar":        os.ModeDir,
+			"/src/bar/bar.go": os.ModeSymlink,
+			"/src/foo":        os.ModeDir,
+			"/src/foo/foo.go": 0,
+			"/src/symdir":     os.ModeSymlink,
+		})
+}
+
+func TestFastWalk_SkipDir(t *testing.T) {
+	testFastWalk(t, map[string]string{
+		"foo/foo.go":   "one",
+		"bar/bar.go":   "two",
+		"skip/skip.go": "skip",
+	},
+		func(path string, typ os.FileMode) error {
+			if typ == os.ModeDir && strings.HasSuffix(path, "skip") {
+				return filepath.SkipDir
+			}
+			return nil
+		},
+		map[string]os.FileMode{
+			"":                os.ModeDir,
+			"/src":            os.ModeDir,
+			"/src/bar":        os.ModeDir,
+			"/src/bar/bar.go": 0,
+			"/src/foo":        os.ModeDir,
+			"/src/foo/foo.go": 0,
+			"/src/skip":       os.ModeDir,
+		})
+}
+
+func TestFastWalk_TraverseSymlink(t *testing.T) {
+	switch runtime.GOOS {
+	case "windows", "plan9":
+		t.Skipf("skipping on %s", runtime.GOOS)
+	}
+
+	testFastWalk(t, map[string]string{
+		"foo/foo.go":   "one",
+		"bar/bar.go":   "two",
+		"skip/skip.go": "skip",
+		"symdir":       "LINK:foo",
+	},
+		func(path string, typ os.FileMode) error {
+			if typ == os.ModeSymlink {
+				return traverseLink
+			}
+			return nil
+		},
+		map[string]os.FileMode{
+			"":                   os.ModeDir,
+			"/src":               os.ModeDir,
+			"/src/bar":           os.ModeDir,
+			"/src/bar/bar.go":    0,
+			"/src/foo":           os.ModeDir,
+			"/src/foo/foo.go":    0,
+			"/src/skip":          os.ModeDir,
+			"/src/skip/skip.go":  0,
+			"/src/symdir":        os.ModeSymlink,
+			"/src/symdir/foo.go": 0,
+		})
+}
+
+var benchDir = flag.String("benchdir", runtime.GOROOT(), "The directory to scan for BenchmarkFastWalk")
+
+func BenchmarkFastWalk(b *testing.B) {
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		err := fastWalk(*benchDir, func(path string, typ os.FileMode) error { return nil })
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}

imports/fastwalk_unix.go

@@ -0,0 +1,122 @@
+// Copyright 2016 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// +build linux darwin freebsd openbsd netbsd
+
+package imports
+
+import (
+	"bytes"
+	"fmt"
+	"os"
+	"syscall"
+	"unsafe"
+)
+
+const blockSize = 8 << 10
+
+// unknownFileMode is a sentinel (and bogus) os.FileMode
+// value used to represent a syscall.DT_UNKNOWN Dirent.Type.
+const unknownFileMode os.FileMode = os.ModeNamedPipe | os.ModeSocket | os.ModeDevice
+
+func readDir(dirName string, fn func(dirName, entName string, typ os.FileMode) error) error {
+	fd, err := syscall.Open(dirName, 0, 0)
+	if err != nil {
+		return err
+	}
+	defer syscall.Close(fd)
+
+	// The buffer must be at least a block long.
+	buf := make([]byte, blockSize) // stack-allocated; doesn't escape
+	bufp := 0                      // starting read position in buf
+	nbuf := 0                      // end valid data in buf
+	for {
+		if bufp >= nbuf {
+			bufp = 0
+			nbuf, err = syscall.ReadDirent(fd, buf)
+			if err != nil {
+				return os.NewSyscallError("readdirent", err)
+			}
+			if nbuf <= 0 {
+				return nil
+			}
+		}
+		consumed, name, typ := parseDirEnt(buf[bufp:nbuf])
+		bufp += consumed
+		if name == "" || name == "." || name == ".." {
+			continue
+		}
+		// Fallback for filesystems (like old XFS) that don't
+		// support Dirent.Type and have DT_UNKNOWN (0) there
+		// instead.
+		if typ == unknownFileMode {
+			fi, err := os.Lstat(dirName + "/" + name)
+			if err != nil {
+				// It got deleted in the meantime.
+				if os.IsNotExist(err) {
+					continue
+				}
+				return err
+			}
+			typ = fi.Mode() & os.ModeType
+		}
+		if err := fn(dirName, name, typ); err != nil {
+			return err
+		}
+	}
+}
+
+func parseDirEnt(buf []byte) (consumed int, name string, typ os.FileMode) {
+	// golang.org/issue/15653
+	dirent := (*syscall.Dirent)(unsafe.Pointer(&buf[0]))
+	if v := unsafe.Offsetof(dirent.Reclen) + unsafe.Sizeof(dirent.Reclen); uintptr(len(buf)) < v {
+		panic(fmt.Sprintf("buf size of %d smaller than dirent header size %d", len(buf), v))
+	}
+	if len(buf) < int(dirent.Reclen) {
+		panic(fmt.Sprintf("buf size %d < record length %d", len(buf), dirent.Reclen))
+	}
+	consumed = int(dirent.Reclen)
+	if direntInode(dirent) == 0 { // File absent in directory.
+		return
+	}
+	switch dirent.Type {
+	case syscall.DT_REG:
+		typ = 0
+	case syscall.DT_DIR:
+		typ = os.ModeDir
+	case syscall.DT_LNK:
+		typ = os.ModeSymlink
+	case syscall.DT_BLK:
+		typ = os.ModeDevice
+	case syscall.DT_FIFO:
+		typ = os.ModeNamedPipe
+	case syscall.DT_SOCK:
+		typ = os.ModeSocket
+	case syscall.DT_UNKNOWN:
+		typ = unknownFileMode
+	default:
+		// Skip weird things.
+		// It's probably a DT_WHT (http://lwn.net/Articles/325369/)
+		// or something. Revisit if/when this package is moved outside
+		// of goimports. goimports only cares about regular files,
+		// symlinks, and directories.
+		return
+	}
+
+	nameBuf := (*[unsafe.Sizeof(dirent.Name)]byte)(unsafe.Pointer(&dirent.Name[0]))
+	nameLen := bytes.IndexByte(nameBuf[:], 0)
+	if nameLen < 0 {
+		panic("failed to find terminating 0 byte in dirent")
+	}
+
+	// Special cases for common things:
+	if nameLen == 1 && nameBuf[0] == '.' {
+		name = "."
+	} else if nameLen == 2 && nameBuf[0] == '.' && nameBuf[1] == '.' {
+		name = ".."
+	} else {
+		name = string(nameBuf[:nameLen])
+	}
+	return
+}

imports/fix.go

@@ -27,6 +27,11 @@ import (
 // Debug controls verbose logging.
 var Debug = false
 
+var (
+	inTests = false      // set true by fix_test.go; if false, no need to use testMu
+	testMu  sync.RWMutex // guards globals reset by tests; used only if inTests
+)
+
 // importToGroup is a list of functions which map from an import path to
 // a group number.
 var importToGroup = []func(importPath string) (num int, ok bool){
@@ -274,6 +279,10 @@ var (
 	// scanGoPathOnce guards calling scanGoPath (for $GOPATH)
 	scanGoPathOnce sync.Once
 
+	// populateIgnoreOnce guards calling populateIgnore
+	populateIgnoreOnce sync.Once
+	ignoredDirs        []os.FileInfo
+
 	dirScanMu sync.RWMutex
 	dirScan   map[string]*pkg // abs dir path => *pkg
 )
@@ -302,22 +311,34 @@ type gate chan struct{}
 func (g gate) enter() { g <- struct{}{} }
 func (g gate) leave() { <-g }
 
-// fsgate protects the OS & filesystem from too much concurrency.
-// Too much disk I/O -> too many threads -> swapping and bad scheduling.
-var fsgate = make(gate, 8)
-
 var visitedSymlinks struct {
 	sync.Mutex
 	m map[string]struct{}
 }
 
-var ignoredDirs []os.FileInfo
+// guarded by populateIgnoreOnce; populates ignoredDirs.
+func populateIgnore() {
+	for _, srcDir := range build.Default.SrcDirs() {
+		if srcDir == filepath.Join(build.Default.GOROOT, "src") {
+			continue
+		}
+		populateIgnoredDirs(srcDir)
+	}
+}
 
 // populateIgnoredDirs reads an optional config file at <path>/.goimportsignore
 // of relative directories to ignore when scanning for go files.
 // The provided path is one of the $GOPATH entries with "src" appended.
 func populateIgnoredDirs(path string) {
-	slurp, err := ioutil.ReadFile(filepath.Join(path, ".goimportsignore"))
+	file := filepath.Join(path, ".goimportsignore")
+	slurp, err := ioutil.ReadFile(file)
+	if Debug {
+		if err != nil {
+			log.Print(err)
+		} else {
+			log.Printf("Read %s", file)
+		}
+	}
 	if err != nil {
 		return
 	}
@@ -327,8 +348,14 @@ func populateIgnoredDirs(path string) {
 		if line == "" || strings.HasPrefix(line, "#") {
 			continue
 		}
-		if fi, err := os.Stat(filepath.Join(path, line)); err == nil {
+		full := filepath.Join(path, line)
+		if fi, err := os.Stat(full); err == nil {
 			ignoredDirs = append(ignoredDirs, fi)
+			if Debug {
+				log.Printf("Directory added to ignore list: %s", full)
+			}
+		} else if Debug {
+			log.Printf("Error statting entry in .goimportsignore: %v", err)
 		}
 	}
 }
@@ -342,22 +369,10 @@ func skipDir(fi os.FileInfo) bool {
 	return false
 }
 
-// shouldTraverse checks if fi, found in dir, is a directory or a symlink to a directory.
+// shouldTraverse reports whether the symlink fi should, found in dir,
-// It makes sure symlinks were never visited before to avoid symlink loops.
+// should be followed.  It makes sure symlinks were never visited
+// before to avoid symlink loops.
 func shouldTraverse(dir string, fi os.FileInfo) bool {
-	if fi.IsDir() {
-		if skipDir(fi) {
-			if Debug {
-				log.Printf("skipping directory %q under %s", fi.Name(), dir)
-			}
-			return false
-		}
-		return true
-	}
-
-	if fi.Mode()&os.ModeSymlink == 0 {
-		return false
-	}
 	path := filepath.Join(dir, fi.Name())
 	target, err := filepath.EvalSymlinks(path)
 	if err != nil {
@@ -395,7 +410,15 @@ func shouldTraverse(dir string, fi os.FileInfo) bool {
 
 var testHookScanDir = func(dir string) {}
 
-func scanGoRoot() { scanGoDirs(true) }
+var scanGoRootDone = make(chan struct{}) // closed when scanGoRoot is done
+
+func scanGoRoot() {
+	go func() {
+		scanGoDirs(true)
+		close(scanGoRootDone)
+	}()
+}
+
 func scanGoPath() { scanGoDirs(false) }
 
 func scanGoDirs(goRoot bool) {
@@ -413,95 +436,69 @@ func scanGoDirs(goRoot bool) {
 	}
 	dirScanMu.Unlock()
 
-	var wg sync.WaitGroup
+	for _, srcDir := range build.Default.SrcDirs() {
-	for _, path := range build.Default.SrcDirs() {
+		isGoroot := srcDir == filepath.Join(build.Default.GOROOT, "src")
-		isGoroot := path == filepath.Join(build.Default.GOROOT, "src")
 		if isGoroot != goRoot {
 			continue
 		}
-		if !goRoot {
+		testHookScanDir(srcDir)
-			populateIgnoredDirs(path)
+		walkFn := func(path string, typ os.FileMode) error {
+			dir := filepath.Dir(path)
+			if typ.IsRegular() {
+				if dir == srcDir {
+					// Doesn't make sense to have regular files
+					// directly in your $GOPATH/src or $GOROOT/src.
+					return nil
 				}
-		fsgate.enter()
+				if !strings.HasSuffix(path, ".go") {
-		testHookScanDir(path)
+					return nil
-		if Debug {
-			log.Printf("scanGoDir, open dir: %v\n", path)
 				}
-		f, err := os.Open(path)
-		if err != nil {
-			fsgate.leave()
-			fmt.Fprintf(os.Stderr, "goimports: scanning directories: %v\n", err)
-			continue
-		}
-		children, err := f.Readdir(-1)
-		f.Close()
-		fsgate.leave()
-		if err != nil {
-			fmt.Fprintf(os.Stderr, "goimports: scanning directory entries: %v\n", err)
-			continue
-		}
-		for _, child := range children {
-			if !shouldTraverse(path, child) {
-				continue
-			}
-			wg.Add(1)
-			go func(path, name string) {
-				defer wg.Done()
-				scanDir(&wg, path, name)
-			}(path, child.Name())
-		}
-	}
-	wg.Wait()
-}
-
-func scanDir(wg *sync.WaitGroup, root, pkgrelpath string) {
-	importpath := filepath.ToSlash(pkgrelpath)
-	dir := filepath.Join(root, importpath)
-
-	fsgate.enter()
-	defer fsgate.leave()
-	if Debug {
-		log.Printf("scanning dir %s", dir)
-	}
-	pkgDir, err := os.Open(dir)
-	if err != nil {
-		return
-	}
-	children, err := pkgDir.Readdir(-1)
-	pkgDir.Close()
-	if err != nil {
-		return
-	}
-	// hasGo tracks whether a directory actually appears to be a
-	// Go source code directory. If $GOPATH == $HOME, and
-	// $HOME/src has lots of other large non-Go projects in it,
-	// then the calls to importPathToName below can be expensive.
-	hasGo := false
-	for _, child := range children {
-		// Avoid .foo, _foo, and testdata directory trees.
-		name := child.Name()
-		if name == "" || name[0] == '.' || name[0] == '_' || name == "testdata" {
-			continue
-		}
-		if strings.HasSuffix(name, ".go") {
-			hasGo = true
-		}
-		if shouldTraverse(dir, child) {
-			wg.Add(1)
-			go func(root, name string) {
-				defer wg.Done()
-				scanDir(wg, root, name)
-			}(root, filepath.Join(importpath, name))
-		}
-	}
-	if hasGo {
 				dirScanMu.Lock()
+				if _, dup := dirScan[dir]; !dup {
+					importpath := filepath.ToSlash(dir[len(srcDir)+len("/"):])
 					dirScan[dir] = &pkg{
 						importPath:      importpath,
 						importPathShort: vendorlessImportPath(importpath),
 						dir:             dir,
 					}
+				}
 				dirScanMu.Unlock()
+				return nil
+			}
+			if typ == os.ModeDir {
+				base := filepath.Base(path)
+				if base == "" || base[0] == '.' || base[0] == '_' || base == "testdata" {
+					return filepath.SkipDir
+				}
+				fi, err := os.Lstat(path)
+				if err == nil && skipDir(fi) {
+					if Debug {
+						log.Printf("skipping directory %q under %s", fi.Name(), dir)
+					}
+					return filepath.SkipDir
+				}
+				return nil
+			}
+			if typ == os.ModeSymlink {
+				base := filepath.Base(path)
+				if strings.HasPrefix(base, ".#") {
+					// Emacs noise.
+					return nil
+				}
+				fi, err := os.Lstat(path)
+				if err != nil {
+					// Just ignore it.
+					return nil
+				}
+				if shouldTraverse(dir, fi) {
+					return traverseLink
+				}
+			}
+			return nil
+		}
+		if err := fastWalk(srcDir, walkFn); err != nil {
+			log.Printf("goimports: scanning directory %v: %v", srcDir, err)
+		}
 	}
 }
 
@@ -616,6 +613,11 @@ var findImport func(pkgName string, symbols map[string]bool, filename string) (f
 // findImportGoPath is the normal implementation of findImport.
 // (Some companies have their own internally.)
 func findImportGoPath(pkgName string, symbols map[string]bool, filename string) (foundPkg string, rename bool, err error) {
+	if inTests {
+		testMu.RLock()
+		defer testMu.RUnlock()
+	}
+
 	// Fast path for the standard library.
 	// In the common case we hopefully never have to scan the GOPATH, which can
 	// be slow with moving disks.
@@ -638,11 +640,22 @@ func findImportGoPath(pkgName string, symbols map[string]bool, filename string)
 	// in the current Go file.  Return rename=true when the other Go files
 	// use a renamed package that's also used in the current file.
 
-	scanGoRootOnce.Do(scanGoRoot)
+	// Read all the $GOPATH/src/.goimportsignore files before scanning directories.
-	if !strings.HasPrefix(filename, build.Default.GOROOT) {
+	populateIgnoreOnce.Do(populateIgnore)
-		scanGoPathOnce.Do(scanGoPath)
-	}
 
+	// Start scanning the $GOROOT asynchronously, then run the
+	// GOPATH scan synchronously if needed, and then wait for the
+	// $GOROOT to finish.
+	//
+	// TODO(bradfitz): run each $GOPATH entry async. But nobody
+	// really has more than one anyway, so low priority.
+	scanGoRootOnce.Do(scanGoRoot) // async
+	if !strings.HasPrefix(filename, build.Default.GOROOT) {
+		scanGoPathOnce.Do(scanGoPath) // blocking
+	}
+	<-scanGoRootDone
+
+	// Find candidate packages, looking only at their directory names first.
 	var candidates []*pkg
 	for _, pkg := range dirScan {
 		if pkgIsCandidate(filename, pkgName, pkg) {
@@ -650,6 +663,10 @@ func findImportGoPath(pkgName string, symbols map[string]bool, filename string)
 		}
 	}
 
+	// 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
+	// there's no "penalty" for vendoring.
 	sort.Sort(byImportPathShortLength(candidates))
 	if Debug {
 		for i, pkg := range candidates {
@@ -664,7 +681,7 @@ func findImportGoPath(pkgName string, symbols map[string]bool, filename string)
 
 	rescv := make([]chan *pkg, len(candidates))
 	for i := range candidates {
-		rescv[i] = make(chan *pkg, 1)
+		rescv[i] = make(chan *pkg)
 	}
 	const maxConcurrentPackageImport = 4
 	loadExportsSem := make(chan struct{}, maxConcurrentPackageImport)
@@ -673,12 +690,20 @@ func findImportGoPath(pkgName string, symbols map[string]bool, filename string)
 		for i, pkg := range candidates {
 			select {
 			case loadExportsSem <- struct{}{}:
+				select {
+				case <-done:
+				default:
+				}
 			case <-done:
 				return
 			}
 			pkg := pkg
 			resc := rescv[i]
 			go func() {
+				if inTests {
+					testMu.RLock()
+					defer testMu.RUnlock()
+				}
 				defer func() { <-loadExportsSem }()
 				exports := loadExports(pkgName, pkg.dir)
 
@@ -690,7 +715,10 @@ func findImportGoPath(pkgName string, symbols map[string]bool, filename string)
 						break
 					}
 				}
-				resc <- pkg
+				select {
+				case resc <- pkg:
+				case <-done:
+				}
 			}()
 		}
 	}()

imports/fix_test.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
+	"strings"
 	"sync"
 	"testing"
 )
@@ -981,12 +982,20 @@ type Buffer2 struct {}
 	})
 }
 
+func init() {
+	inTests = true
+}
+
 func withEmptyGoPath(fn func()) {
+	testMu.Lock()
+
 	dirScanMu.Lock()
+	populateIgnoreOnce = sync.Once{}
 	scanGoRootOnce = sync.Once{}
 	scanGoPathOnce = sync.Once{}
 	dirScan = nil
 	ignoredDirs = nil
+	scanGoRootDone = make(chan struct{})
 	dirScanMu.Unlock()
 
 	oldGOPATH := build.Default.GOPATH
@@ -994,11 +1003,16 @@ func withEmptyGoPath(fn func()) {
 	build.Default.GOPATH = ""
 	visitedSymlinks.m = nil
 	testHookScanDir = func(string) {}
+	testMu.Unlock()
+
 	defer func() {
+		testMu.Lock()
 		testHookScanDir = func(string) {}
 		build.Default.GOPATH = oldGOPATH
 		build.Default.GOROOT = oldGOROOT
+		testMu.Unlock()
 	}()
+
 	fn()
 }
 
@@ -1162,7 +1176,13 @@ func mapToDir(destDir string, files map[string]string) error {
 		if err := os.MkdirAll(filepath.Dir(file), 0755); err != nil {
 			return err
 		}
-		if err := ioutil.WriteFile(file, []byte(contents), 0644); err != nil {
+		var err error
+		if strings.HasPrefix(contents, "LINK:") {
+			err = os.Symlink(strings.TrimPrefix(contents, "LINK:"), file)
+		} else {
+			err = ioutil.WriteFile(file, []byte(contents), 0644)
+		}
+		if err != nil {
 			return err
 		}
 	}