Merge pull request #92 from sourcegraph/context-pool-panic

pool: let WithCancelOnError also cancel on panics

Author: bobheadxi

pool/context_pool.go

@@ -23,6 +23,18 @@ type ContextPool struct {
 // are busy, a call to Go() will block until the task can be started.
 func (p *ContextPool) Go(f func(ctx context.Context) error) {
 	p.errorPool.Go(func() error {
+		if p.cancelOnError {
+			// If we are cancelling on error, then we also want to cancel if a
+			// panic is raised. To do this, we need to recover, cancel, and then
+			// re-throw the caught panic.
+			defer func() {
+				if r := recover(); r != nil {
+					p.cancel()
+					panic(r)
+				}
+			}()
+		}
+
 		err := f(p.ctx)
 		if err != nil && p.cancelOnError {
 			// Leaky abstraction warning: We add the error directly because
@@ -56,7 +68,7 @@ func (p *ContextPool) WithFirstError() *ContextPool {
 }
 
 // WithCancelOnError configures the pool to cancel its context as soon as
-// any task returns an error. By default, the pool's context is not
+// any task returns an error or panics. By default, the pool's context is not
 // canceled until the parent context is canceled.
 //
 // In this case, all errors returned from the pool after the first will

pool/context_pool_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -187,6 +188,27 @@ func TestContextPool(t *testing.T) {
 		require.NotErrorIs(t, err, context.Canceled)
 	})
 
+	t.Run("WithCancelOnError and panic", func(t *testing.T) {
+		t.Parallel()
+		p := New().WithContext(bgctx).WithCancelOnError()
+		var cancelledTasks atomic.Int64
+		p.Go(func(ctx context.Context) error {
+			<-ctx.Done()
+			cancelledTasks.Add(1)
+			return ctx.Err()
+		})
+		p.Go(func(ctx context.Context) error {
+			<-ctx.Done()
+			cancelledTasks.Add(1)
+			return ctx.Err()
+		})
+		p.Go(func(ctx context.Context) error {
+			panic("abort!")
+		})
+		assert.Panics(t, func() { _ = p.Wait() })
+		assert.EqualValues(t, 2, cancelledTasks.Load())
+	})
+
 	t.Run("limit", func(t *testing.T) {
 		t.Parallel()
 		for _, maxConcurrent := range []int{1, 10, 100} {

pool/result_context_pool_test.go

@@ -10,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -100,6 +101,29 @@ func TestResultContextPool(t *testing.T) {
 		require.ErrorIs(t, err, err1)
 	})
 
+	t.Run("WithCancelOnError and panic", func(t *testing.T) {
+		t.Parallel()
+		p := NewWithResults[int]().
+			WithContext(context.Background()).
+			WithCancelOnError()
+		var cancelledTasks atomic.Int64
+		p.Go(func(ctx context.Context) (int, error) {
+			<-ctx.Done()
+			cancelledTasks.Add(1)
+			return 0, ctx.Err()
+		})
+		p.Go(func(ctx context.Context) (int, error) {
+			<-ctx.Done()
+			cancelledTasks.Add(1)
+			return 0, ctx.Err()
+		})
+		p.Go(func(ctx context.Context) (int, error) {
+			panic("abort!")
+		})
+		assert.Panics(t, func() { _, _ = p.Wait() })
+		assert.EqualValues(t, 2, cancelledTasks.Load())
+	})
+
 	t.Run("no WithCancelOnError", func(t *testing.T) {
 		t.Parallel()
 		g := NewWithResults[int]().WithContext(context.Background())