execution: consistently recover in goroutines (#703)

Signed-off-by: Michael Hoffmann <mhoffmann@cloudflare.com>

Author: MichaHoffmann

engine/engine_test.go

@@ -5623,7 +5623,7 @@ func TestEngineRecoversFromPanic(t *testing.T) {
 		testutil.Ok(t, err)
 
 		r := q.Exec(ctx)
-		testutil.Assert(t, r.Err.Error() == "unexpected error: panic!")
+		testutil.Assert(t, r.Err.Error() == "unexpected panic: panic!")
 	})
 
 	t.Run("range", func(t *testing.T) {
@@ -5633,7 +5633,7 @@ func TestEngineRecoversFromPanic(t *testing.T) {
 		testutil.Ok(t, err)
 
 		r := q.Exec(ctx)
-		testutil.Assert(t, r.Err.Error() == "unexpected error: panic!")
+		testutil.Assert(t, r.Err.Error() == "unexpected panic: panic!")
 	})
 }
 

execution/binary/scalar.go

@@ -14,6 +14,7 @@ import (
 	"github.com/thanos-io/promql-engine/query"
 	"github.com/thanos-io/promql-engine/warnings"
 
+	"github.com/efficientgo/core/errors"
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/promql/parser"
@@ -91,12 +92,17 @@ func (o *scalarOperator) Next(ctx context.Context, buf []model.StepVector) (int,
 	var lhsN int
 	var lerrChan = make(chan error, 1)
 	go func() {
+		defer func() {
+			if r := recover(); r != nil {
+				lerrChan <- errors.Newf("unexpected panic: %v", r)
+			}
+			close(lerrChan)
+		}()
 		var err error
 		lhsN, err = o.lhs.Next(ctx, o.lhsBuf)
 		if err != nil {
 			lerrChan <- err
 		}
-		close(lerrChan)
 	}()
 
 	rhsN, rerr := o.rhs.Next(ctx, o.rhsBuf)

execution/binary/vector.go

@@ -106,12 +106,17 @@ func (o *vectorOperator) Next(ctx context.Context, buf []model.StepVector) (int,
 	var lhsN int
 	var lerrChan = make(chan error, 1)
 	go func() {
+		defer func() {
+			if r := recover(); r != nil {
+				lerrChan <- errors.Newf("unexpected panic: %v", r)
+			}
+			close(lerrChan)
+		}()
 		var err error
 		lhsN, err = o.lhs.Next(ctx, o.lhsBuf)
 		if err != nil {
 			lerrChan <- err
 		}
-		close(lerrChan)
 	}()
 
 	rhsN, rerr := o.rhs.Next(ctx, o.rhsBuf)
@@ -153,12 +158,17 @@ func (o *vectorOperator) init(ctx context.Context) error {
 	var highCardSide []labels.Labels
 	var errChan = make(chan error, 1)
 	go func() {
+		defer func() {
+			if r := recover(); r != nil {
+				errChan <- errors.Newf("unexpected panic: %v", r)
+			}
+			close(errChan)
+		}()
 		var err error
 		highCardSide, err = o.lhs.Series(ctx)
 		if err != nil {
 			errChan <- err
 		}
-		close(errChan)
 	}()
 
 	lowCardSide, err := o.rhs.Series(ctx)

execution/exchange/coalesce.go

@@ -124,6 +124,11 @@ func (c *coalesce) Next(ctx context.Context, buf []model.StepVector) (int, error
 		c.wg.Add(1)
 		go func(opIdx int, o model.VectorOperator) {
 			defer c.wg.Done()
+			defer func() {
+				if r := recover(); r != nil {
+					errChan <- errors.Newf("unexpected panic: %v", r)
+				}
+			}()
 
 			n, err := o.Next(ctx, c.tempBufs[opIdx])
 			if err != nil {
@@ -221,16 +226,9 @@ func (c *coalesce) loadSeries(ctx context.Context) error {
 		go func(i int) {
 			defer wg.Done()
 			defer func() {
-				e := recover()
-				if e == nil {
-					return
+				if r := recover(); r != nil {
+					errChan <- errors.Newf("unexpected panic: %v", r)
 				}
-
-				switch err := e.(type) {
-				case error:
-					errChan <- errors.Wrapf(err, "unexpected error")
-				}
-
 			}()
 			series, err := c.operators[i].Series(ctx)
 			if err != nil {

execution/exchange/concurrent.go

@@ -12,6 +12,7 @@ import (
 	"github.com/thanos-io/promql-engine/execution/telemetry"
 	"github.com/thanos-io/promql-engine/query"
 
+	"github.com/efficientgo/core/errors"
 	"github.com/prometheus/prometheus/model/labels"
 )
 
@@ -118,7 +119,12 @@ func (c *concurrencyOperator) Next(ctx context.Context, buf []model.StepVector)
 }
 
 func (c *concurrencyOperator) pull(ctx context.Context) {
-	defer close(c.buffer)
+	defer func() {
+		if r := recover(); r != nil {
+			c.buffer <- maybeStepVector{err: errors.Newf("unexpected panic: %v", r)}
+		}
+		close(c.buffer)
+	}()
 
 	for {
 		select {