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:
parent
75e5ff36f3
commit
81478017b6
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in New Issue