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:
https://github.com/golang/go/issues/16399#issuecomment-293278556
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 <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Joël Stemmer <jstemmer@google.com>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
diff --git a/imports/fastwalk.go b/imports/fastwalk.go
index 157c792..31e6e27 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,13 @@
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 @@
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 @@
// 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 @@
}
return err
}
+
func (w *walker) walk(root string, runUserCallback bool) error {
if runUserCallback {
err := w.fn(root, os.ModeDir)
