Align start-failure cleanup with stream ownership

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

Author: znull

pipe/close_responsibility_test.go

@@ -34,9 +34,9 @@ func (w *writeCloseSpy) Close() error {
 	return nil
 }
 
-// TestGoStageHonorsCloseFlags verifies that a Function stage closes
-// stdin/stdout iff the corresponding close flag is true.
-func TestGoStageHonorsCloseFlags(t *testing.T) {
+// TestGoStageHonorsStreamOwnership verifies that a Function stage closes
+// stdin/stdout iff the corresponding stream is closing.
+func TestGoStageHonorsStreamOwnership(t *testing.T) {
 	cases := []struct {
 		name              string
 		leaveIn, leaveOut bool
@@ -63,8 +63,8 @@ func TestGoStageHonorsCloseFlags(t *testing.T) {
 			))
 			require.NoError(t, s.Wait())
 
-			assert.Equal(t, !tc.leaveIn, in.closed.Load(), "closeStdin=%v", !tc.leaveIn)
-			assert.Equal(t, !tc.leaveOut, out.closed.Load(), "closeStdout=%v", !tc.leaveOut)
+			assert.Equal(t, !tc.leaveIn, in.closed.Load(), "closing stdin=%v", !tc.leaveIn)
+			assert.Equal(t, !tc.leaveOut, out.closed.Load(), "closing stdout=%v", !tc.leaveOut)
 		})
 	}
 }
@@ -88,8 +88,8 @@ func TestStreamConstructorsPreserveOwnershipAndDynamicType(t *testing.T) {
 }
 
 // TestCommandStageHonorsCloseStdin verifies that a command stage closes a
-// non-file stdin (a "late" closer) iff closeStdin is true. An empty
-// reader is used so exec.Cmd's input-copy goroutine sees EOF promptly.
+// non-file stdin (a "late" closer) iff the input stream is closing. An
+// empty reader is used so exec.Cmd's input-copy goroutine sees EOF promptly.
 func TestCommandStageHonorsCloseStdin(t *testing.T) {
 	for _, leave := range []bool{false, true} {
 		name := "owns stdin"
@@ -109,14 +109,14 @@ func TestCommandStageHonorsCloseStdin(t *testing.T) {
 			))
 			require.NoError(t, s.Wait())
 
-			assert.Equal(t, !leave, in.closed.Load(), "closeStdin=%v", !leave)
+			assert.Equal(t, !leave, in.closed.Load(), "closing stdin=%v", !leave)
 		})
 	}
 }
 
 // TestCommandStageHonorsCloseStdout verifies the stdout counterpart: a
 // non-file stdout (routed through the pooled-copy path) is closed iff
-// closeStdout is true.
+// the output stream is closing.
 func TestCommandStageHonorsCloseStdout(t *testing.T) {
 	for _, leave := range []bool{false, true} {
 		name := "owns stdout"
@@ -136,7 +136,7 @@ func TestCommandStageHonorsCloseStdout(t *testing.T) {
 			))
 			require.NoError(t, s.Wait())
 
-			assert.Equal(t, !leave, out.closed.Load(), "closeStdout=%v", !leave)
+			assert.Equal(t, !leave, out.closed.Load(), "closing stdout=%v", !leave)
 		})
 	}
 }

pipe/command_stdout_fastpath_test.go

@@ -18,15 +18,15 @@ import (
 // subprocess can detect when that fd is closed.
 func TestCommandStageStdoutFastPath(t *testing.T) {
 	cases := []struct {
-		name        string
-		closeStdout bool
+		name          string
+		closingStdout bool
 	}{
 		{
-			name:        "raw *os.File with closeStdout",
-			closeStdout: true,
+			name:          "raw *os.File with closing stdout",
+			closingStdout: true,
 		},
 		{
-			name: "raw *os.File without closeStdout",
+			name: "raw *os.File with non-closing stdout",
 		},
 	}
 	for _, tc := range cases {
@@ -43,7 +43,7 @@ func TestCommandStageStdoutFastPath(t *testing.T) {
 			s := CommandStage("true", cmd).(*commandStage)
 
 			stdout := OutputStream{writer: f}
-			if tc.closeStdout {
+			if tc.closingStdout {
 				stdout = ClosingOutput(f)
 			}
 

pipe/pipeline.go

@@ -347,16 +347,19 @@ func (p *Pipeline) Start(ctx context.Context) error {
 		}
 	}
 
-	// Clean up any processes and pipes that have been created. `i` is
-	// the index of the stage that failed to start (whose output pipe
-	// has already been cleaned up if necessary).
+	// Clean up any processes and pipes that have been created. `i` is the
+	// index of the stage that failed to start. If the stage already received
+	// its streams, it owns any closing stream.
 	abort := func(i int, err error, closeFailedStageStdin bool) 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 the failing stage never received its stdin, 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 closeFailedStageStdin {
 			stageStarters[i].stdin.Close()
 		}
+
+		// If stdout was supplied with WithStdoutCloser but the final stage
+		// was never started, then the pipeline still owns that closer.
 		if i < len(p.stages)-1 && p.stdoutCloser != nil {
 			_ = p.stdoutCloser.Close()
 		}

pipe/pipeline_test.go

@@ -83,13 +83,15 @@ func TestPipelineFirstStageFailsToStart(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
 	startErr := errors.New("foo")
+	stdout := &closeTrackingWriter{}
 
-	p := pipe.New()
+	p := pipe.New(pipe.WithStdoutCloser(stdout))
 	p.Add(
 		ErrorStartingStage{startErr},
 		ErrorStartingStage{errors.New("this error should never happen")},
 	)
 	assert.ErrorIs(t, p.Run(ctx), startErr)
+	assert.True(t, stdout.closed, "WithStdoutCloser destination should be closed")
 }
 
 func TestPipelineFirstStageFailsToStartClosesStdoutCloser(t *testing.T) {
@@ -120,6 +122,22 @@ func TestPipelineSecondStageFailsToStart(t *testing.T) {
 	assert.ErrorIs(t, p.Run(ctx), startErr)
 }
 
+func TestPipelineMiddleStageFailsToStartClosesUnstartedStdoutCloser(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	startErr := errors.New("foo")
+	stdout := &closeTrackingWriter{}
+
+	p := pipe.New(pipe.WithStdoutCloser(stdout))
+	p.Add(
+		seqFunction(20000),
+		ErrorStartingStage{startErr},
+		ErrorStartingStage{errors.New("this error should never happen")},
+	)
+	assert.ErrorIs(t, p.Run(ctx), startErr)
+	assert.True(t, stdout.closed, "WithStdoutCloser destination should be closed")
+}
+
 func TestPipelineSingleCommandOutput(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()

pipe/stage.go

@@ -50,18 +50,24 @@ import (
 //	f.Close() // close our copy
 //	cmd.Wait()
 //
-// If the stage is an external command and one of its arguments is not
-// an `*os.File`, then `exec.Cmd` will take care of creating an
-// `os.Pipe()`, copying from the provided argument in/out of the pipe,
-// and eventually closing both ends of the pipe. The stage must close
-// the argument itself, but only _after_ the external command has
+// If the stage is an external command and its stdin is not an
+// `*os.File`, then `exec.Cmd` will take care of creating an
+// `os.Pipe()`, copying from the provided reader into the pipe, and
+// eventually closing both ends of the pipe. The stage must close the
+// provided stdin itself, but only _after_ the external command has
 // finished, like so:
 //
-//	cmd.Stdin = r // Similarly for stdout
+//	cmd.Stdin = r
 //	cmd.Start(…)
 //	cmd.Wait()
 //	r.Close()
 //
+// If the stage is an external command and its stdout is not an
+// `*os.File`, the stage creates a pipe, passes the write end to the
+// command, and copies from the read end to the provided writer. The
+// stage must close the provided stdout itself, but only _after_ the
+// external command and the copy have finished.
+//
 // If the stage is a Go function, then it holds the only copy of
 // stdin/stdout, so it must wait until the function is done before
 // closing them (regardless of their underlying type, like so: