pool: let WithCancelOnError also cancel on panics

Author: bobheadxi

pool/context_pool.go

@@ -2,6 +2,8 @@ package pool
 
 import (
 	"context"
+
+	"github.com/sourcegraph/conc/panics"
 )
 
 // ContextPool is a pool that runs tasks that take a context.
@@ -23,14 +25,31 @@ 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 {
-		err := f(p.ctx)
-		if err != nil && p.cancelOnError {
+		// If we aren't cancelling on error, just return the result of f.
+		if !p.cancelOnError {
+			return f(p.ctx)
+		}
+
+		// If we are cancelling on error, then we also want to cancel on panic.
+		// To do this, we need to recover from any panic f raises.
+		var err error
+		recovered := panics.Try(func() { err = f(p.ctx) })
+		if err != nil || recovered != nil {
 			// Leaky abstraction warning: We add the error directly because
 			// otherwise, canceling could cause another goroutine to exit and
 			// return an error before this error was added, which breaks the
 			// expectations of WithFirstError().
-			p.errorPool.addErr(err)
+			if err != nil {
+				p.errorPool.addErr(err)
+			}
+
 			p.cancel()
+
+			// Now that context is cancelled, if we caught a panic we can
+			// propagate it.
+			if recovered != nil {
+				panic(recovered)
+			}
 			return nil
 		}
 		return err
@@ -56,7 +75,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

@@ -8,6 +8,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/sourcegraph/sourcegraph/lib/errors"
@@ -188,6 +189,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} {