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/pipeline.go

@@ -64,6 +64,7 @@ type Pipeline struct {
 	started uint32
 
 	eventHandler func(e *Event)
+	panicHandler StagePanicHandler
 }
 
 var emptyEventHandler = func(e *Event) {}
@@ -179,6 +180,19 @@ func WithEventHandler(handler func(e *Event)) Option {
 	}
 }
 
+// WithStagePanicHandler sets a panic handler for the stages within a pipeline.
+// When a stage within the pipeline panics, the provided handler will be invoked, enabling
+// clients to capture the panic, such as for observability purposes.
+//
+// Note:
+//   - The client is responsible for deciding whether to recover from the panic or propagate it further.
+//   - Currently, only the Function stage supports this functionality.
+func WithStagePanicHandler(ph StagePanicHandler) 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.(StagePanicHandlerAware); 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,55 @@ 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) {
+		panicked := make(chan bool, 1)
+		p := pipe.New(
+			pipe.WithDir(dir),
+			pipe.WithStagePanicHandler(func(err error) {
+				panicked <- true
+				assert.Equal(t, "this is a panic", err.Error())
+			}),
+		)
+		p.Add(
+			pipe.Print("hello world"),
+			pipe.Function(
+				"farewell",
+				func(_ context.Context, _ pipe.Env, stdin io.Reader, stdout io.Writer) error {
+					panic("this is a panic")
+				},
+			),
+		)
+
+		out, err := p.Output(ctx)
+		assert.True(t, <-panicked)
+		assert.Error(t, err)
+		assert.Empty(t, out)
+	})
 }
 
 func TestPipelineWithFunction(t *testing.T) {