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,41 @@ 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 PanicHandler
 }
 
 func (s *goStage) Name() string {
 	return s.name
 }
 
+func (s *goStage) SetPanicHandler(ph PanicHandler) {
+	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)
+				}
+
+				s.err = s.panicHandler(p)
+				_ = w.Close()
+				if stdin != nil {
+					_ = stdin.Close()
+				}
+				close(s.done)
+			}
+		}()
+
 		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,9 @@
+package pipe
+
+// PanicHandlerAware is an interface that stages can implement to receive
+// a panic handler from the pipeline.
+type PanicHandlerAware interface {
+	SetPanicHandler(PanicHandler)
+}
+
+type PanicHandler func(p any) error