KUBEDR-7763: Don't panic if CanShareWork() is called on a closed work pool

Author: dzaninovic

internal/workshare/workshare_pool.go

@@ -2,6 +2,7 @@
 package workshare
 
 import (
+	"context"
 	"sync"
 	"sync/atomic"
 )
@@ -21,6 +22,7 @@ type workItem[T any] struct {
 type Pool[T any] struct {
 	activeWorkers atomic.Int32
 
+	ctx       context.Context
 	semaphore chan struct{}
 
 	work   chan workItem[T]
@@ -35,12 +37,13 @@ func (w *Pool[T]) ActiveWorkers() int {
 }
 
 // NewPool creates a worker pool that launches a given number of goroutines that can invoke shared work.
-func NewPool[T any](numWorkers int) *Pool[T] {
+func NewPool[T any](ctx context.Context, numWorkers int) *Pool[T] {
 	if numWorkers < 0 {
 		numWorkers = 0
 	}
 
 	w := &Pool[T]{
+		ctx: ctx,
 		// channel must be unbuffered so that it has exactly as many slots as there are goroutines capable of reading from it
 		// this way by pushing to the channel we can be sure that a pre-spun goroutine will pick it up soon.
 		work:      make(chan workItem[T]),

internal/workshare/workshare_test.go

@@ -1,6 +1,7 @@
 package workshare_test
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -119,15 +120,13 @@ func TestDisallowed_WaitAfterClose(t *testing.T) {
 }
 
 func TestDisallowed_UseAfterPoolClose(t *testing.T) {
-	w := workshare.NewPool[int](1)
+	w := workshare.NewPool[int](context.Background(), 1)
 
 	var ag workshare.AsyncGroup[int]
 
 	w.Close()
 
-	require.Panics(t, func() {
-		ag.CanShareWork(w)
-	})
+	require.False(t, ag.CanShareWork(w))
 
 	require.Panics(t, func() {
 		ag.RunAsync(w, func(c *workshare.Pool[int], request int) {
@@ -138,7 +137,7 @@ func TestDisallowed_UseAfterPoolClose(t *testing.T) {
 
 //nolint:thelper
 func testComputeTreeSum(t *testing.T, numWorkers int) {
-	w := workshare.NewPool[*computeTreeSumRequest](numWorkers)
+	w := workshare.NewPool[*computeTreeSumRequest](context.Background(), numWorkers)
 	defer w.Close()
 
 	n := buildTree(6)
@@ -151,7 +150,7 @@ func testComputeTreeSum(t *testing.T, numWorkers int) {
 var treeToWalk = buildTree(6)
 
 func BenchmarkComputeTreeSum(b *testing.B) {
-	w := workshare.NewPool[*computeTreeSumRequest](10)
+	w := workshare.NewPool[*computeTreeSumRequest](context.Background(), 10)
 	defer w.Close()
 
 	b.ResetTimer()

internal/workshare/workshare_waitgroup.go

@@ -44,8 +44,12 @@ package workshare
 
 import (
 	"sync"
+
+	"github.com/kopia/kopia/repo/logging"
 )
 
+var log = logging.Module("workshare")
+
 // AsyncGroup launches and awaits asynchronous work through a WorkerPool.
 // It provides API designed to minimize allocations while being reasonably easy to use.
 // AsyncGroup is very lightweight and NOT safe for concurrent use, all interactions must
@@ -121,7 +125,8 @@ func (g *AsyncGroup[T]) CanShareWork(w *Pool[T]) bool {
 	// check pool is not closed
 	select {
 	case <-w.closed:
-		panic("invalid usage - CanShareWork() after workshare.AsyncGroup has been closed")
+		log(w.ctx).Warn("CanShareWork() is called after workshare.Pool has been closed")
+		return false
 
 	default:
 		select {

snapshot/snapshotfs/dir_rewriter.go

@@ -338,7 +338,7 @@ func NewDirRewriter(ctx context.Context, rep repo.RepositoryWriter, opts DirRewr
 	}
 
 	return &DirRewriter{
-		ws:    workshare.NewPool[*dirRewriterRequest](opts.Parallel - 1),
+		ws:    workshare.NewPool[*dirRewriterRequest](ctx, opts.Parallel-1),
 		opts:  opts,
 		rep:   rep,
 		cache: cache,

snapshot/snapshotfs/snapshot_tree_walker.go

@@ -193,7 +193,7 @@ func NewTreeWalker(ctx context.Context, options TreeWalkerOptions) (*TreeWalker,
 
 	return &TreeWalker{
 		options:  options,
-		wp:       workshare.NewPool[any](options.Parallelism - 1),
+		wp:       workshare.NewPool[any](ctx, options.Parallelism-1),
 		enqueued: s,
 	}, nil
 }

snapshot/snapshotfs/upload.go

@@ -1267,7 +1267,7 @@ func (u *Uploader) Upload(
 		Source: sourceInfo,
 	}
 
-	u.workerPool = workshare.NewPool[*uploadWorkItem](parallel - 1)
+	u.workerPool = workshare.NewPool[*uploadWorkItem](ctx, parallel-1)
 	defer u.workerPool.Close()
 
 	u.stats = &snapshot.Stats{}