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>
This commit is contained in:
Brad Fitzpatrick 2017-04-25 16:26:45 +00:00
parent 75e5ff36f3
commit 81478017b6
1 changed files with 19 additions and 4 deletions

View File

@ -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)