run pool completion after a worker error, not by chance

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

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=

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=

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=

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=

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=

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=

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=

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=

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=

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=

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 {

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)
+	}
+}