imports: wait for fastWalk workers to finish before returning

In some cases walkFn is being called after the fastWalk function has
returned. This often happens when an error was encountered early on in
scanning directories with many entries.

It is caused by fastWalk not waiting for its workers to complete their
work. A sync.WaitGroup is used to wait for all workers to finish when
the function returns.

Updates golang/go#16399

Change-Id: I695d30c18e4878b789520b9d8a650f9688d896ac
Reviewed-on: https://go-review.googlesource.com/40092
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
diff --git a/imports/fastwalk.go b/imports/fastwalk.go
index 157c792..c8a7949 100644
--- a/imports/fastwalk.go
+++ b/imports/fastwalk.go
@@ -19,6 +19,7 @@
 	"os"
 	"path/filepath"
 	"runtime"
+	"sync"
 )
 
 // traverseLink is a sentinel error for fastWalk, similar to filepath.SkipDir.
@@ -48,6 +49,12 @@
 	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.
+	var wg sync.WaitGroup
+	defer wg.Wait()
+
 	w := &walker{
 		fn:       walkFn,
 		enqueuec: make(chan walkItem, numWorkers), // buffered for performance
@@ -60,7 +67,8 @@
 	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,7 +111,8 @@
 
 // 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: