Extend the Function stage with panic handler support.

Stages can be extended to allow panics to be caught
and reported. The panic will still bubble up once
the panic handler has resolved.

Author: jcmanzo

pipe/function.go

@@ -32,19 +32,42 @@ func Function(name string, f StageFunc) Stage {
 // goStage is a `Stage` that does its work by running an arbitrary
 // `stageFunc` in a goroutine.
 type goStage struct {
-	name string
-	f    StageFunc
-	done chan struct{}
-	err  error
+	name         string
+	f            StageFunc
+	done         chan struct{}
+	err          error
+	panicHandler StagePanicHandler
 }
 
 func (s *goStage) Name() string {
 	return s.name
 }
 
+func (s *goStage) SetPanicHandler(ph StagePanicHandler) {
+	s.panicHandler = ph
+}
+
 func (s *goStage) Start(ctx context.Context, env Env, stdin io.ReadCloser) (io.ReadCloser, error) {
 	r, w := io.Pipe()
+
 	go func() {
+		defer func() {
+			if p := recover(); p != nil {
+				if s.panicHandler == nil {
+					// Nothing to do, just panic
+					panic(p)
+				}
+
+				_ = w.Close()
+				if stdin != nil {
+					_ = stdin.Close()
+				}
+				close(s.done)
+
+				s.err = s.panicHandler(p)
+			}
+		}()
+
 		s.err = s.f(ctx, env, stdin, w)
 		if err := w.Close(); err != nil && s.err == nil {
 			s.err = fmt.Errorf("error closing output pipe for stage %q: %w", s.Name(), err)

pipe/panic.go

@@ -0,0 +1,12 @@
+package pipe
+
+// StagePanicHandlerAware is an interface that Stages can implement to receive
+// a panic handler from the pipeline. This is particularly useful for stages
+// that execute work in a separate goroutine and need to manage panics occurring
+// within that goroutine.
+type StagePanicHandlerAware interface {
+	SetPanicHandler(StagePanicHandler)
+}
+
+// StagePanicHandler is a function that handles panics in the pipeline and its stages.
+type StagePanicHandler func(p any) error