imports: wait for fastWalk workers to finish before returning (take 2)

This is Joël Stemmer's https://golang.org/cl/40092 again, but with
a fix to prevent workers from deadlocking on send if the caller had
already started to shut down. See:

golang/go#16399 (comment)

Updates golang/go#16399
Fixes golang/go#20109 (it looks like)

Change-Id: I3d1cf6f24563d02e1369a4496c2d37dcc1f5e5b8
Reviewed-on: https://go-review.googlesource.com/41681
Run-TryBot: Brad Fitzpatrick <[email protected]>
TryBot-Result: Gobot Gobot <[email protected]>
Reviewed-by: Joël Stemmer <[email protected]>
Reviewed-by: Brad Fitzpatrick <[email protected]>

Author: bradfitz

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)