Stage interface redesign: pipe type negotiation, pooled buffer copies

README.md

@@ -1,9 +1,9 @@
-# go-pipe [![GoDoc](https://pkg.go.dev/badge/github.com/github/docs)](https://pkg.go.dev/github.com/github/go-pipe)
+# go-pipe [![GoDoc](https://pkg.go.dev/badge/github.com/github/docs)](https://pkg.go.dev/github.com/github/go-pipe/v2)
 A package used to easily build command pipelines in your Go applications
 
 # Important
 We have not thoroughly tested this package on OSs other than Linux, especially Windows. At this time, using this package on Windows based systems is considered experimental and will be supported only on a best effort basis.
 
 # Links
 
-* [Docs](https://pkg.go.dev/github.com/github/go-pipe)
+* [Docs](https://pkg.go.dev/github.com/github/go-pipe/v2)

go.mod

@@ -1,4 +1,4 @@
-module github.com/github/go-pipe
+module github.com/github/go-pipe/v2
 
 go 1.24.0
 

internal/ptree/ptree_test.go

@@ -9,7 +9,7 @@ import (
 	"strconv"
 	"testing"
 
-	"github.com/github/go-pipe/internal/ptree"
+	"github.com/github/go-pipe/v2/internal/ptree"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )

pipe/command.go

@@ -20,9 +20,13 @@ var errProcessInfoMissing = errors.New("cmd.Process is nil")
 // commandStage is a pipeline `Stage` based on running an external
 // command and piping the data through its stdin and stdout.
 type commandStage struct {
-	name   string
-	stdin  io.Closer
-	cmd    *exec.Cmd
+	name string
+	cmd  *exec.Cmd
+
+	// lateClosers is a list of things that have to be closed once the
+	// command has finished.
+	lateClosers []io.Closer
+
 	done   chan struct{}
 	wg     errgroup.Group
 	stderr bytes.Buffer
@@ -32,6 +36,10 @@ type commandStage struct {
 	ctxErr atomic.Value
 }
 
+var (
+	_ Stage = (*commandStage)(nil)
+)
+
 // Command returns a pipeline `Stage` based on the specified external
 // `command`, run with the given command-line `args`. Its stdin and
 // stdout are handled as usual, and its stderr is collected and
@@ -61,33 +69,100 @@ func (s *commandStage) Name() string {
 	return s.name
 }
 
+func (s *commandStage) Preferences() StagePreferences {
+	return StagePreferences{
+		StdinPreference:  IOPreferenceFile,
+		StdoutPreference: IOPreferenceFile,
+	}
+}
+
 func (s *commandStage) Start(
-	ctx context.Context, env Env, stdin io.ReadCloser,
-) (io.ReadCloser, error) {
+	ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser, _ StartOptions,
+) error {
 	if s.cmd.Dir == "" {
 		s.cmd.Dir = env.Dir
 	}
 
 	s.setupEnv(ctx, env)
 
+	// Things that have to be closed as soon as the command has
+	// started:
+	var earlyClosers []io.Closer
+
+	// 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 {
-		// See the long comment in `Pipeline.Start()` for the
-		// explanation of this special case.
 		switch stdin := stdin.(type) {
-		case nopCloser:
-			s.cmd.Stdin = stdin.Reader
-		case nopCloserWriterTo:
-			s.cmd.Stdin = stdin.Reader
+		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)
 		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)
 		}
-		// Also keep a copy so that we can close it when the command exits:
-		s.stdin = stdin
 	}
 
-	stdout, err := s.cmd.StdoutPipe()
-	if err != nil {
-		return nil, err
+	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)
+			}
+		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
+			// pooled buffer rather than letting exec.Cmd allocate a
+			// fresh 32KB buffer for its internal io.Copy.
+			ec, err := s.setupPooledStdout(stdout)
+			if err != nil {
+				return err
+			}
+			earlyClosers = append(earlyClosers, ec)
+			s.lateClosers = append(s.lateClosers, stdout)
+		}
+	}
+
+	closeEarlyClosers := func() {
+		for _, closer := range earlyClosers {
+			_ = closer.Close()
+		}
+	}
+
+	// On error, Close any pipes we created and wait for the goroutines to
+	// exit before propagating the error.
+	cleanupOnStartFailure := func() {
+		closeEarlyClosers()
+		_ = s.wg.Wait()
+		_ = s.closeLateClosers()
 	}
 
 	// If the caller hasn't arranged otherwise, read the command's
@@ -99,7 +174,8 @@ func (s *commandStage) Start(
 		// can be sure.
 		p, err := s.cmd.StderrPipe()
 		if err != nil {
-			return nil, err
+			cleanupOnStartFailure()
+			return err
 		}
 		s.wg.Go(func() error {
 			_, err := io.Copy(&s.stderr, p)
@@ -116,9 +192,12 @@ func (s *commandStage) Start(
 	s.runInOwnProcessGroup()
 
 	if err := s.cmd.Start(); err != nil {
-		return nil, err
+		cleanupOnStartFailure()
+		return err
 	}
 
+	closeEarlyClosers()
+
 	// Arrange for the process to be killed (gently) if the context
 	// expires before the command exits normally:
 	go func() {
@@ -130,7 +209,7 @@ func (s *commandStage) Start(
 		}
 	}()
 
-	return stdout, nil
+	return nil
 }
 
 // setupEnv sets or modifies the environment that will be passed to
@@ -219,21 +298,55 @@ func (s *commandStage) Wait() error {
 
 	// Make sure that any stderr is copied before `s.cmd.Wait()`
 	// closes the read end of the pipe:
-	wErr := s.wg.Wait()
+	wgErr := s.wg.Wait()
 
 	err := s.cmd.Wait()
 	err = s.filterCmdError(err)
 
-	if err == nil && wErr != nil {
-		err = wErr
+	if err == nil && wgErr != nil {
+		err = wgErr
 	}
 
-	if s.stdin != nil {
-		cErr := s.stdin.Close()
-		if cErr != nil && err == nil {
-			return cErr
-		}
+	if closeErr := s.closeLateClosers(); err == nil {
+		err = closeErr
 	}
 
 	return err
 }
+
+func (s *commandStage) closeLateClosers() error {
+	var err error
+	for _, closer := range s.lateClosers {
+		if closeErr := closer.Close(); closeErr != nil && err == nil {
+			err = closeErr
+		}
+	}
+	s.lateClosers = nil
+	return err
+}
+
+// setupPooledStdout creates an `os.Pipe()`, sets it as `cmd.Stdout`,
+// and starts a goroutine that copies from the read end to `dst` using
+// a pooled buffer (or `dst.ReadFrom` when `dst` implements it). The
+// returned closer is the write end of the pipe; the caller must add
+// it to `earlyClosers` so it is closed once the command has started.
+//
+// The buffer-pool optimization works for command stages whose stdout is
+// not an `*os.File`. Without it, `exec.Cmd` would set up its own pipe
+// and run `io.Copy` with a freshly allocated 32KB buffer per invocation.
+func (s *commandStage) setupPooledStdout(dst io.Writer) (io.Closer, error) {
+	pr, pw, err := os.Pipe()
+	if err != nil {
+		return nil, err
+	}
+	s.cmd.Stdout = pw
+	s.wg.Go(func() error {
+		defer pr.Close()
+		_, err := pooledCopy(dst, pr)
+		if err != nil && !errors.Is(err, os.ErrClosed) {
+			return err
+		}
+		return nil
+	})
+	return pw, nil
+}

pipe/command_linux.go

@@ -5,7 +5,7 @@ package pipe
 import (
 	"context"
 
-	"github.com/github/go-pipe/internal/ptree"
+	"github.com/github/go-pipe/v2/internal/ptree"
 )
 
 // On linux, we can limit or observe memory usage in command stages.

pipe/command_nil_panic_test.go

@@ -33,7 +33,7 @@ func TestKillWithFailedStart(t *testing.T) {
 
 	stage := Command("/this/path/does/not/exist/invalid_command_12345")
 
-	_, err := stage.Start(ctx, Env{}, nil)
+	err := stage.Start(ctx, Env{}, nil, nil, StartOptions{})
 	if err == nil {
 		t.Fatal("Expected start to fail, but it succeeded")
 	}

pipe/command_starterror_test.go

@@ -0,0 +1,63 @@
+package pipe_test
+
+import (
+	"bytes"
+	"context"
+	"os/exec"
+	"sync/atomic"
+	"testing"
+
+	"github.com/github/go-pipe/v2/pipe"
+)
+
+// TestCommandStageStartFailureNoRace verifies that when `cmd.Start()`
+// fails (e.g. command not found), the goroutine that
+// `setupPooledStdout` spawned does not leak past `Pipeline.Run()`.
+// `bytes.Buffer.ReadFrom` writes to the buffer's slice header via
+// `grow()` before its first `Read()`, so a leaked goroutine races
+// with the caller's access to the destination buffer once Run
+// returns the error. Run a tight loop so `-race` is likely to catch
+// any regression.
+func TestCommandStageStartFailureNoRace(t *testing.T) {
+	for i := 0; i < 50; i++ {
+		var buf bytes.Buffer
+		p := pipe.New(pipe.WithStdout(&buf))
+		p.Add(pipe.CommandStage("nope", exec.Command("this-binary-does-not-exist-xyz123")))
+		if err := p.Run(context.Background()); err == nil {
+			t.Fatalf("expected error from non-existent command, got nil")
+		}
+		_ = buf.String()
+	}
+}
+
+// trackingWriteCloser is a non-`*os.File` `io.WriteCloser` that records
+// whether it has been closed. Because it isn't an `*os.File`, a command
+// stage routes it through `setupPooledStdout` and closes it as a "late
+// closer" (i.e. only after the command finishes / cleanup runs).
+type trackingWriteCloser struct {
+	closed atomic.Bool
+}
+
+func (w *trackingWriteCloser) Write(p []byte) (int, error) { return len(p), nil }
+
+func (w *trackingWriteCloser) Close() error {
+	w.closed.Store(true)
+	return nil
+}
+
+// TestCommandStageStartFailureClosesLateClosers verifies that a
+// `WithStdoutCloser` on the last stage is closed even when `cmd.Start()`
+// fails. The closer is registered as a "late closer," which is normally
+// drained by `Wait()`; since `Wait()` never runs when `Start()` fails,
+// the start-failure cleanup path must close it instead.
+func TestCommandStageStartFailureClosesLateClosers(t *testing.T) {
+	w := &trackingWriteCloser{}
+	p := pipe.New(pipe.WithStdoutCloser(w))
+	p.Add(pipe.CommandStage("nope", exec.Command("this-binary-does-not-exist-xyz123")))
+	if err := p.Run(context.Background()); err == nil {
+		t.Fatalf("expected error from non-existent command, got nil")
+	}
+	if !w.closed.Load() {
+		t.Fatalf("expected late closer to be closed after Start() failure")
+	}
+}