From 5f54c75cf817ee5059f37453bc7cfc571931d398 Mon Sep 17 00:00:00 2001 From: Dmitry Verkhoturov Date: Wed, 19 Aug 2026 06:15:18 +0100 Subject: [PATCH 1/4] run pool completion after a worker error, not by chance a failing worker cancels the errgroup context, so a surviving worker exits with context.Canceled and suppresses poolCompleteFn, while the same worker exiting through the drained channels reports nil and runs it. which one happens is up to the scheduler, so a chained pool got its next stage closed about half the time. decide the suppression on the context passed to Go, which only the caller can cancel, instead of on lastErr. --- README.md | 3 ++- pool.go | 15 +++++++++++---- pool_test.go | 33 +++++++++++++++++++++++++++++++++ 3 files changed, 46 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index fd8e912..4365627 100644 --- a/README.md +++ b/README.md @@ -526,7 +526,8 @@ p := pool.New[string](5, worker). The completion callback executes when: - All workers have completed processing - Errors occurred but pool continued (`WithContinueOnError()`) -- Skipped only on `context.Canceled` (still runs on `context.DeadlineExceeded`) +- Skipped only when the context passed to `Go` was cancelled (still runs on `context.DeadlineExceeded`, + and on a worker error, which cancels the pool's internal context but not the caller's) Important notes: - Use `Submit` when sending items from a single goroutine diff --git a/pool.go b/pool.go index 8aea9c1..826aecb 100644 --- a/pool.go +++ b/pool.go @@ -436,10 +436,17 @@ func (p *WorkerGroup[T]) finishWorker(ctx context.Context, id int, worker Worker activeWorkers := p.activeWorkers.Add(-1) - // pool completion should be called when this is the last worker - // regardless of error state, except for context cancellation - if activeWorkers == 0 && p.poolCompleteFn != nil && !errors.Is(lastErr, context.Canceled) { - if e := p.poolCompleteFn(ctx); e != nil { + // pool completion should be called when this is the last worker regardless of error state, + // except when the caller cancelled. checked on callerCtx rather than on lastErr, which also + // carries the errgroup's cancellation after a peer worker failed + if activeWorkers == 0 && p.poolCompleteFn != nil && !errors.Is(p.callerCtx.Err(), context.Canceled) { + completeCtx := ctx + if p.callerCtx.Err() == nil { + // ctx may be cancelled because a peer worker failed, which should not stop the + // callback from closing the next pool. values, metrics among them, are kept + completeCtx = context.WithoutCancel(ctx) + } + if e := p.poolCompleteFn(completeCtx); e != nil { if lastErr == nil { lastErr = fmt.Errorf("complete pool func for %d failed: %w", id, e) } diff --git a/pool_test.go b/pool_test.go index 994b152..855b838 100644 --- a/pool_test.go +++ b/pool_test.go @@ -1547,6 +1547,39 @@ func TestPool_PoolCompletion(t *testing.T) { assert.False(t, completeCalled.Load(), "pool completion must not run on a cancelled pool") } }) + + t.Run("worker error still runs pool completion", func(t *testing.T) { + // a failing worker cancels the errgroup context, which the surviving worker sees on its + // next select. that must not be taken for the caller cancelling the pool + var completeCalled atomic.Bool + var completeCtxErr error + errFailed := errors.New("failed") + + p := New[string](2, WorkerFunc[string](func(_ context.Context, v string) error { + if v == "fail" { + return errFailed + } + return nil + })).WithBatchSize(0).WithPoolCompleteFn(func(ctx context.Context) error { + completeCalled.Store(true) + completeCtxErr = ctx.Err() + return nil + }) + require.NoError(t, p.Go(context.Background())) + + p.Submit("ok") + p.Submit("fail") + + <-p.ctx.Done() // the failing worker has returned and the errgroup cancelled the pool context + + // the channels are still open, so the surviving worker can only leave through wCtx.Done + require.Eventually(t, func() bool { return p.activeWorkers.Load() == 0 }, time.Second, time.Millisecond) + + err := p.Close(context.Background()) + require.ErrorIs(t, err, errFailed) + assert.True(t, completeCalled.Load(), "pool completion must run when the caller did not cancel") + assert.NoError(t, completeCtxErr, "the callback must get a context it can still work with") + }) } func TestPool_ChainedBatching(t *testing.T) { From 4660fc2dd2d386b2ef1c817f3099345968dc28ab Mon Sep 17 00:00:00 2001 From: Dmitry Verkhoturov Date: Thu, 20 Aug 2026 11:07:33 +0100 Subject: [PATCH 2/4] keep the completion callback cancellable and report its error the callback ran on context.WithoutCancel whenever the caller was still active, including a normal drain where nothing was cancelled, so a blocking callback could not be stopped by the caller at all. cancel is now stripped only when the pool context is already cancelled and the caller's is not, and the caller's cancellation is bridged into the callback context. poolCompleteFn's error was dropped whenever a worker had already failed. it is kept on the group and joined with the errgroup result, which retains the first worker error only and would otherwise lose it. --- README.md | 5 +++++ pool.go | 34 ++++++++++++++++++++-------- pool_test.go | 62 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 92 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 4365627..a634715 100644 --- a/README.md +++ b/README.md @@ -529,6 +529,11 @@ The completion callback executes when: - Skipped only when the context passed to `Go` was cancelled (still runs on `context.DeadlineExceeded`, and on a worker error, which cancels the pool's internal context but not the caller's) +After a worker error the callback receives a context that is no longer cancelled by that error, so it can +still close the next pool in a chain. Cancelling the context passed to `Go` cancels it as well, so a +callback that blocks can always be stopped. An error returned by the callback is reported alongside the +worker error rather than replacing it. + Important notes: - Use `Submit` when sending items from a single goroutine - Use `Send` when workers need to submit items to another pool diff --git a/pool.go b/pool.go index 826aecb..f23cab6 100644 --- a/pool.go +++ b/pool.go @@ -45,6 +45,11 @@ type WorkerGroup[T any] struct { ctx context.Context // errgroup context, cancelled by a failing worker as well as by the caller callerCtx context.Context // context passed to Go, cancelled only by the caller + // completeErr holds what poolCompleteFn returned. kept apart from the errgroup, which retains + // the first worker error only and would drop a completion failure reported after it. + // written by the last worker to finish, read once eg.Wait has returned + completeErr error + sendMu sync.Mutex } @@ -441,15 +446,19 @@ func (p *WorkerGroup[T]) finishWorker(ctx context.Context, id int, worker Worker // carries the errgroup's cancellation after a peer worker failed if activeWorkers == 0 && p.poolCompleteFn != nil && !errors.Is(p.callerCtx.Err(), context.Canceled) { completeCtx := ctx - if p.callerCtx.Err() == nil { - // ctx may be cancelled because a peer worker failed, which should not stop the - // callback from closing the next pool. values, metrics among them, are kept - completeCtx = context.WithoutCancel(ctx) + if ctx.Err() != nil && p.callerCtx.Err() == nil { + // ctx is cancelled because a peer worker failed, which should not stop the callback + // from closing the next pool. values, metrics among them, are kept, and the caller's + // own cancellation is bridged back so a blocking callback still ends when asked to + var cancel context.CancelFunc + completeCtx, cancel = context.WithCancel(context.WithoutCancel(ctx)) + defer cancel() + //nolint:contextcheck // callerCtx is deliberate, it is the only context left that the caller still controls + stop := context.AfterFunc(p.callerCtx, cancel) + defer stop() } if e := p.poolCompleteFn(completeCtx); e != nil { - if lastErr == nil { - lastErr = fmt.Errorf("complete pool func for %d failed: %w", id, e) - } + p.completeErr = fmt.Errorf("complete pool func for %d failed: %w", id, e) } } @@ -459,6 +468,13 @@ func (p *WorkerGroup[T]) finishWorker(ctx context.Context, id int, worker Worker return nil } +// waitWorkers waits for every worker and adds the pool completion error, which the errgroup +// cannot carry because it keeps the first error only. +func (p *WorkerGroup[T]) waitWorkers() error { + err := p.eg.Wait() + return errors.Join(err, p.completeErr) +} + // Close pool. Has to be called by consumer as the indication of "all records submitted". // The call is blocking till all processing completed by workers or context is cancelled. // After this call pool can't be reused. Returns an error if any happened during the run. @@ -483,7 +499,7 @@ func (p *WorkerGroup[T]) Close(ctx context.Context) error { // wait for workers with context respect done := make(chan error, 1) go func() { - done <- p.eg.Wait() + done <- p.waitWorkers() }() select { @@ -538,7 +554,7 @@ func (p *WorkerGroup[T]) Wait(ctx context.Context) error { // wait for workers with context respect done := make(chan error, 1) go func() { - done <- p.eg.Wait() + done <- p.waitWorkers() }() select { diff --git a/pool_test.go b/pool_test.go index 855b838..97485e3 100644 --- a/pool_test.go +++ b/pool_test.go @@ -1580,6 +1580,68 @@ func TestPool_PoolCompletion(t *testing.T) { assert.True(t, completeCalled.Load(), "pool completion must run when the caller did not cancel") assert.NoError(t, completeCtxErr, "the callback must get a context it can still work with") }) + + t.Run("caller cancellation reaches a running completion callback", func(t *testing.T) { + // the callback runs on a context stripped of the peer worker's cancellation, so the + // caller's own cancellation has to be bridged into it or a blocking callback never ends + started, unblocked := make(chan struct{}), make(chan struct{}) + errFailed := errors.New("failed") + + callerCtx, cancelCaller := context.WithCancel(context.Background()) + defer cancelCaller() + + p := New[string](2, WorkerFunc[string](func(_ context.Context, v string) error { + if v == "fail" { + return errFailed + } + return nil + })).WithBatchSize(0).WithPoolCompleteFn(func(ctx context.Context) error { + close(started) + <-ctx.Done() + close(unblocked) + return nil + }) + require.NoError(t, p.Go(callerCtx)) + + p.Submit("ok") + p.Submit("fail") + + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("pool completion never started") + } + + cancelCaller() + + select { + case <-unblocked: + case <-time.After(time.Second): + t.Fatal("pool completion did not observe the caller cancelling") + } + + require.ErrorIs(t, p.Close(context.Background()), errFailed) + }) + + t.Run("completion error reported alongside a worker error", func(t *testing.T) { + // lastErr is already set on this path, the completion failure must not be swallowed + errFailed, errComplete := errors.New("failed"), errors.New("complete failed") + + p := New[string](2, WorkerFunc[string](func(_ context.Context, v string) error { + if v == "fail" { + return errFailed + } + return nil + })).WithBatchSize(0).WithPoolCompleteFn(func(context.Context) error { return errComplete }) + require.NoError(t, p.Go(context.Background())) + + p.Submit("ok") + p.Submit("fail") + + err := p.Close(context.Background()) + require.ErrorIs(t, err, errFailed, "the worker error must still be reported") + require.ErrorIs(t, err, errComplete, "the completion error must not be dropped") + }) } func TestPool_ChainedBatching(t *testing.T) { From 4dfd9571b3a327e89562769353732212194bb0dc Mon Sep 17 00:00:00 2001 From: Dmitry Verkhoturov Date: Thu, 20 Aug 2026 11:08:14 +0100 Subject: [PATCH 3/4] record the published pool checksums in the example go.sum files with the replace directive active go mod tidy resolves the module from ../.. and records no checksum for it, so an example copied out of the tree failed to build with "missing go.sum entry" until the user ran go get. the v0.9.2 module and go.mod hashes are now present in all ten. note for later bumps: go mod tidy run with the replace directive in place drops these lines again, they have to be regenerated with it temporarily removed. --- examples/basic/go.sum | 2 ++ examples/chunking/go.sum | 2 ++ examples/collector_errors/go.sum | 2 ++ examples/collectors_chain/go.sum | 2 ++ examples/direct_chain/go.sum | 2 ++ examples/middleware/go.sum | 2 ++ examples/parallel_files/go.sum | 2 ++ examples/pool_completion/go.sum | 2 ++ examples/tokenizer_stateful/go.sum | 2 ++ examples/tokenizer_stateless/go.sum | 2 ++ 10 files changed, 20 insertions(+) diff --git a/examples/basic/go.sum b/examples/basic/go.sum index df47ecd..f35fbcf 100644 --- a/examples/basic/go.sum +++ b/examples/basic/go.sum @@ -1,3 +1,5 @@ +github.com/go-pkgz/pool v0.9.2 h1:VJ9rJDYTFKbp1/wml/7XlLBa8huL5/IeK+1aUf23ugw= +github.com/go-pkgz/pool v0.9.2/go.mod h1:HpVwnbSym5sbYVU/N460+GBixeMgThvbAhMSPsTlkZE= github.com/stretchr/testify v1.12.0 h1:K6Mr6jO9JICuend/5xzTM03ydSV3vdNRYAdPSukj8uI= github.com/stretchr/testify v1.12.0/go.mod h1:bOYBZb5qJ00vPzWfIqBUZPaxK8jWiXc6d3ErP4Ca9Gw= golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= diff --git a/examples/chunking/go.sum b/examples/chunking/go.sum index df47ecd..f35fbcf 100644 --- a/examples/chunking/go.sum +++ b/examples/chunking/go.sum @@ -1,3 +1,5 @@ +github.com/go-pkgz/pool v0.9.2 h1:VJ9rJDYTFKbp1/wml/7XlLBa8huL5/IeK+1aUf23ugw= +github.com/go-pkgz/pool v0.9.2/go.mod h1:HpVwnbSym5sbYVU/N460+GBixeMgThvbAhMSPsTlkZE= github.com/stretchr/testify v1.12.0 h1:K6Mr6jO9JICuend/5xzTM03ydSV3vdNRYAdPSukj8uI= github.com/stretchr/testify v1.12.0/go.mod h1:bOYBZb5qJ00vPzWfIqBUZPaxK8jWiXc6d3ErP4Ca9Gw= golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= diff --git a/examples/collector_errors/go.sum b/examples/collector_errors/go.sum index df47ecd..f35fbcf 100644 --- a/examples/collector_errors/go.sum +++ b/examples/collector_errors/go.sum @@ -1,3 +1,5 @@ +github.com/go-pkgz/pool v0.9.2 h1:VJ9rJDYTFKbp1/wml/7XlLBa8huL5/IeK+1aUf23ugw= +github.com/go-pkgz/pool v0.9.2/go.mod h1:HpVwnbSym5sbYVU/N460+GBixeMgThvbAhMSPsTlkZE= github.com/stretchr/testify v1.12.0 h1:K6Mr6jO9JICuend/5xzTM03ydSV3vdNRYAdPSukj8uI= github.com/stretchr/testify v1.12.0/go.mod h1:bOYBZb5qJ00vPzWfIqBUZPaxK8jWiXc6d3ErP4Ca9Gw= golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= diff --git a/examples/collectors_chain/go.sum b/examples/collectors_chain/go.sum index df47ecd..f35fbcf 100644 --- a/examples/collectors_chain/go.sum +++ b/examples/collectors_chain/go.sum @@ -1,3 +1,5 @@ +github.com/go-pkgz/pool v0.9.2 h1:VJ9rJDYTFKbp1/wml/7XlLBa8huL5/IeK+1aUf23ugw= +github.com/go-pkgz/pool v0.9.2/go.mod h1:HpVwnbSym5sbYVU/N460+GBixeMgThvbAhMSPsTlkZE= github.com/stretchr/testify v1.12.0 h1:K6Mr6jO9JICuend/5xzTM03ydSV3vdNRYAdPSukj8uI= github.com/stretchr/testify v1.12.0/go.mod h1:bOYBZb5qJ00vPzWfIqBUZPaxK8jWiXc6d3ErP4Ca9Gw= golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= diff --git a/examples/direct_chain/go.sum b/examples/direct_chain/go.sum index df47ecd..f35fbcf 100644 --- a/examples/direct_chain/go.sum +++ b/examples/direct_chain/go.sum @@ -1,3 +1,5 @@ +github.com/go-pkgz/pool v0.9.2 h1:VJ9rJDYTFKbp1/wml/7XlLBa8huL5/IeK+1aUf23ugw= +github.com/go-pkgz/pool v0.9.2/go.mod h1:HpVwnbSym5sbYVU/N460+GBixeMgThvbAhMSPsTlkZE= github.com/stretchr/testify v1.12.0 h1:K6Mr6jO9JICuend/5xzTM03ydSV3vdNRYAdPSukj8uI= github.com/stretchr/testify v1.12.0/go.mod h1:bOYBZb5qJ00vPzWfIqBUZPaxK8jWiXc6d3ErP4Ca9Gw= golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= diff --git a/examples/middleware/go.sum b/examples/middleware/go.sum index 745539a..ef5595a 100644 --- a/examples/middleware/go.sum +++ b/examples/middleware/go.sum @@ -1,3 +1,5 @@ +github.com/go-pkgz/pool v0.9.2 h1:VJ9rJDYTFKbp1/wml/7XlLBa8huL5/IeK+1aUf23ugw= +github.com/go-pkgz/pool v0.9.2/go.mod h1:HpVwnbSym5sbYVU/N460+GBixeMgThvbAhMSPsTlkZE= github.com/stretchr/testify v1.12.0 h1:K6Mr6jO9JICuend/5xzTM03ydSV3vdNRYAdPSukj8uI= github.com/stretchr/testify v1.12.0/go.mod h1:bOYBZb5qJ00vPzWfIqBUZPaxK8jWiXc6d3ErP4Ca9Gw= golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= diff --git a/examples/parallel_files/go.sum b/examples/parallel_files/go.sum index df47ecd..f35fbcf 100644 --- a/examples/parallel_files/go.sum +++ b/examples/parallel_files/go.sum @@ -1,3 +1,5 @@ +github.com/go-pkgz/pool v0.9.2 h1:VJ9rJDYTFKbp1/wml/7XlLBa8huL5/IeK+1aUf23ugw= +github.com/go-pkgz/pool v0.9.2/go.mod h1:HpVwnbSym5sbYVU/N460+GBixeMgThvbAhMSPsTlkZE= github.com/stretchr/testify v1.12.0 h1:K6Mr6jO9JICuend/5xzTM03ydSV3vdNRYAdPSukj8uI= github.com/stretchr/testify v1.12.0/go.mod h1:bOYBZb5qJ00vPzWfIqBUZPaxK8jWiXc6d3ErP4Ca9Gw= golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= diff --git a/examples/pool_completion/go.sum b/examples/pool_completion/go.sum index df47ecd..f35fbcf 100644 --- a/examples/pool_completion/go.sum +++ b/examples/pool_completion/go.sum @@ -1,3 +1,5 @@ +github.com/go-pkgz/pool v0.9.2 h1:VJ9rJDYTFKbp1/wml/7XlLBa8huL5/IeK+1aUf23ugw= +github.com/go-pkgz/pool v0.9.2/go.mod h1:HpVwnbSym5sbYVU/N460+GBixeMgThvbAhMSPsTlkZE= github.com/stretchr/testify v1.12.0 h1:K6Mr6jO9JICuend/5xzTM03ydSV3vdNRYAdPSukj8uI= github.com/stretchr/testify v1.12.0/go.mod h1:bOYBZb5qJ00vPzWfIqBUZPaxK8jWiXc6d3ErP4Ca9Gw= golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= diff --git a/examples/tokenizer_stateful/go.sum b/examples/tokenizer_stateful/go.sum index df47ecd..f35fbcf 100644 --- a/examples/tokenizer_stateful/go.sum +++ b/examples/tokenizer_stateful/go.sum @@ -1,3 +1,5 @@ +github.com/go-pkgz/pool v0.9.2 h1:VJ9rJDYTFKbp1/wml/7XlLBa8huL5/IeK+1aUf23ugw= +github.com/go-pkgz/pool v0.9.2/go.mod h1:HpVwnbSym5sbYVU/N460+GBixeMgThvbAhMSPsTlkZE= github.com/stretchr/testify v1.12.0 h1:K6Mr6jO9JICuend/5xzTM03ydSV3vdNRYAdPSukj8uI= github.com/stretchr/testify v1.12.0/go.mod h1:bOYBZb5qJ00vPzWfIqBUZPaxK8jWiXc6d3ErP4Ca9Gw= golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= diff --git a/examples/tokenizer_stateless/go.sum b/examples/tokenizer_stateless/go.sum index df47ecd..f35fbcf 100644 --- a/examples/tokenizer_stateless/go.sum +++ b/examples/tokenizer_stateless/go.sum @@ -1,3 +1,5 @@ +github.com/go-pkgz/pool v0.9.2 h1:VJ9rJDYTFKbp1/wml/7XlLBa8huL5/IeK+1aUf23ugw= +github.com/go-pkgz/pool v0.9.2/go.mod h1:HpVwnbSym5sbYVU/N460+GBixeMgThvbAhMSPsTlkZE= github.com/stretchr/testify v1.12.0 h1:K6Mr6jO9JICuend/5xzTM03ydSV3vdNRYAdPSukj8uI= github.com/stretchr/testify v1.12.0/go.mod h1:bOYBZb5qJ00vPzWfIqBUZPaxK8jWiXc6d3ErP4Ca9Gw= golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4= From 5cf0e8afe0a6348deeeae50b0134896e1077e631 Mon Sep 17 00:00:00 2001 From: Dmitry Verkhoturov Date: Sun, 11 Oct 2026 11:58:31 +0200 Subject: [PATCH 4/4] take the completion context's cancellation from the caller on every path finishWorker chose between the worker context and a detached one by reading ctx.Err(). a failing worker counts itself out before its error reaches the errgroup, so the last worker could start the callback on a live worker context and the peer's error then cancelled it mid-flight while the caller's context was still live. the callback context now always takes cancellation, deadline and cause from the context passed to Go and values from the worker context. --- README.md | 8 ++--- pool.go | 39 ++++++++++++++++-------- pool_test.go | 86 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 116 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index a634715..fccdc75 100644 --- a/README.md +++ b/README.md @@ -529,10 +529,10 @@ The completion callback executes when: - Skipped only when the context passed to `Go` was cancelled (still runs on `context.DeadlineExceeded`, and on a worker error, which cancels the pool's internal context but not the caller's) -After a worker error the callback receives a context that is no longer cancelled by that error, so it can -still close the next pool in a chain. Cancelling the context passed to `Go` cancels it as well, so a -callback that blocks can always be stopped. An error returned by the callback is reported alongside the -worker error rather than replacing it. +The callback's context takes its cancellation and deadline from the context passed to `Go` and its values, +metrics among them, from the worker context. A worker error doesn't cancel it, so the callback can still +close the next pool in a chain, and cancelling the context passed to `Go` stops a callback that blocks. An +error returned by the callback is reported alongside the worker error and doesn't replace it. Important notes: - Use `Submit` when sending items from a single goroutine diff --git a/pool.go b/pool.go index f23cab6..b2913f0 100644 --- a/pool.go +++ b/pool.go @@ -445,19 +445,8 @@ func (p *WorkerGroup[T]) finishWorker(ctx context.Context, id int, worker Worker // except when the caller cancelled. checked on callerCtx rather than on lastErr, which also // carries the errgroup's cancellation after a peer worker failed if activeWorkers == 0 && p.poolCompleteFn != nil && !errors.Is(p.callerCtx.Err(), context.Canceled) { - completeCtx := ctx - if ctx.Err() != nil && p.callerCtx.Err() == nil { - // ctx is cancelled because a peer worker failed, which should not stop the callback - // from closing the next pool. values, metrics among them, are kept, and the caller's - // own cancellation is bridged back so a blocking callback still ends when asked to - var cancel context.CancelFunc - completeCtx, cancel = context.WithCancel(context.WithoutCancel(ctx)) - defer cancel() - //nolint:contextcheck // callerCtx is deliberate, it is the only context left that the caller still controls - stop := context.AfterFunc(p.callerCtx, cancel) - defer stop() - } - if e := p.poolCompleteFn(completeCtx); e != nil { + //nolint:contextcheck // callerCtx is deliberate, it is the only context the caller alone controls + if e := p.poolCompleteFn(newCompletionCtx(p.callerCtx, ctx)); e != nil { p.completeErr = fmt.Errorf("complete pool func for %d failed: %w", id, e) } } @@ -468,6 +457,30 @@ func (p *WorkerGroup[T]) finishWorker(ctx context.Context, id int, worker Worker return nil } +// completionCtx is the context of the pool completion callback. cancellation, deadline and cause +// come from the context passed to Go and values, metrics among them, from the worker context. the +// worker context is also cancelled by a failing peer worker, which may happen before or after the +// last worker starts the callback, so it can't decide whether the callback keeps running +type completionCtx struct { + context.Context // caller's context: Deadline, Done, Err + values context.Context +} + +// newCompletionCtx makes the completion context from the caller's and the worker's contexts +func newCompletionCtx(callerCtx, workerCtx context.Context) completionCtx { + return completionCtx{Context: callerCtx, values: context.WithoutCancel(workerCtx)} +} + +// Value returns the worker context's value for key, falling back to the caller's context. the +// worker context is derived from the caller's, so the fallback only finds what WithoutCancel hides, +// the caller's cancellation, which context.Cause and derived contexts look up through Value +func (c completionCtx) Value(key any) any { + if v := c.values.Value(key); v != nil { + return v + } + return c.Context.Value(key) +} + // waitWorkers waits for every worker and adds the pool completion error, which the errgroup // cannot carry because it keeps the first error only. func (p *WorkerGroup[T]) waitWorkers() error { diff --git a/pool_test.go b/pool_test.go index 97485e3..01fbdba 100644 --- a/pool_test.go +++ b/pool_test.go @@ -1642,6 +1642,82 @@ func TestPool_PoolCompletion(t *testing.T) { require.ErrorIs(t, err, errFailed, "the worker error must still be reported") require.ErrorIs(t, err, errComplete, "the completion error must not be dropped") }) + + t.Run("peer error landing during completion does not cancel it", func(t *testing.T) { + // a failing worker counts itself out before its error reaches the errgroup, so the last + // worker can start completion on a live pool context that the peer's error then cancels. + // drives finishWorker directly, as nothing outside it can hold the peer in that window + type ctxKey struct{} + errPeer, errCaller := errors.New("peer failed"), errors.New("caller stopped") + started, check, checked := make(chan struct{}), make(chan struct{}), make(chan struct{}) + var afterPeerErr, afterCallerCause error + var value any + + p := New[string](1, WorkerFunc[string](func(context.Context, string) error { return nil })). + WithPoolCompleteFn(func(ctx context.Context) error { + close(started) + <-check + afterPeerErr, value = ctx.Err(), ctx.Value(ctxKey{}) + close(checked) + <-ctx.Done() + afterCallerCause = context.Cause(ctx) + return nil + }) + callerCtx, callerCancel := context.WithCancelCause(context.Background()) + defer callerCancel(nil) + p.callerCtx = callerCtx + p.activeWorkers.Store(1) + workerCtx, peerFailed := context.WithCancelCause(context.WithValue(callerCtx, ctxKey{}, "worker")) + defer peerFailed(nil) + + done := make(chan error, 1) + go func() { done <- p.finishWorker(workerCtx, 0, p.worker, nil, 0) }() + waitClosed(t, started, "pool completion never started") + peerFailed(errPeer) + close(check) + waitClosed(t, checked, "pool completion never checked its context") + callerCancel(errCaller) + + require.NoError(t, <-done) + assert.NoError(t, afterPeerErr, "a peer error must not cancel a running completion callback") + assert.Equal(t, "worker", value, "the callback must keep the worker context's values") + assert.ErrorIs(t, afterCallerCause, errCaller, "the callback context must carry the caller's cause") + }) + + t.Run("completion keeps the caller's deadline", func(t *testing.T) { + for _, fail := range []bool{false, true} { + t.Run(fmt.Sprintf("worker error %v", fail), func(t *testing.T) { + errFailed := errors.New("failed") + var gotDeadline, hasDeadline atomic.Bool + callerCtx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + want, _ := callerCtx.Deadline() + + p := New[string](2, WorkerFunc[string](func(context.Context, string) error { + if fail { + return errFailed + } + return nil + })).WithBatchSize(0).WithPoolCompleteFn(func(ctx context.Context) error { + dl, ok := ctx.Deadline() + hasDeadline.Store(ok) + gotDeadline.Store(dl.Equal(want)) + return nil + }) + require.NoError(t, p.Go(callerCtx)) + p.Submit("item") + + err := p.Close(context.Background()) + if fail { + require.ErrorIs(t, err, errFailed) + } else { + require.NoError(t, err) + } + assert.True(t, hasDeadline.Load(), "the callback must see the caller's deadline") + assert.True(t, gotDeadline.Load(), "the deadline must be the caller's") + }) + } + }) } func TestPool_ChainedBatching(t *testing.T) { @@ -1917,3 +1993,13 @@ func TestPool_WorkerCompleteFnErrorWithExistingError(t *testing.T) { // worker complete function should have been called assert.True(t, completeCalled.Load(), "worker complete function should be called") } + +// waitClosed fails the test if ch is not closed within a second +func waitClosed(t *testing.T, ch <-chan struct{}, msg string) { + t.Helper() + select { + case <-ch: + case <-time.After(time.Second): + t.Fatal(msg) + } +}