diff --git a/imports/fastwalk.go b/imports/fastwalk.go index 157c7922..31e6e27b 100644 --- a/imports/fastwalk.go +++ b/imports/fastwalk.go @@ -19,6 +19,7 @@ import ( "os" "path/filepath" "runtime" + "sync" ) // traverseLink is a sentinel error for fastWalk, similar to filepath.SkipDir. @@ -48,6 +49,13 @@ func fastWalk(root string, walkFn func(path string, typ os.FileMode) error) erro if n := runtime.NumCPU(); n > numWorkers { numWorkers = n } + + // Make sure to wait for all workers to finish, otherwise + // walkFn could still be called after returning. This Wait call + // runs after close(e.donec) below. + var wg sync.WaitGroup + defer wg.Wait() + w := &walker{ fn: walkFn, enqueuec: make(chan walkItem, numWorkers), // buffered for performance @@ -58,9 +66,10 @@ func fastWalk(root string, walkFn func(path string, typ os.FileMode) error) erro 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() + wg.Add(1) + go w.doWork(&wg) } todo := []walkItem{{dir: root}} out := 0 @@ -103,13 +112,18 @@ func fastWalk(root string, walkFn func(path string, typ os.FileMode) error) erro // doWork reads directories as instructed (via workc) and runs the // user's callback function. -func (w *walker) doWork() { +func (w *walker) doWork(wg *sync.WaitGroup) { + defer wg.Done() for { select { case <-w.donec: return case it := <-w.workc: - w.resc <- w.walk(it.dir, !it.callbackDone) + select { + case <-w.donec: + return + case w.resc <- w.walk(it.dir, !it.callbackDone): + } } } } @@ -157,6 +171,7 @@ func (w *walker) onDirEnt(dirName, baseName string, typ os.FileMode) error { } return err } + func (w *walker) walk(root string, runUserCallback bool) error { if runUserCallback { err := w.fn(root, os.ModeDir)