Stage: change the interface to make stdin/stdout handling more flexible

The old `Stage` interface, and in particular its `Start()` method, was
not ideal. `Start()` was responsible for creating its own stdout,
without knowledge of what will be consuming it.

In practice, there are only two main stages:

* `commandStage` ultimately runs a subprocess, which needs an
  `*os.File` as both stdin and stdout. The old code created its stdout
  using `cmd.StdoutPipe()`, which creates an `*os.File`.

* `goStage` runs a Go function, which should be happy with any kind of
  `io.ReadCloser` / `io.WriteCloser` for its stdin and stdout. The old
  code created its stdout using `io.Pipe()`, which _doesn't_ return an
  `*os.File`.

There are some scenarios where the old behavior was not ideal:

1. If a `goStage` was followed by a `commandStage`, the `commandStage`
   would had to consume the non-`*os.File` stdin that was created by
   the former. But since an external command requires an `*os.File`,
   `exec.Cmd` had to create an `os.Pipe()` internally and create an
   extra goroutine to copy from the `io.Reader` to the pipe. This is
   not only wasteful, but also meant that the `goStage` was not
   informed when the subprocess terminated or closed its stdin. (For
   example, the copy goroutine could block waiting to read from the
   `io.Reader`.)

2. If `Pipeline.stdout` was set, then an extra stage was always needed
   to copy from the output of the last stage to `Pipeline.stdout`. But:

   * If the last stage was a `commandStage` and `Pipeline.stdout` was
     an `*os.File`, then this copy was unnecessary; the subprocess
     could instead have written directly to the corresponding file
     descriptor. This was wasteful, and also lead to cases where the
     subprocess couldn't detect that `Pipeline.stdout` had been
     closed.

   * If the last stage was a `goStage`, then the copy was also
     unnecessary; the stage could have written directly to
     `Pipeline.stdout` whatever its type.

Problem (1) could have been fixed by changing `goStage` to always use
`os.Pipe()` to create its stdout pipe. But that would be wasteful if
two `goStage`s were adjacent, in which case they could use a cheaper
`io.Pipe()` instead. And it wouldn't solve problem (2) at all.

Both problems can only be solved by considering both the producer
_and_ the consumer of the stdin and stdout of any stage. If either end
is a `commandStage`, then it is preferable to us `os.Pipe()`. If both
ends are `goStage`s, then it is preferable to use `io.Pipe()`. And if
`Pipeline.Stdout` is set, the last stage should write directly into it
whenever possible.

This PR solves the problem by changing the `Stage` interface to add a
`Preferences()` method and change the signature of the `Start()`
method:

    Preferences() StagePreferences
    Start(
        ctx context.Context, env Env,
	stdin io.ReadCloser, stdout io.WriteCloser,
    ) error

The first indicates what kind of stdin/stdout the stage prefers, and
the second starts up the stage with a `stdin` and `stdout` that are
provided by the caller, rather than letting the stage return its own
stdout.

Now, when a stage is added to a `Pipeline`, then `Pipeline.Start()`
uses the first method to figure out what kind of pipes are preferred
between this stage and its neighbors, then the second is called to
start the stage with the preferred type of pipe if possible. It also
passes `Pipeline.stdout` into the last stage rather than copying the
data an extra time.

Note that this is a backwards-incompatible change, and thus will
require a change to v2. Any clients that implement their own `Stage`
will have to change their stage to conform to the new interface.

However, clients that only create stages using the functions in this
package (e.g., `pipe.Command()`, `pipe.CommandStage()`,
`pipe.Function()`, `pipe.LinewiseFunction()`, etc.) should continue to
work without changes, since those functions' signatures haven't
changed. Such clients will get the benefit of the new behavior. For
example, the benchmarks `BenchmarkMoreDataBuffered` and
`BenchmarkMoreDataUnbuffered` (admittedly, worst cases for the old
code) are sped up by roughly 2.25x and 6.6x, respectively:

```
snare:~/github/proj/go-pipe/git(main-bench)$ /bin/time go test -bench=. -benchtime=10s ./pipe/pipeline_test.go
goos: linux
goarch: amd64
cpu: Intel(R) Xeon(R) W-2255 CPU @ 3.70GHz
BenchmarkSingleProgram-20         	    8497	   1383275 ns/op
BenchmarkTenPrograms-20           	    2186	   5388075 ns/op
BenchmarkTenFunctions-20          	   37605	    324808 ns/op
BenchmarkTenMixedStages-20        	    3380	   3565218 ns/op
BenchmarkMoreDataUnbuffered-20    	      25	 423838490 ns/op
BenchmarkMoreDataBuffered-20      	      44	 261734773 ns/op
PASS
ok  	command-line-arguments	76.120s
172.91user 91.15system 1:16.56elapsed 344%CPU (0avgtext+0avgdata 114080maxresident)k
0inputs+7768outputs (40major+3819487minor)pagefaults 0swaps

snare:~/github/proj/go-pipe/git(version-2)$ /bin/time go test -bench=. -benchtime=10s ./pipe/pipeline_test.go
goos: linux
goarch: amd64
cpu: Intel(R) Xeon(R) W-2255 CPU @ 3.70GHz
BenchmarkSingleProgram-20         	    8458	   1366214 ns/op
BenchmarkTenPrograms-20           	    2233	   5296019 ns/op
BenchmarkTenFunctions-20          	   42453	    289761 ns/op
BenchmarkTenMixedStages-20        	    3398	   3497226 ns/op
BenchmarkMoreDataUnbuffered-20    	     177	  64410211 ns/op
BenchmarkMoreDataBuffered-20      	     100	 115728132 ns/op
PASS
ok  	command-line-arguments	82.751s
175.42user 142.81system 1:23.21elapsed 382%CPU (0avgtext+0avgdata 114080maxresident)k
0inputs+7776outputs (42major+3883888minor)pagefaults 0swaps
```

Also, look how much simpler `testMemoryLimit()` has become, since it
doesn't need the awkward workaround that was previously required.

In terms of backwards compatibility, some applications might notice a
difference with the new pipe structure. The difference should usually
be an improvement, for example lower resource consumption and less
risk of deadlock. It is conceivable that some applications were in
some way relying on the delayed completion of pipelines when an
`io.Pipe` was closed, though I'm having trouble imagining scenarios
like that in the real world.

# Conflicts:
#	pipe/function.go
#	pipe/iocopier.go
#	pipe/memorylimit_test.go
#	pipe/pipeline.go

Author: mhagger

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,80 @@ func (s *commandStage) Name() string {
 	return s.name
 }
 
+func (s *commandStage) Preferences() StagePreferences {
+	prefs := StagePreferences{
+		StdinPreference:  IOPreferenceFile,
+		StdoutPreference: IOPreferenceFile,
+	}
+	if s.cmd.Stdin != nil {
+		prefs.StdinPreference = IOPreferenceNil
+	}
+	if s.cmd.Stdout != nil {
+		prefs.StdoutPreference = IOPreferenceNil
+	}
+
+	return prefs
+}
+
 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,
+) 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 command 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 readerNopCloser:
+			// In this case, we shouldn't close it. But unwrap it for
+			// efficiency's sake:
 			s.cmd.Stdin = stdin.Reader
 		case readerWriterToNopCloser:
+			// In this case, we shouldn't close it. But unwrap it for
+			// efficiency's sake:
 			s.cmd.Stdin = stdin.Reader
+		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:
+			// In this case, we shouldn't close it. But unwrap it for
+			// efficiency's sake:
+			s.cmd.Stdout = stdout.Writer
+		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:
+			s.cmd.Stdout = stdout
+			s.lateClosers = append(s.lateClosers, stdout)
+		}
 	}
 
 	// If the caller hasn't arranged otherwise, read the command's
@@ -99,7 +154,7 @@ func (s *commandStage) Start(
 		// can be sure.
 		p, err := s.cmd.StderrPipe()
 		if err != nil {
-			return nil, err
+			return err
 		}
 		s.wg.Go(func() error {
 			_, err := io.Copy(&s.stderr, p)
@@ -116,7 +171,11 @@ func (s *commandStage) Start(
 	s.runInOwnProcessGroup()
 
 	if err := s.cmd.Start(); err != nil {
-		return nil, err
+		return err
+	}
+
+	for _, closer := range earlyClosers {
+		_ = closer.Close()
 	}
 
 	// Arrange for the process to be killed (gently) if the context
@@ -130,7 +189,7 @@ func (s *commandStage) Start(
 		}
 	}()
 
-	return stdout, nil
+	return nil
 }
 
 // setupEnv sets or modifies the environment that will be passed to
@@ -219,19 +278,18 @@ 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
+	for _, closer := range s.lateClosers {
+		if closeErr := closer.Close(); closeErr != nil && err == nil {
+			err = closeErr
 		}
 	}
 

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)
 	if err == nil {
 		t.Fatal("Expected start to fail, but it succeeded")
 	}

pipe/command_test.go

@@ -78,7 +78,8 @@ func TestCopyEnvWithOverride(t *testing.T) {
 	for _, ex := range examples {
 		t.Run(ex.label, func(t *testing.T) {
 			assert.ElementsMatch(t, ex.expectedResult,
-				copyEnvWithOverrides(ex.env, ex.overrides))
+				copyEnvWithOverrides(ex.env, ex.overrides),
+			)
 		})
 	}
 }

pipe/function.go

@@ -9,7 +9,7 @@ import (
 // StageFunc is a function that can be used to power a `goStage`. It
 // should read its input from `stdin` and write its output to
 // `stdout`. `stdin` and `stdout` will be closed automatically (if
-// necessary) once the function returns.
+// non-nil) once the function returns.
 //
 // Neither `stdin` nor `stdout` are necessarily buffered. If the
 // `StageFunc` requires buffering, it needs to arrange that itself.
@@ -32,57 +32,62 @@ 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
-	panicHandler StagePanicHandler
+	name string
+	f    StageFunc
+	done chan struct{}
+	err  error
 }
 
+var (
+	_ Stage = (*goStage)(nil)
+)
+
 func (s *goStage) Name() string {
 	return s.name
 }
 
-func (s *goStage) SetPanicHandler(ph StagePanicHandler) {
-	s.panicHandler = ph
+func (s *goStage) Preferences() StagePreferences {
+	return StagePreferences{
+		StdinPreference:  IOPreferenceUndefined,
+		StdoutPreference: IOPreferenceUndefined,
+	}
 }
 
-func (s *goStage) Start(ctx context.Context, env Env, stdin io.ReadCloser) (io.ReadCloser, error) {
-	r, w := io.Pipe()
+func (s *goStage) Start(
+	ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser,
+) error {
+	var r io.Reader = stdin
+	if stdin, ok := stdin.(readerNopCloser); ok {
+		r = stdin.Reader
+	}
+
+	var w io.Writer = stdout
+	if stdout, ok := stdout.(writerNopCloser); ok {
+		w = stdout.Writer
+	}
 
 	go func() {
-		defer func() {
-			// 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)
-				}
+		s.err = s.f(ctx, env, r, w)
+
+		if stdout != nil {
+			if err := stdout.Close(); err != nil && s.err == nil {
+				s.err = fmt.Errorf("error closing stdout for stage %q: %w", s.Name(), err)
 			}
-			close(s.done)
-		}()
+		}
 
-		defer s.recoverPanic()
+		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)
+			}
+		}
 
-		s.err = s.f(ctx, env, stdin, w)
+		close(s.done)
 	}()
 
-	return r, nil
+	return nil
 }
 
 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)
-	}
-}