From 5f54c75cf817ee5059f37453bc7cfc571931d398 Mon Sep 17 00:00:00 2001 From: Dmitry Verkhoturov Date: Wed, 19 Aug 2026 06:15:18 +0100 Subject: [PATCH 1/3] 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/3] 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/3] 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=