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

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,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 {

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) {