Improves panic-handling

- StagePanicHandler siganture is updated to take in `any` to avoid
discarding information. The return
- Panic is wrapped after calling StagePanicHandler and error is returned
for the stage so that we can verify the pipeline's error contains the
stage panic.
- Skip recovery if panic-handler is not set to ensure previous behavior
remains unchanged.
- Cleanup resources on any type of exit from stage.Start to ensure.
- Move recovery logic to a separate method.

Author: jcmanzo

pipe/function.go

@@ -52,34 +52,21 @@ func (s *goStage) Start(ctx context.Context, env Env, stdin io.ReadCloser) (io.R
 
 	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()
+			// Cleanup resources on exit
+			if err := w.Close(); err != nil && s.err == nil {
+				s.err = fmt.Errorf("error closing output pipe for stage %q: %w", s.Name(), err)
+			}
+			if stdin != nil {
+				if err := stdin.Close(); err != nil && s.err == nil {
+					s.err = fmt.Errorf("error closing stdin for stage %q: %w", s.Name(), err)
 				}
-				close(s.done)
-
-				err := FromPanicValue(p)
-				s.err = err
-				s.panicHandler(err)
 			}
+			close(s.done)
 		}()
 
+		defer s.recoverPanic()
+
 		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)
-		}
-		if stdin != nil {
-			if err := stdin.Close(); err != nil && s.err == nil {
-				s.err = fmt.Errorf("error closing stdin for stage %q: %w", s.Name(), err)
-			}
-		}
-		close(s.done)
 	}()
 
 	return r, nil
@@ -89,3 +76,13 @@ func (s *goStage) Wait() error {
 	<-s.done
 	return s.err
 }
+
+func (s *goStage) recoverPanic() {
+	if s.panicHandler == nil {
+		return
+	}
+
+	if p := recover(); p != nil {
+		s.err = s.panicHandler(p)
+	}
+}

pipe/panic.go

@@ -1,7 +1,5 @@
 package pipe
 
-import "fmt"
-
 // 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
@@ -11,17 +9,4 @@ type StagePanicHandlerAware interface {
 }
 
 // StagePanicHandler is a function that handles panics in a pipeline's stages.
-type StagePanicHandler func(err error)
-
-// FromPanicValue converts a panic value to an error. If the panic value is
-// already an error, it returns it directly. Otherwise, it wraps the value in
-// a generic error.
-func FromPanicValue(p any) error {
-	var err error
-	if e, ok := p.(error); ok {
-		err = e
-	} else {
-		err = fmt.Errorf("%v", p)
-	}
-	return err
-}
+type StagePanicHandler func(p any) error

pipe/pipeline_test.go

@@ -462,11 +462,11 @@ func TestFunction(t *testing.T) {
 	})
 
 	t.Run("panic with handler", func(t *testing.T) {
-		panickedMessage := make(chan error, 1)
 		p := pipe.New(
 			pipe.WithDir(dir),
-			pipe.WithStagePanicHandler(func(err error) {
-				panickedMessage <- err
+			pipe.WithStagePanicHandler(func(p any) error {
+				err := fmt.Errorf("panic handled: %v", p)
+				return err
 			}),
 		)
 		p.Add(
@@ -480,8 +480,7 @@ func TestFunction(t *testing.T) {
 		)
 
 		out, err := p.Output(ctx)
-		assert.Error(t, <-panickedMessage)
-		assert.Error(t, err)
+		assert.ErrorContains(t, err, "panic handled")
 		assert.Empty(t, out)
 	})
 }