diff --git a/README.md b/README.md index fd8e912..a634715 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) + +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 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..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 } @@ -436,13 +441,24 @@ 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) { + 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 { + p.completeErr = fmt.Errorf("complete pool func for %d failed: %w", id, e) } } @@ -452,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. @@ -476,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 { @@ -531,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 994b152..97485e3 100644 --- a/pool_test.go +++ b/pool_test.go @@ -1547,6 +1547,101 @@ 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") + }) } func TestPool_ChainedBatching(t *testing.T) {