Decouple close-responsibility from stdin/stdout dynamic type

Rather than inferring who is responsible for closing a pipeline stage's
stdin/stdout from the dynamic type, communicate that explicitly in
StartOptions.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: znull

pipe/close_responsibility_test.go

@@ -0,0 +1,136 @@
+package pipe
+
+import (
+	"context"
+	"io"
+	"os/exec"
+	"strings"
+	"sync/atomic"
+	"testing"
+)
+
+// readCloseSpy records whether Close was called.
+type readCloseSpy struct {
+	io.Reader
+	closed atomic.Bool
+}
+
+func (r *readCloseSpy) Close() error {
+	r.closed.Store(true)
+	return nil
+}
+
+// writeCloseSpy records whether Close was called.
+type writeCloseSpy struct {
+	io.Writer
+	closed atomic.Bool
+}
+
+func (w *writeCloseSpy) Close() error {
+	w.closed.Store(true)
+	return nil
+}
+
+// TestGoStageHonorsLeaveOpenFlags verifies that a Function stage closes
+// stdin/stdout iff the corresponding StartOptions.Leave*Open flag is unset.
+func TestGoStageHonorsLeaveOpenFlags(t *testing.T) {
+	cases := []struct {
+		name              string
+		leaveIn, leaveOut bool
+	}{
+		{"own both", false, false},
+		{"leave stdin open", true, false},
+		{"leave stdout open", false, true},
+		{"leave both open", true, true},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			in := &readCloseSpy{Reader: strings.NewReader("hi")}
+			out := &writeCloseSpy{Writer: io.Discard}
+
+			s := Function("f", func(_ context.Context, _ Env, stdin io.Reader, stdout io.Writer) error {
+				_, err := io.Copy(stdout, stdin)
+				return err
+			})
+
+			if err := s.Start(context.Background(), Env{}, in, out, StartOptions{
+				LeaveStdinOpen:  tc.leaveIn,
+				LeaveStdoutOpen: tc.leaveOut,
+			}); err != nil {
+				t.Fatalf("Start: %v", err)
+			}
+			if err := s.Wait(); err != nil {
+				t.Fatalf("Wait: %v", err)
+			}
+
+			if got, want := in.closed.Load(), !tc.leaveIn; got != want {
+				t.Errorf("stdin closed = %v, want %v (LeaveStdinOpen=%v)", got, want, tc.leaveIn)
+			}
+			if got, want := out.closed.Load(), !tc.leaveOut; got != want {
+				t.Errorf("stdout closed = %v, want %v (LeaveStdoutOpen=%v)", got, want, tc.leaveOut)
+			}
+		})
+	}
+}
+
+// TestCommandStageHonorsLeaveStdinOpen verifies that a command stage closes a
+// non-file stdin (a "late" closer) iff LeaveStdinOpen is unset. An empty
+// reader is used so exec.Cmd's input-copy goroutine sees EOF promptly.
+func TestCommandStageHonorsLeaveStdinOpen(t *testing.T) {
+	for _, leave := range []bool{false, true} {
+		name := "owns stdin"
+		if leave {
+			name = "leaves stdin open"
+		}
+		t.Run(name, func(t *testing.T) {
+			in := &readCloseSpy{Reader: strings.NewReader("")}
+
+			cmd := exec.Command("true")
+			s := CommandStage("true", cmd).(*commandStage)
+
+			if err := s.Start(context.Background(), Env{}, in, nil, StartOptions{
+				LeaveStdinOpen: leave,
+			}); err != nil {
+				t.Fatalf("Start: %v", err)
+			}
+			if err := s.Wait(); err != nil {
+				t.Fatalf("Wait: %v", err)
+			}
+
+			if got, want := in.closed.Load(), !leave; got != want {
+				t.Errorf("stdin closed = %v, want %v (LeaveStdinOpen=%v)", got, want, leave)
+			}
+		})
+	}
+}
+
+// TestCommandStageHonorsLeaveStdoutOpen verifies the stdout counterpart: a
+// non-file stdout (routed through the pooled-copy path) is closed iff
+// LeaveStdoutOpen is unset.
+func TestCommandStageHonorsLeaveStdoutOpen(t *testing.T) {
+	for _, leave := range []bool{false, true} {
+		name := "owns stdout"
+		if leave {
+			name = "leaves stdout open"
+		}
+		t.Run(name, func(t *testing.T) {
+			out := &writeCloseSpy{Writer: io.Discard}
+
+			cmd := exec.Command("true")
+			s := CommandStage("true", cmd).(*commandStage)
+
+			if err := s.Start(context.Background(), Env{}, nil, out, StartOptions{
+				LeaveStdoutOpen: leave,
+			}); err != nil {
+				t.Fatalf("Start: %v", err)
+			}
+			if err := s.Wait(); err != nil {
+				t.Fatalf("Wait: %v", err)
+			}
+
+			if got, want := out.closed.Load(), !leave; got != want {
+				t.Errorf("stdout closed = %v, want %v (LeaveStdoutOpen=%v)", got, want, leave)
+			}
+		})
+	}
+}

pipe/command.go

@@ -77,7 +77,7 @@ func (s *commandStage) Preferences() StagePreferences {
 }
 
 func (s *commandStage) Start(
-	ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser, _ StartOptions,
+	ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser, opts StartOptions,
 ) error {
 	if s.cmd.Dir == "" {
 		s.cmd.Dir = env.Dir
@@ -92,62 +92,44 @@ func (s *commandStage) Start(
 	// See the type comment for `Stage` and the long comment in
 	// `Pipeline.WithStdin()` for the explanation of this unwrapping
 	// and closing behavior.
-
 	if stdin != nil {
-		switch stdin := stdin.(type) {
-		case readerNopCloser:
-			// In this case, we shouldn't close it. But unwrap it for
-			// efficiency's sake:
-			s.cmd.Stdin = UnwrapReader(stdin)
-		case *os.File:
-			// In this case, we can close stdin as soon as the command
-			// has started:
-			s.cmd.Stdin = stdin
-			earlyClosers = append(earlyClosers, stdin)
+		// For a non-wrapped value this is a no-op.
+		reader := UnwrapReader(stdin)
+		s.cmd.Stdin = reader
+
+		switch {
+		case opts.LeaveStdinOpen:
+			// leave it open.
 		default:
-			// In this case, we need to close `stdin`, but we should
-			// only do so after the command has finished:
-			s.cmd.Stdin = stdin
-			s.lateClosers = append(s.lateClosers, stdin)
+			if _, ok := reader.(*os.File); ok {
+				// We can close our copy as soon as the command has started
+				earlyClosers = append(earlyClosers, stdin)
+			} else {
+				// We need to close `stdin`, but only after the command has finished
+				s.lateClosers = append(s.lateClosers, stdin)
+			}
 		}
 	}
 
 	if stdout != nil {
-		// See the long comment in `Pipeline.Start()` for the
-		// explanation of this special case.
-		switch stdout := stdout.(type) {
-		case writerNopCloser:
-			// We shouldn't close the wrapped writer. Unwrap it; if
-			// it's an `*os.File`, exec.Cmd can pass the fd directly
-			// to the child. Otherwise route the copy through our own
-			// pipe so we can use a pooled buffer.
-			writer := UnwrapWriter(stdout)
-			if f, ok := writer.(*os.File); ok {
-				s.cmd.Stdout = f
-			} else {
-				ec, err := s.setupPooledStdout(writer)
-				if err != nil {
-					return err
-				}
-				earlyClosers = append(earlyClosers, ec)
+		writer := UnwrapWriter(stdout)
+		if f, ok := writer.(*os.File); ok {
+			s.cmd.Stdout = f
+			if !opts.LeaveStdoutOpen {
+				earlyClosers = append(earlyClosers, stdout)
 			}
-		case *os.File:
-			// In this case, we can close stdout as soon as the command
-			// has started:
-			s.cmd.Stdout = stdout
-			earlyClosers = append(earlyClosers, stdout)
-		default:
-			// In this case, we need to close `stdout`, but we should
-			// only do so after the command has finished. We also
-			// route the copy through our own pipe so we can use a
+		} else {
+			// Route the copy through our own pipe so we can use a
 			// pooled buffer rather than letting exec.Cmd allocate a
 			// fresh 32KB buffer for its internal io.Copy.
-			ec, err := s.setupPooledStdout(stdout)
+			ec, err := s.setupPooledStdout(writer)
 			if err != nil {
 				return err
 			}
 			earlyClosers = append(earlyClosers, ec)
-			s.lateClosers = append(s.lateClosers, stdout)
+			if !opts.LeaveStdoutOpen {
+				s.lateClosers = append(s.lateClosers, stdout)
+			}
 		}
 	}
 

pipe/function.go

@@ -74,12 +74,12 @@ func (s *goStage) Start(
 					s.err = opts.PanicHandler(p)
 				}
 			}
-			if stdout != nil {
+			if stdout != nil && !opts.LeaveStdoutOpen {
 				if err := stdout.Close(); err != nil && s.err == nil {
 					s.err = fmt.Errorf("error closing stdout for stage %q: %w", s.Name(), err)
 				}
 			}
-			if stdin != nil {
+			if stdin != nil && !opts.LeaveStdinOpen {
 				if err := stdin.Close(); err != nil && s.err == nil {
 					s.err = fmt.Errorf("error closing stdin for stage %q: %w", s.Name(), err)
 				}

pipe/nop_closer.go

@@ -25,10 +25,7 @@ func (readerNopCloser) Close() error {
 	return nil
 }
 
-// writerNopCloser is a WriteCloser that wraps a provided `io.Writer`, but
-// whose `Close()` method does nothing. It should be unwrapped (via
-// [UnwrapWriter]) before use where fast-path interfaces such as
-// `io.ReaderFrom` are relevant.
+// writerNopCloser is the stdout counterpart of [readerNopCloser]
 type writerNopCloser struct {
 	io.Writer
 }

pipe/pipeline.go

@@ -60,6 +60,8 @@ type Pipeline struct {
 	stages []Stage
 	cancel func()
 
+	leaveStdoutOpen bool // only matters when stdout is non-nil
+
 	// Atomically written and read value, nonzero if the pipeline has
 	// been started. This is only used for lifecycle sanity checks but
 	// does not guarantee that clients are using the class correctly.
@@ -152,6 +154,7 @@ func WithStdin(stdin io.Reader) Option {
 func WithStdout(stdout io.Writer) Option {
 	return func(p *Pipeline) {
 		p.stdout = writerNopCloser{stdout}
+		p.leaveStdoutOpen = true
 	}
 }
 
@@ -160,6 +163,7 @@ func WithStdout(stdout io.Writer) Option {
 func WithStdoutCloser(stdout io.WriteCloser) Option {
 	return func(p *Pipeline) {
 		p.stdout = stdout
+		p.leaveStdoutOpen = false
 	}
 }
 
@@ -258,6 +262,19 @@ type stageStarter struct {
 	stdout io.WriteCloser
 }
 
+// startOptions builds the StartOptions for the stage at index i. It sets
+// LeaveStdinOpen/LeaveStdoutOpen for the first and last stages, as appropriate.
+func (p *Pipeline) startOptions(i int) StartOptions {
+	opts := StartOptions{PanicHandler: p.panicHandler}
+	if i == 0 && p.stdin != nil {
+		opts.LeaveStdinOpen = true
+	}
+	if i == len(p.stages)-1 && p.stdout != nil {
+		opts.LeaveStdoutOpen = p.leaveStdoutOpen
+	}
+	return opts
+}
+
 // Start starts the commands in the pipeline. If `Start()` exits
 // without an error, `Wait()` must also be called, to allow all
 // resources to be freed.
@@ -320,7 +337,7 @@ func (p *Pipeline) Start(ctx context.Context) error {
 		// Close the pipe that the previous stage was writing to.
 		// That should cause it to exit even if it's not minding
 		// its context.
-		if stageStarters[i].stdin != nil {
+		if stageStarters[i].stdin != nil && !p.startOptions(i).LeaveStdinOpen {
 			_ = stageStarters[i].stdin.Close()
 		}
 
@@ -361,7 +378,7 @@ func (p *Pipeline) Start(ctx context.Context) error {
 		} else {
 			nextSS.stdin, ss.stdout = io.Pipe()
 		}
-		if err := s.Start(ctx, p.env, ss.stdin, ss.stdout, StartOptions{PanicHandler: p.panicHandler}); err != nil {
+		if err := s.Start(ctx, p.env, ss.stdin, ss.stdout, p.startOptions(i)); err != nil {
 			nextSS.stdin.Close()
 			ss.stdout.Close()
 			return abort(i, err)
@@ -376,7 +393,7 @@ func (p *Pipeline) Start(ctx context.Context) error {
 		s := p.stages[i]
 		ss := &stageStarters[i]
 
-		if err := s.Start(ctx, p.env, ss.stdin, ss.stdout, StartOptions{PanicHandler: p.panicHandler}); err != nil {
+		if err := s.Start(ctx, p.env, ss.stdin, ss.stdout, p.startOptions(i)); err != nil {
 			return abort(i, err)
 		}
 	}
@@ -387,6 +404,7 @@ func (p *Pipeline) Start(ctx context.Context) error {
 func (p *Pipeline) Output(ctx context.Context) ([]byte, error) {
 	var buf bytes.Buffer
 	p.stdout = writerNopCloser{&buf}
+	p.leaveStdoutOpen = true
 	err := p.Run(ctx)
 	return buf.Bytes(), err
 }

pipe/stage.go

@@ -10,13 +10,13 @@ import (
 //
 // Who closes stdin and stdout?
 //
-// A `Stage` as a whole needs to be responsible for closing its end of
+// A `Stage` as a whole is responsible for closing its end of
 // stdin and stdout (assuming that `Start()` returns successfully).
 // Its doing so tells the previous/next stage that it is done
 // reading/writing data, which can affect their behavior. Therefore,
 // it should close each one as soon as it is done with it. If the
 // caller wants to suppress the closing of stdin/stdout, it can always
-// wrap the corresponding argument in a "nopCloser".
+// indicate otherwise via `StartOptions`.
 //
 // How this should be done depends on whether stdin/stdout are of type
 // `*os.File`.
@@ -63,18 +63,17 @@ import (
 //	}()
 //
 // From the point of view of the pipeline as a whole, if stdin is
-// provided by the user (`WithStdin()`), then we don't want to close
-// it at all, whether it's an `*os.File` or not. For this reason,
-// stdin has to be wrapped using a `readerNopCloser` before being
-// passed into the first stage. For efficiency reasons, the first
-// stage should ideally unwrap its stdin argument (using
-// [UnwrapReader]) before actually using it. If the wrapped value is
-// an `*os.File` and the stage is a command stage, then unwrapping is
-// also important to get the right semantics.
+// provided by the user (`WithStdin()`), then we don't want the first
+// stage to close it at all, whether it's an `*os.File` or not. The
+// pipeline communicates this by setting `StartOptions.LeaveStdinOpen`
+// when it starts that stage. stdin is still wrapped in a
+// `readerNopCloser` before being passed in, but only so that a bare
+// `io.Reader` satisfies `io.ReadCloser`, and so that a command stage
+// can recover the underlying object via [UnwrapReader].
 //
 // For stdout, it depends on whether the user supplied it using
-// `WithStdout()` or `WithStdoutCloser()`. If the former, then the
-// considerations are the same as for stdin.
+// `WithStdout()` or `WithStdoutCloser()`. [UnwrapWriter] plays the same
+// role for stdout that [UnwrapReader] plays for stdin.
 //
 // [1] It's theoretically possible for a command to pass the open file
 //     descriptor to another, longer-lived process, in which case the
@@ -118,6 +117,10 @@ type StartOptions struct {
 	// handler), converting it into an error. Stage types that don't run
 	// user code in a library-spawned goroutine ignore it.
 	PanicHandler StagePanicHandler
+
+	// LeaveStd{in,out}Open tell the stage that it must NOT close stdin/stdout
+	LeaveStdinOpen  bool
+	LeaveStdoutOpen bool
 }
 
 // StagePanicHandler is a function that handles panics in a pipeline's stages.