Allow configuring a pipeline with a panic handler.

Clients set the panic handler at the pipeline level.
The pipeline is then responsible for passing the handler
to stages that support the PanicHandlerAware interface.

Author: jcmanzo

pipe/panic.go

@@ -6,4 +6,5 @@ type PanicHandlerAware interface {
 	SetPanicHandler(PanicHandler)
 }
 
+// PanicHandler is a function that handles panics in the pipeline and its stages.
 type PanicHandler func(p any) error

pipe/pipeline.go

@@ -64,6 +64,7 @@ type Pipeline struct {
 	started uint32
 
 	eventHandler func(e *Event)
+	panicHandler PanicHandler
 }
 
 var emptyEventHandler = func(e *Event) {}
@@ -179,6 +180,19 @@ func WithEventHandler(handler func(e *Event)) Option {
 	}
 }
 
+// WithPanicHandler sets a panic handler for the pipeline, allowing clients to handle
+// and observe panics that occur during processing. When a process within the pipeline
+// panics, the provided handler will be invoked, enabling clients to capture the panic,
+// such as for observability purposes.
+//
+// Note: While the handler allows for additional processing of the panic (e.g., logging),
+// the panic will still propagate and be raised after the handler is executed.
+func WithPanicHandler(ph PanicHandler) Option {
+	return func(p *Pipeline) {
+		p.panicHandler = ph
+	}
+}
+
 func (p *Pipeline) hasStarted() bool {
 	return atomic.LoadUint32(&p.started) != 0
 }
@@ -265,6 +279,12 @@ func (p *Pipeline) Start(ctx context.Context) error {
 	}
 
 	for i, s := range p.stages {
+		if p.panicHandler != nil {
+			if phs, ok := s.(PanicHandlerAware); ok {
+				phs.SetPanicHandler(p.panicHandler)
+			}
+		}
+
 		var err error
 		stdout, err := s.Start(ctx, p.env, nextStdin)
 		if err != nil {

pipe/pipeline_test.go

@@ -436,28 +436,54 @@ func TestFunction(t *testing.T) {
 
 	dir := t.TempDir()
 
-	p := pipe.New(pipe.WithDir(dir))
-	p.Add(
-		pipe.Print("hello world"),
-		pipe.Function(
-			"farewell",
-			func(_ context.Context, _ pipe.Env, stdin io.Reader, stdout io.Writer) error {
-				buf, err := io.ReadAll(stdin)
-				if err != nil {
+	t.Run("successful function", func(t *testing.T) {
+		p := pipe.New(pipe.WithDir(dir))
+		p.Add(
+			pipe.Print("hello world"),
+			pipe.Function(
+				"farewell",
+				func(_ context.Context, _ pipe.Env, stdin io.Reader, stdout io.Writer) error {
+					buf, err := io.ReadAll(stdin)
+					if err != nil {
+						return err
+					}
+					if string(buf) != "hello world" {
+						return fmt.Errorf("expected \"hello world\"; got %q", string(buf))
+					}
+					_, err = stdout.Write([]byte("goodbye, cruel world"))
 					return err
-				}
-				if string(buf) != "hello world" {
-					return fmt.Errorf("expected \"hello world\"; got %q", string(buf))
-				}
-				_, err = stdout.Write([]byte("goodbye, cruel world"))
-				return err
-			},
-		),
-	)
+				},
+			),
+		)
 
-	out, err := p.Output(ctx)
-	assert.NoError(t, err)
-	assert.EqualValues(t, "goodbye, cruel world", out)
+		out, err := p.Output(ctx)
+		assert.NoError(t, err)
+		assert.EqualValues(t, "goodbye, cruel world", out)
+	})
+
+	t.Run("panic with handler", func(t *testing.T) {
+		expectedErr := fmt.Errorf("recovered from panic: oh no!")
+		p := pipe.New(
+			pipe.WithDir(dir),
+			pipe.WithPanicHandler(func(panicValue any) error {
+				assert.Equal(t, "oh no!", panicValue)
+				return expectedErr
+			}),
+		)
+		p.Add(
+			pipe.Print("hello world"),
+			pipe.Function(
+				"farewell",
+				func(_ context.Context, _ pipe.Env, stdin io.Reader, stdout io.Writer) error {
+					panic("oh no!")
+				},
+			),
+		)
+
+		out, err := p.Output(ctx)
+		assert.ErrorIs(t, err, expectedErr)
+		assert.Empty(t, out)
+	})
 }
 
 func TestPipelineWithFunction(t *testing.T) {