diff --git a/README.md b/README.md index fd8e912..fccdc75 100644 --- a/README.md +++ b/README.md @@ -526,7 +526,13 @@ 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) + +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/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= diff --git a/pool.go b/pool.go index 8aea9c1..b2913f0 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 } @@ -436,13 +441,13 @@ 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 { - if lastErr == nil { - lastErr = fmt.Errorf("complete pool func for %d failed: %w", id, e) - } + // 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) { + //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) } } @@ -452,6 +457,37 @@ 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 { + 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. @@ -476,7 +512,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 { @@ -531,7 +567,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 994b152..01fbdba 100644 --- a/pool_test.go +++ b/pool_test.go @@ -1547,6 +1547,177 @@ 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") + }) + + 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") + }) + + 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) { @@ -1822,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) + } +}