fix(snapshots): fixed random deadlock when Uploader results in a failure (#2020)

* fix(snapshots): fixed random deadlock when Uploader results in a failure

The deadlock was caused by not properly waiting for all asynchronous
work to complete before closing the worker pool.

Introduced `workshare.AsyncGroup.Close()` and some assertions.

* fixed select race

* linter fix

* pr feedback

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)

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
+		}
 	}
 }

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

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() {

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.