diff --git a/internal/workshare/workshare_test.go b/internal/workshare/workshare_test.go index 1b41b5f87..c5131a557 100644 --- a/internal/workshare/workshare_test.go +++ b/internal/workshare/workshare_test.go @@ -102,6 +102,44 @@ func TestComputeTreeSumNegative(t *testing.T) { testComputeTreeSum(t, -1) } +func TestDisallowed_DoubleWait(t *testing.T) { + var ag workshare.AsyncGroup + + ag.Wait() + require.Panics(t, func() { + ag.Wait() + }) +} + +func TestDisallowed_WaitAfterClose(t *testing.T) { + var ag workshare.AsyncGroup + + ag.Close() + require.Panics(t, func() { + ag.Wait() + }) + + ag.Close() // no-op +} + +func TestDisallowed_UseAfterPoolClose(t *testing.T) { + w := workshare.NewPool(1) + + var ag workshare.AsyncGroup + + w.Close() + + require.Panics(t, func() { + ag.CanShareWork(w) + }) + + require.Panics(t, func() { + ag.RunAsync(w, func(c *workshare.Pool, request interface{}) { + t.Fatal("should not be called") + }, nil) + }) +} + // nolint:thelper func testComputeTreeSum(t *testing.T, numWorkers int) { w := workshare.NewPool(numWorkers) diff --git a/internal/workshare/workshare_waitgroup.go b/internal/workshare/workshare_waitgroup.go index ffdbfbd1f..212a76158 100644 --- a/internal/workshare/workshare_waitgroup.go +++ b/internal/workshare/workshare_waitgroup.go @@ -14,6 +14,7 @@ // // func visitNode(p *workshare.Pool, n *someNode) error { // var wg workshare.AsyncGroup +// defer wg.Close() // // for _, child := range n.children { // if wg.CanShareWork(p) { @@ -52,10 +53,22 @@ import ( type AsyncGroup struct { wg *sync.WaitGroup requests []interface{} + waited bool + isClosed bool } // Wait waits for scheduled asynchronous work to complete and returns all asynchronously processed inputs. func (g *AsyncGroup) Wait() []interface{} { + if g.isClosed { + panic("Wait() can't be called after Close()") + } + + if g.waited { + panic("Wait() can't be called more than once") + } + + g.waited = true + if g.wg == nil { return nil } @@ -65,6 +78,19 @@ func (g *AsyncGroup) Wait() []interface{} { return g.requests } +// Close ensures all asynchronous work has been awaited for. +func (g *AsyncGroup) Close() { + if g.isClosed { + return + } + + if !g.waited { + g.Wait() // ensure we wait + } + + g.isClosed = true +} + // RunAsync starts the asynchronous work to process the provided request, // the user must call Wait() after all RunAsync() have been scheduled. func (g *AsyncGroup) RunAsync(w *Pool, process ProcessFunc, request interface{}) { @@ -78,10 +104,15 @@ func (g *AsyncGroup) RunAsync(w *Pool, process ProcessFunc, request interface{}) g.requests = append(g.requests, request) } - w.work <- workItem{ + select { + case w.work <- workItem{ process: process, request: request, wg: g.wg, + }: + + case <-w.closed: + panic("invalid usage - RunAsync() after workshare.AsyncGroup has been closed") } } @@ -89,15 +120,22 @@ func (g *AsyncGroup) RunAsync(w *Pool, process ProcessFunc, request interface{}) // If the function returns true, the use MUST call RunAsync() exactly once. This pattern avoids // allocations required to create asynchronous input if the worker pool is full. func (g *AsyncGroup) CanShareWork(w *Pool) bool { + // check pool is not closed select { - case w.semaphore <- struct{}{}: - // we successfully added token to the channel, because we have exactly the same number - // of workers as the capacity of the channel, one worker will wake up to process - // item from w.work, which will be added by RunAsync(). - return true + case <-w.closed: + panic("invalid usage - CanShareWork() after workshare.AsyncGroup has been closed") default: - // all workers are busy. - return false + select { + case w.semaphore <- struct{}{}: + // we successfully added token to the channel, because we have exactly the same number + // of workers as the capacity of the channel, one worker will wake up to process + // item from w.work, which will be added by RunAsync(). + return true + + default: + // all workers are busy. + return false + } } } diff --git a/snapshot/snapshotfs/dir_rewriter.go b/snapshot/snapshotfs/dir_rewriter.go index f24217e3a..f7dc8be81 100644 --- a/snapshot/snapshotfs/dir_rewriter.go +++ b/snapshot/snapshotfs/dir_rewriter.go @@ -143,6 +143,9 @@ func (rw *DirRewriter) processDirectoryEntries(ctx context.Context, parentPath s wg workshare.AsyncGroup ) + // ensure we wait for all work items before returning + defer wg.Close() + for _, child := range entries { if wg.CanShareWork(rw.ws) { // see if we can run this child in a goroutine diff --git a/snapshot/snapshotfs/snapshot_tree_walker.go b/snapshot/snapshotfs/snapshot_tree_walker.go index 0adfe2e7b..9757bd03b 100644 --- a/snapshot/snapshotfs/snapshot_tree_walker.go +++ b/snapshot/snapshotfs/snapshot_tree_walker.go @@ -110,7 +110,7 @@ func (w *TreeWalker) processDirEntry(ctx context.Context, dir fs.Directory, entr } var ag workshare.AsyncGroup - defer ag.Wait() + defer ag.Close() err := dir.IterateEntries(ctx, func(c context.Context, ent fs.Entry) error { if w.TooManyErrors() { diff --git a/snapshot/snapshotfs/upload.go b/snapshot/snapshotfs/upload.go index 64557ec62..0eb8d4ee6 100644 --- a/snapshot/snapshotfs/upload.go +++ b/snapshot/snapshotfs/upload.go @@ -538,6 +538,9 @@ func (u *Uploader) processChildren( ) error { var wg workshare.AsyncGroup + // ensure we wait for all work items before returning + defer wg.Close() + // ignore errCancel because a more serious error may be reported in wg.Wait() // we'll check for cancelation later.