From c4146f48cc4deb4f8728f97e6c03f4b37b969b18 Mon Sep 17 00:00:00 2001 From: Michael Haggerty Date: Sun, 17 Dec 2023 11:05:24 +0100 Subject: [PATCH 01/23] Shush the linter --- pipe/scanner.go | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/pipe/scanner.go b/pipe/scanner.go index b56b58c..5ec16e8 100644 --- a/pipe/scanner.go +++ b/pipe/scanner.go @@ -56,11 +56,7 @@ func ScannerFunction( return err } } - if err := scanner.Err(); err != nil { - return err - } - - return nil + return scanner.Err() // `p.AddFunction()` arranges for `stdout` to be closed. }, ) From de3225f32aa7d76da3ce015e325c7835d89ab125 Mon Sep 17 00:00:00 2001 From: Michael Haggerty Date: Wed, 8 Apr 2026 19:38:01 +0200 Subject: [PATCH 02/23] pipeline_test.go: remove unnecessary tmpdirs and simplify test setup Ported from version-2 branch commits: - 95dc2e8 pipeline_test.go: get rid of a bunch of unnecessary tmpdirs - 5fdc22a TestPipelineStdinThatIsNeverClosed(): create stdin more simply - c2c9802 pipeline_test.go: use WithStdoutCloser() to close stdout pipes Tests that don't run external commands (or whose commands don't need a specific working directory) don't need t.TempDir(). --- pipe/pipeline_test.go | 101 +++++++++--------------------------------- 1 file changed, 22 insertions(+), 79 deletions(-) diff --git a/pipe/pipeline_test.go b/pipe/pipeline_test.go index bebc931..6335483 100644 --- a/pipe/pipeline_test.go +++ b/pipe/pipeline_test.go @@ -29,12 +29,9 @@ func TestMain(m *testing.M) { func TestPipelineFirstStageFailsToStart(t *testing.T) { t.Parallel() ctx := context.Background() - - dir := t.TempDir() - startErr := errors.New("foo") - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( ErrorStartingStage{startErr}, ErrorStartingStage{errors.New("this error should never happen")}, @@ -45,12 +42,9 @@ func TestPipelineFirstStageFailsToStart(t *testing.T) { func TestPipelineSecondStageFailsToStart(t *testing.T) { t.Parallel() ctx := context.Background() - - dir := t.TempDir() - startErr := errors.New("foo") - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( seqFunction(20000), ErrorStartingStage{startErr}, @@ -61,10 +55,7 @@ func TestPipelineSecondStageFailsToStart(t *testing.T) { func TestPipelineSingleCommandOutput(t *testing.T) { t.Parallel() ctx := context.Background() - - dir := t.TempDir() - - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add(pipe.Command("echo", "hello world")) out, err := p.Output(ctx) if assert.NoError(t, err) { @@ -75,12 +66,9 @@ func TestPipelineSingleCommandOutput(t *testing.T) { func TestPipelineSingleCommandWithStdout(t *testing.T) { t.Parallel() ctx := context.Background() - - dir := t.TempDir() - stdout := &bytes.Buffer{} - p := pipe.New(pipe.WithDir(dir), pipe.WithStdout(stdout)) + p := pipe.New(pipe.WithStdout(stdout)) p.Add(pipe.Command("echo", "hello world")) if assert.NoError(t, p.Run(ctx)) { assert.Equal(t, "hello world\n", stdout.String()) @@ -158,10 +146,7 @@ func TestPipelineStdinThatIsNeverClosed(t *testing.T) { func TestNontrivialPipeline(t *testing.T) { t.Parallel() ctx := context.Background() - - dir := t.TempDir() - - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( pipe.Command("echo", "hello world"), pipe.Command("sed", "s/hello/goodbye/"), @@ -210,9 +195,6 @@ func TestPipelineReadFromSlowly2(t *testing.T) { t.Parallel() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() - - dir := t.TempDir() - r, w := io.Pipe() var buf []byte @@ -236,7 +218,7 @@ func TestPipelineReadFromSlowly2(t *testing.T) { } }() - p := pipe.New(pipe.WithDir(dir), pipe.WithStdout(w)) + p := pipe.New(pipe.WithStdout(w)) p.Add(pipe.Command("seq", "100")) assert.NoError(t, p.Run(ctx)) @@ -252,10 +234,7 @@ func TestPipelineReadFromSlowly2(t *testing.T) { func TestPipelineTwoCommandsPiping(t *testing.T) { t.Parallel() ctx := context.Background() - - dir := t.TempDir() - - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add(pipe.Command("echo", "hello world")) assert.Panics(t, func() { p.Add(pipe.Command("")) }) out, err := p.Output(ctx) @@ -282,10 +261,7 @@ func TestPipelineDir(t *testing.T) { func TestPipelineExit(t *testing.T) { t.Parallel() ctx := context.Background() - - dir := t.TempDir() - - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( pipe.Command("false"), pipe.Command("true"), @@ -316,11 +292,10 @@ func TestPipelineInterrupted(t *testing.T) { } t.Parallel() - dir := t.TempDir() stdout := &bytes.Buffer{} - p := pipe.New(pipe.WithDir(dir), pipe.WithStdout(stdout)) + p := pipe.New(pipe.WithStdout(stdout)) p.Add(pipe.Command("sleep", "10")) ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond) @@ -339,11 +314,10 @@ func TestPipelineCanceled(t *testing.T) { } t.Parallel() - dir := t.TempDir() stdout := &bytes.Buffer{} - p := pipe.New(pipe.WithDir(dir), pipe.WithStdout(stdout)) + p := pipe.New(pipe.WithStdout(stdout)) p.Add(pipe.Command("sleep", "10")) ctx, cancel := context.WithCancel(context.Background()) @@ -367,9 +341,8 @@ func TestLittleEPIPE(t *testing.T) { } t.Parallel() - dir := t.TempDir() - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( pipe.Command("sh", "-c", "sleep 1; echo foo"), pipe.Command("true"), @@ -391,9 +364,8 @@ func TestBigEPIPE(t *testing.T) { } t.Parallel() - dir := t.TempDir() - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( pipe.Command("seq", "100000"), pipe.Command("true"), @@ -415,9 +387,8 @@ func TestIgnoredSIGPIPE(t *testing.T) { } t.Parallel() - dir := t.TempDir() - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( pipe.IgnoreError(pipe.Command("seq", "100000"), pipe.IsSIGPIPE), pipe.Command("echo", "foo"), @@ -433,11 +404,8 @@ func TestIgnoredSIGPIPE(t *testing.T) { func TestFunction(t *testing.T) { t.Parallel() ctx := context.Background() - - dir := t.TempDir() - t.Run("successful function", func(t *testing.T) { - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( pipe.Print("hello world"), pipe.Function( @@ -463,7 +431,6 @@ func TestFunction(t *testing.T) { t.Run("panic with handler", func(t *testing.T) { p := pipe.New( - pipe.WithDir(dir), pipe.WithStagePanicHandler(func(p any) error { err := fmt.Errorf("panic handled: %v", p) return err @@ -488,10 +455,7 @@ func TestFunction(t *testing.T) { func TestPipelineWithFunction(t *testing.T) { t.Parallel() ctx := context.Background() - - dir := t.TempDir() - - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( pipe.Command("echo", "-n", "hello world"), pipe.Function( @@ -552,10 +516,7 @@ func seqFunction(n int) pipe.Stage { func TestPipelineWithLinewiseFunction(t *testing.T) { t.Parallel() ctx := context.Background() - - dir := t.TempDir() - - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() // Print the numbers from 1 to 20 (generated from scratch): p.Add( seqFunction(20), @@ -694,10 +655,7 @@ func TestScannerFinishEarly(t *testing.T) { func TestPrintln(t *testing.T) { t.Parallel() ctx := context.Background() - - dir := t.TempDir() - - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add(pipe.Println("Look Ma, no hands!")) out, err := p.Output(ctx) if assert.NoError(t, err) { @@ -708,10 +666,7 @@ func TestPrintln(t *testing.T) { func TestPrintf(t *testing.T) { t.Parallel() ctx := context.Background() - - dir := t.TempDir() - - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add(pipe.Printf("Strangely recursive: %T", p)) out, err := p.Output(ctx) if assert.NoError(t, err) { @@ -903,11 +858,8 @@ func TestErrors(t *testing.T) { func BenchmarkSingleProgram(b *testing.B) { ctx := context.Background() - - dir := b.TempDir() - for i := 0; i < b.N; i++ { - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( pipe.Command("true"), ) @@ -917,11 +869,8 @@ func BenchmarkSingleProgram(b *testing.B) { func BenchmarkTenPrograms(b *testing.B) { ctx := context.Background() - - dir := b.TempDir() - for i := 0; i < b.N; i++ { - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( pipe.Command("echo", "hello world"), pipe.Command("cat"), @@ -943,16 +892,13 @@ func BenchmarkTenPrograms(b *testing.B) { func BenchmarkTenFunctions(b *testing.B) { ctx := context.Background() - - dir := b.TempDir() - cp := func(_ context.Context, _ pipe.Env, stdin io.Reader, stdout io.Writer) error { _, err := io.Copy(stdout, stdin) return err } for i := 0; i < b.N; i++ { - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( pipe.Println("hello world"), pipe.Function("copy1", cp), @@ -974,16 +920,13 @@ func BenchmarkTenFunctions(b *testing.B) { func BenchmarkTenMixedStages(b *testing.B) { ctx := context.Background() - - dir := b.TempDir() - cp := func(_ context.Context, _ pipe.Env, stdin io.Reader, stdout io.Writer) error { _, err := io.Copy(stdout, stdin) return err } for i := 0; i < b.N; i++ { - p := pipe.New(pipe.WithDir(dir)) + p := pipe.New() p.Add( pipe.Command("echo", "hello world"), pipe.Function("copy1", cp), From 9fd5b2c8f99bc5b78b55342a64fb81bf0a94f2f1 Mon Sep 17 00:00:00 2001 From: Michael Haggerty Date: Sat, 16 Dec 2023 17:56:03 +0100 Subject: [PATCH 03/23] Add some benchmarks of moving a bunch of data through a pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add some benchmarks that move MB-scale data through pipelines consisting of alternating commands and functions, one in small writes, and one buffered into larger writes, then processing it one line at a time. This is not so efficient, because every transition from `Function` → `Command` requires an extra (hidden) goroutine that copies the data from an `io.Reader` to a `*os.File`. We can make this faster! --- pipe/pipeline_test.go | 91 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 91 insertions(+) diff --git a/pipe/pipeline_test.go b/pipe/pipeline_test.go index 6335483..8559e52 100644 --- a/pipe/pipeline_test.go +++ b/pipe/pipeline_test.go @@ -946,6 +946,97 @@ func BenchmarkTenMixedStages(b *testing.B) { } } +func BenchmarkMoreDataUnbuffered(b *testing.B) { + ctx := context.Background() + + cp := func(_ context.Context, _ pipe.Env, stdin io.Reader, stdout io.Writer) error { + _, err := io.Copy(stdout, stdin) + return err + } + + for i := 0; i < b.N; i++ { + count := 0 + p := pipe.New() + p.Add( + pipe.Function( + "seq", + func(ctx context.Context, _ pipe.Env, stdin io.Reader, stdout io.Writer) error { + for i := 1; i <= 100000; i++ { + fmt.Fprintln(stdout, i) + } + return nil + }, + ), + pipe.Command("cat"), + pipe.Function("copy2", cp), + pipe.Command("cat"), + pipe.Function("copy3", cp), + pipe.Command("cat"), + pipe.Function("copy4", cp), + pipe.Command("cat"), + pipe.Function("copy5", cp), + pipe.Command("cat"), + pipe.LinewiseFunction( + "count", + func(ctx context.Context, _ pipe.Env, line []byte, stdout *bufio.Writer) error { + count++ + return nil + }, + ), + ) + err := p.Run(ctx) + if assert.NoError(b, err) { + assert.EqualValues(b, 100000, count) + } + } +} + +func BenchmarkMoreDataBuffered(b *testing.B) { + ctx := context.Background() + + cp := func(_ context.Context, _ pipe.Env, stdin io.Reader, stdout io.Writer) error { + _, err := io.Copy(stdout, stdin) + return err + } + + for i := 0; i < b.N; i++ { + count := 0 + p := pipe.New() + p.Add( + pipe.Function( + "seq", + func(ctx context.Context, _ pipe.Env, stdin io.Reader, stdout io.Writer) error { + out := bufio.NewWriter(stdout) + for i := 1; i <= 1000000; i++ { + fmt.Fprintln(out, i) + } + return out.Flush() + }, + ), + pipe.Command("cat"), + pipe.Function("copy2", cp), + pipe.Command("cat"), + pipe.Function("copy3", cp), + pipe.Command("cat"), + pipe.Function("copy4", cp), + pipe.Command("cat"), + pipe.Function("copy5", cp), + pipe.Command("cat"), + pipe.LinewiseFunction( + "count", + func(ctx context.Context, _ pipe.Env, line []byte, stdout *bufio.Writer) error { + count++ + return nil + }, + ), + ) + err := p.Run(ctx) + if assert.NoError(b, err) { + assert.EqualValues(b, 1000000, count) + } + } +} + func genErr(err error) pipe.StageFunc { return func(_ context.Context, _ pipe.Env, _ io.Reader, _ io.Writer) error { return err From ca78f76246868a46aeb74b44e39ac3406cd02219 Mon Sep 17 00:00:00 2001 From: Michael Haggerty Date: Fri, 15 Dec 2023 10:58:55 +0100 Subject: [PATCH 04/23] Simplify the `NopCloser`s MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Rename * `newNopCloser()` → `newReaderNopCloser()` * `nopCloser` → `readerNopCloser` * `nopCloserWriterTo` → `readerWriterToNopCloser` * `nopWriteCloser` → `writerNopCloser` to help keep readers and writers straight and because only the `Close()` part is a NOP. * Move `writerNopCloser` to `nop_closer.go` to be with its siblings. --- pipe/command.go | 4 ++-- pipe/nop_closer.go | 45 ++++++++++++++++++++++++++++++++------------- pipe/pipeline.go | 14 +++----------- 3 files changed, 37 insertions(+), 26 deletions(-) diff --git a/pipe/command.go b/pipe/command.go index 2113902..a20cda8 100644 --- a/pipe/command.go +++ b/pipe/command.go @@ -74,9 +74,9 @@ func (s *commandStage) Start( // See the long comment in `Pipeline.Start()` for the // explanation of this special case. switch stdin := stdin.(type) { - case nopCloser: + case readerNopCloser: s.cmd.Stdin = stdin.Reader - case nopCloserWriterTo: + case readerWriterToNopCloser: s.cmd.Stdin = stdin.Reader default: s.cmd.Stdin = stdin diff --git a/pipe/nop_closer.go b/pipe/nop_closer.go index d435d0a..8c66c72 100644 --- a/pipe/nop_closer.go +++ b/pipe/nop_closer.go @@ -6,29 +6,48 @@ package pipe import "io" -// newNopCloser returns a ReadCloser with a no-op Close method wrapping -// the provided io.Reader r. -// If r implements io.WriterTo, the returned io.ReadCloser will implement io.WriterTo -// by forwarding calls to r. -func newNopCloser(r io.Reader) io.ReadCloser { +// newReaderNopCloser returns a ReadCloser with a no-op Close method, +// wrapping the provided io.Reader `r`. If `r` implements +// `io.WriterTo`, the returned `io.ReadCloser` will also implement +// `io.WriterTo` by forwarding calls to `r`. +func newReaderNopCloser(r io.Reader) io.ReadCloser { if _, ok := r.(io.WriterTo); ok { - return nopCloserWriterTo{r} + return readerWriterToNopCloser{r} } - return nopCloser{r} + return readerNopCloser{r} } -type nopCloser struct { +// readerNopCloser is a ReadCloser that wraps a provided `io.Reader`, +// but whose `Close()` method does nothing. We don't need to check +// whether the wrapped reader also implements `io.WriterTo`, since +// it's always unwrapped before use. +type readerNopCloser struct { io.Reader } -func (nopCloser) Close() error { return nil } +func (readerNopCloser) Close() error { + return nil +} -type nopCloserWriterTo struct { +// readerWriterToNopCloser is like `readerNopCloser` except that it +// also implements `io.WriterTo` by delegating `WriteTo()` to the +// wrapped `io.Reader` (which must also implement `io.WriterTo`). +type readerWriterToNopCloser struct { io.Reader } -func (nopCloserWriterTo) Close() error { return nil } +func (readerWriterToNopCloser) Close() error { return nil } + +func (r readerWriterToNopCloser) WriteTo(w io.Writer) (n int64, err error) { + return r.Reader.(io.WriterTo).WriteTo(w) +} + +// writerNopCloser is a WriteCloser that wraps a provided `io.Writer`, +// but whose `Close()` method does nothing. +type writerNopCloser struct { + io.Writer +} -func (c nopCloserWriterTo) WriteTo(w io.Writer) (n int64, err error) { - return c.Reader.(io.WriterTo).WriteTo(w) +func (w writerNopCloser) Close() error { + return nil } diff --git a/pipe/pipeline.go b/pipe/pipeline.go index 8bc3a37..7391a2c 100644 --- a/pipe/pipeline.go +++ b/pipe/pipeline.go @@ -70,14 +70,6 @@ type Pipeline struct { var emptyEventHandler = func(_ *Event) {} -type nopWriteCloser struct { - io.Writer -} - -func (w nopWriteCloser) Close() error { - return nil -} - type NewPipeFn func(opts ...Option) *Pipeline // NewPipeline returns a Pipeline struct with all of the `options` @@ -114,7 +106,7 @@ func WithStdin(stdin io.Reader) Option { // WithStdout assigns stdout to the last command in the pipeline. func WithStdout(stdout io.Writer) Option { return func(p *Pipeline) { - p.stdout = nopWriteCloser{stdout} + p.stdout = writerNopCloser{stdout} } } @@ -277,7 +269,7 @@ func (p *Pipeline) Start(ctx context.Context) error { // own `nopCloser`, which behaves like `io.NopCloser`, except // that `pipe.CommandStage` knows how to unwrap it before // passing it to `exec.Cmd`. - nextStdin = newNopCloser(p.stdin) + nextStdin = newReaderNopCloser(p.stdin) } for i, s := range p.stages { @@ -325,7 +317,7 @@ func (p *Pipeline) Start(ctx context.Context) error { func (p *Pipeline) Output(ctx context.Context) ([]byte, error) { var buf bytes.Buffer - p.stdout = nopWriteCloser{&buf} + p.stdout = writerNopCloser{&buf} err := p.Run(ctx) return buf.Bytes(), err } From 2f59602106c3cdf78f8ddff1875dafa5d2d7a5d7 Mon Sep 17 00:00:00 2001 From: Michael Haggerty Date: Fri, 15 Dec 2023 21:03:26 +0100 Subject: [PATCH 05/23] 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. --- pipe/command.go | 102 +++++++++++---- pipe/command_nil_panic_test.go | 2 +- pipe/command_test.go | 3 +- pipe/function.go | 73 ++++++----- pipe/iocopier.go | 62 --------- pipe/memorylimit.go | 132 ++++--------------- pipe/memorylimit_test.go | 168 +++--------------------- pipe/pipeline.go | 228 +++++++++++++++++++-------------- pipe/pipeline_test.go | 78 ++++++----- pipe/stage.go | 138 ++++++++++++++++++-- 10 files changed, 472 insertions(+), 514 deletions(-) delete mode 100644 pipe/iocopier.go diff --git a/pipe/command.go b/pipe/command.go index a20cda8..3d61782 100644 --- a/pipe/command.go +++ b/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 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 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 } } diff --git a/pipe/command_nil_panic_test.go b/pipe/command_nil_panic_test.go index 740af73..d19ce44 100644 --- a/pipe/command_nil_panic_test.go +++ b/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") } diff --git a/pipe/command_test.go b/pipe/command_test.go index ca5a8c0..531f11f 100644 --- a/pipe/command_test.go +++ b/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), + ) }) } } diff --git a/pipe/function.go b/pipe/function.go index e8d9522..fe8abd0 100644 --- a/pipe/function.go +++ b/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) - } -} diff --git a/pipe/iocopier.go b/pipe/iocopier.go deleted file mode 100644 index 78a9143..0000000 --- a/pipe/iocopier.go +++ /dev/null @@ -1,62 +0,0 @@ -package pipe - -import ( - "context" - "errors" - "io" - "os" -) - -// ioCopier is a stage that copies its stdin to a specified -// `io.Writer`. It generates no stdout itself. -type ioCopier struct { - w io.WriteCloser - done chan struct{} - err error -} - -func newIOCopier(w io.WriteCloser) *ioCopier { - return &ioCopier{ - w: w, - done: make(chan struct{}), - } -} - -func (s *ioCopier) Name() string { - return "ioCopier" -} - -// This method always returns `nil, nil`. -func (s *ioCopier) Start(_ context.Context, _ Env, r io.ReadCloser) (io.ReadCloser, error) { - go func() { - _, err := io.Copy(s.w, r) - // We don't consider `ErrClosed` an error (FIXME: is this - // correct?): - if err != nil && !errors.Is(err, os.ErrClosed) { - s.err = err - } - if err := r.Close(); err != nil && s.err == nil { - s.err = err - } - if err := s.w.Close(); err != nil && s.err == nil { - s.err = err - } - close(s.done) - }() - - // FIXME: if `s.w.Write()` is blocking (e.g., because there is a - // downstream process that is not reading from the other side), - // there's no way to terminate the copy when the context expires. - // This is not too bad, because the `io.Copy()` call will exit by - // itself when its input is closed. - // - // We could, however, be smarter about exiting more quickly if the - // context expires but `s.w.Write()` is not blocking. - - return nil, nil -} - -func (s *ioCopier) Wait() error { - <-s.done - return s.err -} diff --git a/pipe/memorylimit.go b/pipe/memorylimit.go index 8e91dc1..48e1bb6 100644 --- a/pipe/memorylimit.go +++ b/pipe/memorylimit.go @@ -11,12 +11,12 @@ import ( const memoryPollInterval = time.Second -// ErrMemoryLimitExceeded is the error that will be used to kill a process, if -// necessary, from MemoryLimit. +// ErrMemoryLimitExceeded is the error that will be used to kill a +// process, if necessary, from MemoryLimit. var ErrMemoryLimitExceeded = errors.New("memory limit exceeded") -// LimitableStage is the superset of Stage that must be implemented by stages -// passed to MemoryLimit and MemoryObserver. +// LimitableStage is the superset of `Stage` that must be implemented +// by stages passed to MemoryLimit and MemoryObserver. type LimitableStage interface { Stage @@ -89,102 +89,6 @@ func killAtLimit(byteLimit uint64, eventHandler func(e *Event)) memoryWatchFunc } } -// MemoryLimitWithObserver combines MemoryLimit and MemoryObserver into a single -// stage that uses one goroutine instead of two. It watches the memory usage of -// the stage, kills the process if it exceeds byteLimit, and logs peak memory -// usage when the stage exits. -// -// Use this instead of MemoryLimit(MemoryObserver(stage, h), limit, h) to save -// one goroutine per pipeline stage. -func MemoryLimitWithObserver(stage Stage, byteLimit uint64, eventHandler func(e *Event)) Stage { - limitableStage, ok := stage.(LimitableStage) - if !ok { - eventHandler(&Event{ - Command: stage.Name(), - Msg: "invalid pipe.MemoryLimitWithObserver usage", - Err: fmt.Errorf("invalid pipe.MemoryLimitWithObserver usage"), - }) - return stage - } - - return &memoryWatchStage{ - nameSuffix: " with memory limit", - stage: limitableStage, - watch: killAtLimitAndObserve(byteLimit, eventHandler), - } -} - -func killAtLimitAndObserve(byteLimit uint64, eventHandler func(e *Event)) memoryWatchFunc { - return func(ctx context.Context, stage LimitableStage) { - var ( - maxRSS uint64 - samples, errCount, consecutiveErrors int - killed bool - ) - - t := time.NewTicker(memoryPollInterval) - defer t.Stop() - - for { - select { - case <-ctx.Done(): - eventHandler(&Event{ - Command: stage.Name(), - Msg: "peak memory usage", - Context: map[string]interface{}{ - "max_rss_bytes": maxRSS, - "samples": samples, - "errors": errCount, - }, - }) - return - case <-t.C: - if killed { - continue - } - - rss, err := stage.GetRSSAnon(ctx) - if err != nil { - if !errors.Is(err, errProcessInfoMissing) { - errCount++ - consecutiveErrors++ - if consecutiveErrors == 2 { - eventHandler(&Event{ - Command: stage.Name(), - Msg: "error getting RSS", - Err: err, - }) - } - } else { - consecutiveErrors = 0 - } - continue - } - - consecutiveErrors = 0 - samples++ - if rss > maxRSS { - maxRSS = rss - } - - if rss >= byteLimit { - eventHandler(&Event{ - Command: stage.Name(), - Msg: "stage exceeded allowed memory use", - Err: fmt.Errorf("stage exceeded allowed memory use"), - Context: map[string]interface{}{ - "limit": byteLimit, - "used": rss, - }, - }) - stage.Kill(ErrMemoryLimitExceeded) - killed = true - } - } - } - } -} - // MemoryObserver watches memory use of the stage and logs the maximum // value when the stage exits. func MemoryObserver(stage Stage, eventHandler func(e *Event)) Stage { @@ -271,12 +175,24 @@ func (m *memoryWatchStage) Name() string { return m.stage.Name() + m.nameSuffix } -func (m *memoryWatchStage) Start(ctx context.Context, env Env, stdin io.ReadCloser) (io.ReadCloser, error) { - io, err := m.stage.Start(ctx, env, stdin) - if err != nil { - return nil, err +func (m *memoryWatchStage) Preferences() StagePreferences { + return m.stage.Preferences() +} + +func (m *memoryWatchStage) Start( + ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser, +) error { + if err := m.stage.Start(ctx, env, stdin, stdout); err != nil { + return err } + m.monitor(ctx) + + return nil +} + +// monitor starts up a goroutine that monitors the memory of `m`. +func (m *memoryWatchStage) monitor(ctx context.Context) { ctx, cancel := context.WithCancel(ctx) m.cancel = cancel m.wg.Add(1) @@ -285,14 +201,14 @@ func (m *memoryWatchStage) Start(ctx context.Context, env Env, stdin io.ReadClos m.watch(ctx, m.stage) m.wg.Done() }() - - return io, nil } func (m *memoryWatchStage) Wait() error { - err := m.stage.Wait() + if err := m.stage.Wait(); err != nil { + return err + } m.stopWatching() - return err + return nil } func (m *memoryWatchStage) GetRSSAnon(ctx context.Context) (uint64, error) { diff --git a/pipe/memorylimit_test.go b/pipe/memorylimit_test.go index 37d3edd..e84f465 100644 --- a/pipe/memorylimit_test.go +++ b/pipe/memorylimit_test.go @@ -8,6 +8,7 @@ import ( "log" "os" "strings" + "syscall" "testing" "time" @@ -112,165 +113,36 @@ func TestMemoryLimitTreeMem(t *testing.T) { require.ErrorContains(t, err, "memory limit exceeded") } -type closeWrapper struct { - io.Writer - close func() error -} - -func (w closeWrapper) Close() error { - return w.close() -} - -func TestMemoryLimitWithObserverSimple(t *testing.T) { - t.Parallel() - msg, err := testMemoryLimitWithObserver(t, 400, 10_000_000, pipe.Command("less")) - assert.Contains(t, msg, "exceeded allowed memory") - assert.Contains(t, msg, "limit=10000000") - require.ErrorContains(t, err, "memory limit exceeded") -} - -func TestMemoryLimitWithObserverTreeMem(t *testing.T) { - t.Parallel() - msg, err := testMemoryLimitWithObserver(t, 400, 10_000_000, pipe.Command("sh", "-c", "less; :")) - assert.Contains(t, msg, "exceeded allowed memory") - assert.Contains(t, msg, "limit=10000000") - require.ErrorContains(t, err, "memory limit exceeded") -} - -func TestMemoryLimitWithObserverBelowLimit(t *testing.T) { - t.Parallel() - rss := testMemoryLimitWithObserverBelowLimit(t, 400, pipe.Command("less")) - require.Greater(t, rss, 400_000_000) -} - -func TestMemoryLimitWithObserverBelowLimitTreeMem(t *testing.T) { - t.Parallel() - rss := testMemoryLimitWithObserverBelowLimit(t, 400, pipe.Command("sh", "-c", "less; :")) - require.Greater(t, rss, 400_000_000) -} - -func TestMemoryLimitWithObserverLogsPeakOnKill(t *testing.T) { - t.Parallel() - msg, err := testMemoryLimitWithObserver(t, 400, 10_000_000, pipe.Command("less")) - // Verify both limit-exceeded AND peak memory are logged (matching - // the behavior of MemoryLimit(MemoryObserver(...))) - assert.Contains(t, msg, "exceeded allowed memory") - assert.Contains(t, msg, "peak memory usage") - require.ErrorContains(t, err, "memory limit exceeded") -} - -func testMemoryLimitWithObserverBelowLimit(t *testing.T, mbs int, stage pipe.Stage) int { - ctx := context.Background() - - stdinReader, stdinWriter := io.Pipe() - - devNull, err := os.OpenFile("/dev/null", os.O_WRONLY, 0) - require.NoError(t, err) - - buf := &bytes.Buffer{} - logger := log.New(buf, "testMemoryLimitWithObserver", log.Ldate|log.Ltime) - - // Use a high limit so it won't be hit — we want to verify the observer part - p := pipe.New(pipe.WithDir("/"), pipe.WithStdin(stdinReader), pipe.WithStdout(devNull)) - p.Add(pipe.MemoryLimitWithObserver(stage, 100*1024*1024*1024, LogEventHandler(logger))) - require.NoError(t, p.Start(ctx)) - - var bytes [1_000_000]byte - for i := 0; i < mbs; i++ { - n, err := stdinWriter.Write(bytes[:]) - require.NoError(t, err) - require.Equal(t, len(bytes), n) - } - - time.Sleep(2 * time.Second) - - require.NoError(t, stdinWriter.Close()) - require.NoError(t, p.Wait()) - - // Verify that peak memory usage was logged (the observer part) - output := buf.String() - assert.Contains(t, output, "peak memory usage") - - return maxBytes(output) -} - -func testMemoryLimitWithObserver(t *testing.T, mbs int, limit uint64, stage pipe.Stage) (string, error) { - ctx := context.Background() - - stdinReader, stdinWriter := io.Pipe() - - devNull, err := os.OpenFile("/dev/null", os.O_WRONLY, 0) - require.NoError(t, err) - - closedErr := fmt.Errorf("stdout was closed") - stdout := closeWrapper{ - Writer: devNull, - close: func() error { - require.NoError(t, stdinReader.CloseWithError(closedErr)) - return nil - }, - } - - buf := &bytes.Buffer{} - logger := log.New(buf, "testMemoryLimitWithObserver", log.Ldate|log.Ltime) - - p := pipe.New(pipe.WithDir("/"), pipe.WithStdin(stdinReader), pipe.WithStdoutCloser(stdout)) - p.Add(pipe.MemoryLimitWithObserver(stage, limit, LogEventHandler(logger))) - require.NoError(t, p.Start(ctx)) - - var bytes [1_000_000]byte - for i := 0; i < mbs; i++ { - _, err := stdinWriter.Write(bytes[:]) - if err != nil { - require.ErrorIs(t, err, closedErr) - } - } - - require.NoError(t, stdinWriter.Close()) - err = p.Wait() - - return buf.String(), err -} - func testMemoryLimit(t *testing.T, mbs int, limit uint64, stage pipe.Stage) (string, error) { ctx := context.Background() - stdinReader, stdinWriter := io.Pipe() - devNull, err := os.OpenFile("/dev/null", os.O_WRONLY, 0) require.NoError(t, err) - // io.Pipe doesn't know if anything is listening on the other end, so once - // our process is expectedly killed then we'll end up blocked trying to - // write to it. To workaround this, make sure we close the pipe reader when - // we've detected that the process has exited (i.e. when stdout has been - // closed). This will cause our write to immediately fail with this error. - closedErr := fmt.Errorf("stdout was closed") - stdout := closeWrapper{ - Writer: devNull, - close: func() error { - require.NoError(t, stdinReader.CloseWithError(closedErr)) - return nil - }, - } - buf := &bytes.Buffer{} logger := log.New(buf, "testMemoryObserver", log.Ldate|log.Ltime) - p := pipe.New(pipe.WithDir("/"), pipe.WithStdin(stdinReader), pipe.WithStdoutCloser(stdout)) - p.Add(pipe.MemoryLimit(stage, limit, LogEventHandler(logger))) + p := pipe.New(pipe.WithDir("/"), pipe.WithStdoutCloser(devNull)) + p.Add( + pipe.Function( + "write-to-less", + func(ctx context.Context, _ pipe.Env, _ io.Reader, stdout io.Writer) error { + // Write some nonsense data to less. + var bytes [1_000_000]byte + for i := 0; i < mbs; i++ { + _, err := stdout.Write(bytes[:]) + if err != nil { + require.ErrorIs(t, err, syscall.EPIPE) + } + } + + return nil + }, + ), + pipe.MemoryLimit(stage, limit, LogEventHandler(logger)), + ) require.NoError(t, p.Start(ctx)) - // Write some nonsense data to less. - var bytes [1_000_000]byte - for i := 0; i < mbs; i++ { - _, err := stdinWriter.Write(bytes[:]) - if err != nil { - require.ErrorIs(t, err, closedErr) - } - } - - require.NoError(t, stdinWriter.Close()) err = p.Wait() return buf.String(), err diff --git a/pipe/pipeline.go b/pipe/pipeline.go index 7391a2c..4fdd4a5 100644 --- a/pipe/pipeline.go +++ b/pipe/pipeline.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "io" + "os" "sync/atomic" ) @@ -28,7 +29,6 @@ type Env struct { // and is not reported to the caller. // //revive:disable:error-naming -//nolint:staticcheck // ST1012: FinishEarly is the intentional name for this sentinel error var FinishEarly = errors.New("finish stage early") //revive:enable:error-naming @@ -54,7 +54,7 @@ type ContextValuesFunc func(context.Context) []EnvVar type Pipeline struct { env Env - stdin io.Reader + stdin io.ReadCloser stdout io.WriteCloser stages []Stage cancel func() @@ -65,10 +65,9 @@ type Pipeline struct { started uint32 eventHandler func(e *Event) - panicHandler StagePanicHandler } -var emptyEventHandler = func(_ *Event) {} +var emptyEventHandler = func(e *Event) {} type NewPipeFn func(opts ...Option) *Pipeline @@ -99,7 +98,51 @@ func WithDir(dir string) Option { // WithStdin assigns stdin to the first command in the pipeline. func WithStdin(stdin io.Reader) Option { return func(p *Pipeline) { - p.stdin = stdin + // We don't want the first stage to close `stdin`, and it is + // not even necessarily an `io.ReadCloser`. So wrap it in a + // fake `io.ReadCloser` whose `Close()` method doesn't do + // anything. + // + // We could use `io.NopCloser()` for this purpose, but that + // would have a subtle problem. If the first stage is a + // `Command`, then it wants to set the `exec.Cmd`'s `Stdin` to + // an `io.Reader` corresponding to `p.stdin`. If `Cmd.Stdin` + // is an `*os.File`, then `exec.Cmd` will pass the file + // descriptor to the subcommand directly; there is no need to + // create a pipe and copy the data into the input side of the + // pipe. But if `p.stdin` is not an `*os.File`, then this + // optimization is prevented. And even worse, it also has the + // side effect that the goroutine that copies from `Cmd.Stdin` + // into the pipe doesn't terminate until that fd is closed by + // the writing side. + // + // That isn't always what we want. Consider, for example, the + // following snippet, where the subcommand's stdin is set to + // the stdin of the enclosing Go program, but wrapped with + // `io.NopCloser`: + // + // cmd := exec.Command("ls") + // cmd.Stdin = io.NopCloser(os.Stdin) + // cmd.Stdout = os.Stdout + // cmd.Stderr = os.Stderr + // cmd.Run() + // + // In this case, we don't want the Go program to wait for + // `os.Stdin` to close (because `ls` isn't even trying to read + // from its stdin). But it does: `exec.Cmd` doesn't recognize + // that `Cmd.Stdin` is an `*os.File`, so it sets up a pipe and + // copies the data itself, and this goroutine doesn't + // terminate until `cmd.Stdin` (i.e., the Go program's own + // stdin) is closed. But if, for example, the Go program is + // run from an interactive shell session, that might never + // happen, in which case the program will fail to terminate, + // even after `ls` exits. + // + // So instead, in this special case, we wrap `stdin` in our + // own `nopCloser`, which behaves like `io.NopCloser`, except + // that `pipe.CommandStage` knows how to unwrap it before + // passing it to `exec.Cmd`. + p.stdin = newReaderNopCloser(stdin) } } @@ -173,20 +216,6 @@ func WithEventHandler(handler func(e *Event)) Option { } } -// WithStagePanicHandler sets a panic handler for the stages within a pipeline. -// When a pipeline stage panics, the provided handler will be invoked, allowing -// the client to handle the panic in whatever way they see fit. -// -// Note: -// - Only the Function stage supports this functionality. -// - The client is responsible for deciding whether to recover from the panic or panicking again. -// - If a panic handler is not set, the panic will be propagated normally. -func WithStagePanicHandler(ph StagePanicHandler) Option { - return func(p *Pipeline) { - p.panicHandler = ph - } -} - func (p *Pipeline) hasStarted() bool { return atomic.LoadUint32(&p.started) != 0 } @@ -212,6 +241,12 @@ func (p *Pipeline) AddWithIgnoredError(em ErrorMatcher, stages ...Stage) { } } +type stageStarter struct { + prefs StagePreferences + stdin io.ReadCloser + stdout io.WriteCloser +} + // 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. @@ -223,93 +258,94 @@ func (p *Pipeline) Start(ctx context.Context) error { atomic.StoreUint32(&p.started, 1) ctx, p.cancel = context.WithCancel(ctx) - var nextStdin io.ReadCloser + // We need to decide how to start the stages, especially what + // pipes to use to connect adjacent stages (`os.Pipe()` vs. + // `io.Pipe()`) based on the two stages' preferences. + stageStarters := make([]stageStarter, len(p.stages), len(p.stages)+1) + + // Collect information about each stage's type and preferences: + for i, s := range p.stages { + stageStarters[i].prefs = s.Preferences() + } + if p.stdin != nil { - // We don't want the first stage to actually close this, and - // `p.stdin` is not even necessarily an `io.ReadCloser`. So - // wrap it in a fake `io.ReadCloser` whose `Close()` method - // doesn't do anything. - // - // We could use `io.NopCloser()` for this purpose, but it has - // a subtle problem. If the first stage is a `Command`, then - // it wants to set the `exec.Cmd`'s `Stdin` to an `io.Reader` - // corresponding to `p.stdin`. If `Cmd.Stdin` is an - // `*os.File`, then the file descriptor can be passed to the - // subcommand directly; there is no need for this process to - // create a pipe and copy the data into the input side of the - // pipe. But if `p.stdin` is not an `*os.File`, then this - // optimization is prevented. And even worse, it also has the - // side effect that the goroutine that copies from `Cmd.Stdin` - // into the pipe doesn't terminate until that fd is closed by - // the writing side. - // - // That isn't always what we want. Consider, for example, the - // following snippet, where the subcommand's stdin is set to - // the stdin of the enclosing Go program, but wrapped with - // `io.NopCloser`: - // - // cmd := exec.Command("ls") - // cmd.Stdin = io.NopCloser(os.Stdin) - // cmd.Stdout = os.Stdout - // cmd.Stderr = os.Stderr - // cmd.Run() - // - // In this case, we don't want the Go program to wait for - // `os.Stdin` to close (because `ls` isn't even trying to read - // from its stdin). But it does: `exec.Cmd` doesn't recognize - // that `Cmd.Stdin` is an `*os.File`, so it sets up a pipe and - // copies the data itself, and this goroutine doesn't - // terminate until `cmd.Stdin` (i.e., the Go program's own - // stdin) is closed. But if, for example, the Go program is - // run from an interactive shell session, that might never - // happen, in which case the program will fail to terminate, - // even after `ls` exits. - // - // So instead, in this special case, we wrap `p.stdin` in our - // own `nopCloser`, which behaves like `io.NopCloser`, except - // that `pipe.CommandStage` knows how to unwrap it before - // passing it to `exec.Cmd`. - nextStdin = newReaderNopCloser(p.stdin) + // Arrange for the input of the 0th stage to come from + // `p.stdin`: + stageStarters[0].stdin = p.stdin } - for i, s := range p.stages { - if phs, ok := s.(StagePanicHandlerAware); ok && p.panicHandler != nil { - phs.SetPanicHandler(p.panicHandler) + if p.stdout != nil { + i := len(p.stages) - 1 + ss := &stageStarters[i] + ss.stdout = p.stdout + } + + // 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). + abort := func(i int, err error) 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 { + _ = stageStarters[i].stdin.Close() } - var err error - stdout, err := s.Start(ctx, p.env, nextStdin) - if err != nil { - // 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 nextStdin != nil { - _ = nextStdin.Close() - } + // Kill and wait for any stages that have been started + // already to finish: + p.cancel() + for _, s := range p.stages[:i] { + _ = s.Wait() + } + p.eventHandler(&Event{ + Command: p.stages[i].Name(), + Msg: "failed to start pipeline stage", + Err: err, + }) + return fmt.Errorf( + "starting pipeline stage %q: %w", p.stages[i].Name(), err, + ) + } - // Kill and wait for any stages that have been started - // already to finish: - p.cancel() - for _, s := range p.stages[:i] { - _ = s.Wait() + // Loop over all but the last stage, starting them. By the time we + // get to a stage, its stdin will have already been determined, + // but we still need to figure out its stdout and set the stdin + // that will be used for the subsequent stage. + for i, s := range p.stages[:len(p.stages)-1] { + ss := &stageStarters[i] + nextSS := &stageStarters[i+1] + + // We need to generate a pipe pair for this stage to use + // to communicate with its successor: + if ss.prefs.StdoutPreference == IOPreferenceFile || + nextSS.prefs.StdinPreference == IOPreferenceFile { + // Use an OS-level pipe for the communication: + var err error + nextSS.stdin, ss.stdout, err = os.Pipe() + if err != nil { + return abort(i, err) } - p.eventHandler(&Event{ - Command: s.Name(), - Msg: "failed to start pipeline stage", - Err: err, - }) - return fmt.Errorf("starting pipeline stage %q: %w", s.Name(), err) + } else { + nextSS.stdin, ss.stdout = io.Pipe() + } + if err := s.Start(ctx, p.env, ss.stdin, ss.stdout); err != nil { + nextSS.stdin.Close() + ss.stdout.Close() + return abort(i, err) } - nextStdin = stdout } - // If the pipeline was configured with a `stdout`, add a synthetic - // stage to copy the last stage's stdout to that writer: - if p.stdout != nil { - c := newIOCopier(p.stdout) - p.stages = append(p.stages, c) - // `ioCopier.Start()` never fails: - _, _ = c.Start(ctx, p.env, nextStdin) + // The last stage needs special handling, because its stdout + // doesn't need to flow into another stage (it's already set in + // `ss.stdout` if it's needed). + { + i := len(p.stages) - 1 + s := p.stages[i] + ss := &stageStarters[i] + + if err := s.Start(ctx, p.env, ss.stdin, ss.stdout); err != nil { + return abort(i, err) + } } return nil diff --git a/pipe/pipeline_test.go b/pipe/pipeline_test.go index 8559e52..0a5d17d 100644 --- a/pipe/pipeline_test.go +++ b/pipe/pipeline_test.go @@ -75,7 +75,7 @@ func TestPipelineSingleCommandWithStdout(t *testing.T) { } } -func TestPipelineStdinFileThatIsNeverClosed(t *testing.T) { +func TestPipelineStdinOSPipeThatIsNeverClosed(t *testing.T) { t.Parallel() // Make sure that the subprocess terminates on its own, as opposed @@ -93,7 +93,10 @@ func TestPipelineStdinFileThatIsNeverClosed(t *testing.T) { var stdout bytes.Buffer - p := pipe.New(pipe.WithStdin(r), pipe.WithStdout(&stdout)) + p := pipe.New( + pipe.WithStdin(r), + pipe.WithStdout(&stdout), + ) // Note that this command doesn't read from its stdin, so it will // terminate regardless of whether `w` gets closed: p.Add(pipe.Command("true")) @@ -103,7 +106,7 @@ func TestPipelineStdinFileThatIsNeverClosed(t *testing.T) { assert.NoError(t, p.Run(ctx)) } -func TestPipelineStdinThatIsNeverClosed(t *testing.T) { +func TestPipelineIOPipeStdinThatIsNeverClosed(t *testing.T) { t.Skip("test not run because it currently deadlocks") t.Parallel() @@ -157,7 +160,33 @@ func TestNontrivialPipeline(t *testing.T) { } } -func TestPipelineReadFromSlowly(t *testing.T) { +func TestOSPipePipelineReadFromSlowly(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + r, w, err := os.Pipe() + require.NoError(t, err) + + var buf []byte + readErr := make(chan error, 1) + + go func() { + time.Sleep(200 * time.Millisecond) + var err error + buf, err = io.ReadAll(r) + readErr <- err + }() + + p := pipe.New(pipe.WithStdoutCloser(w)) + p.Add(pipe.Command("echo", "hello world")) + assert.NoError(t, p.Run(ctx)) + + assert.NoError(t, <-readErr) + assert.Equal(t, "hello world\n", string(buf)) +} + +func TestIOPipePipelineReadFromSlowly(t *testing.T) { t.Parallel() ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) defer cancel() @@ -428,28 +457,6 @@ func TestFunction(t *testing.T) { assert.NoError(t, err) assert.EqualValues(t, "goodbye, cruel world", out) }) - - t.Run("panic with handler", func(t *testing.T) { - p := pipe.New( - pipe.WithStagePanicHandler(func(p any) error { - err := fmt.Errorf("panic handled: %v", p) - return err - }), - ) - p.Add( - pipe.Print("hello world"), - pipe.Function( - "farewell", - func(_ context.Context, _ pipe.Env, _ io.Reader, _ io.Writer) error { - panic("this is a panic") - }, - ), - ) - - out, err := p.Output(ctx) - assert.ErrorContains(t, err, "panic handled") - assert.Empty(t, out) - }) } func TestPipelineWithFunction(t *testing.T) { @@ -488,10 +495,23 @@ func (s ErrorStartingStage) Name() string { return "errorStartingStage" } +func (s ErrorStartingStage) Preferences() pipe.StagePreferences { + return pipe.StagePreferences{ + StdinPreference: pipe.IOPreferenceUndefined, + StdoutPreference: pipe.IOPreferenceUndefined, + } +} + func (s ErrorStartingStage) Start( - _ context.Context, _ pipe.Env, _ io.ReadCloser, -) (io.ReadCloser, error) { - return io.NopCloser(&bytes.Buffer{}), s.err + _ context.Context, _ pipe.Env, stdin io.ReadCloser, stdout io.WriteCloser, +) error { + if stdin != nil { + _ = stdin.Close() + } + if stdout != nil { + _ = stdout.Close() + } + return s.err } func (s ErrorStartingStage) Wait() error { diff --git a/pipe/stage.go b/pipe/stage.go index f3d74d9..8b6dff5 100644 --- a/pipe/stage.go +++ b/pipe/stage.go @@ -5,30 +5,142 @@ import ( "io" ) -// Stage is an element of a `Pipeline`. +// Stage is an element of a `Pipeline`. It reads from standard input +// and writes to standard output. +// +// Who closes stdin and stdout? +// +// A `Stage` as a whole needs to be 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". +// +// How this should be done depends on whether stdin/stdout are of type +// `*os.File`. +// +// If a stage is an external command, then the subprocess ultimately +// needs its own copies of `*os.File` file descriptors for its stdin +// and stdout. The external command will "always" [1] close those when +// it exits. +// +// If the stage is an external command and one of the arguments is an +// `*os.File`, then it can set the corresponding field of `exec.Cmd` +// to that argument directly. This has the result that `exec.Cmd` +// duplicates that file descriptor and passes the dup to the +// subprocess. Therefore, the stage must close its copy of that +// argument as soon as the external command has started, because the +// external command will keep its own copy open as long as necessary +// (and no longer!). It should use roughly the following sequence: +// +// cmd.Stdin = f // Similarly for stdout +// cmd.Start(…) +// 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 +// finished, like so: +// +// cmd.Stdin = r // Similarly for stdout +// cmd.Start(…) +// cmd.Wait() +// r.Close() +// +// 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: +// +// go func() { +// f(…, stdin, stdout) +// stdin.Close() +// stdout.Close() +// }() +// +// 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, it's +// advantageous for the first stage should ideally unwrap its stdin +// argument 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. +// +// 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. +// +// [1] It's theoretically possible for a command to pass the open file +// descriptor to another, longer-lived process, in which case the +// file descriptor wouldn't necessarily get closed when the +// command finishes. But that's ill-behaved in a command that is +// being used in a pipeline, so we'll ignore that possibility. + type Stage interface { // Name returns the name of the stage. Name() string + // Preferences() returns this stage's preferences regarding how it + // should be run. + Preferences() StagePreferences + // Start starts the stage in the background, in the environment - // described by `env`, and using `stdin` as input. (`stdin` should - // be set to `nil` if the stage is to receive no input, which - // might be the case for the first stage in a pipeline.) It - // returns an `io.ReadCloser` from which the stage's output can be - // read (or `nil` if it generates no output, which should only be - // the case for the last stage in a pipeline). It is the stages' - // responsibility to close `stdin` (if it is not nil) when it has - // read all of the input that it needs, and to close the write end - // of its output reader when it is done, as that is generally how - // the subsequent stage knows that it has received all of its - // input and can finish its work, too. + // described by `env`, using `stdin` to provide its input and + // `stdout` to collect its output. (`stdin`/`stdout` might be set + // to `nil` if the stage is to receive no input, which might be + // the case for the first/last stage in a pipeline.) See the + // `Stage` type comment for more information about responsibility + // for closing stdin and stdout. // // If `Start()` returns without an error, `Wait()` must also be // called, to allow all resources to be freed. - Start(ctx context.Context, env Env, stdin io.ReadCloser) (io.ReadCloser, error) + Start(ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser) error // Wait waits for the stage to be done, either because it has // finished or because it has been killed due to the expiration of // the context passed to `Start()`. Wait() error } + +// StagePreferences is the way that a `Stage` indicates its +// preferences about how it is run. This is used within +// `pipe.Pipeline` to decide when to use `os.Pipe()` vs. `io.Pipe()` +// for creating the pipes between stages. +type StagePreferences struct { + StdinPreference IOPreference + StdoutPreference IOPreference +} + +// IOPreference describes what type of stdin / stdout a stage would +// prefer. +// +// External commands prefer `*os.File`s (such as those produced by +// `os.Pipe()`) as their stdin and stdout, because those can be passed +// directly by the external process without any extra copying and also +// simplify the semantics around process termination. Go function +// stages are typically happy with any `io.ReadCloser` (such as one +// produced by `io.Pipe()`), which can be more efficient because +// traffic through an `io.Pipe()` happens entirely in userspace. +type IOPreference int + +const ( + // IOPreferenceUndefined indicates that the stage doesn't care + // what form the specified stdin / stdout takes (i.e., any old + // `io.ReadCloser` / `io.WriteCloser` is just fine). + IOPreferenceUndefined IOPreference = iota + + // IOPreferenceFile indicates that the stage would prefer for the + // specified stdin / stdout to be an `*os.File`, to avoid copying. + IOPreferenceFile + + // IOPreferenceNil indicates that the stage does not use the + // specified stdin / stdout, so `nil` should be passed in. This + // should only happen at the beginning / end of a pipeline. + IOPreferenceNil +) From cc9cd67da05479151291bb8e6bbef8efedaf9461 Mon Sep 17 00:00:00 2001 From: Michael Haggerty Date: Sun, 17 Dec 2023 13:42:27 +0100 Subject: [PATCH 06/23] Add some tests that `Pipeline.Start()` picks the right stdin/stdout The most complicated code dealing with the change to `Stage.Start()` is the selection of which types of stdin/stderr to pass to stages, and that's also the main advantage of the new interface. So add a bunch of tests that the correct types (especially, `io.Pipe()` vs. `os.Pipe()`) are indeed being selected. --- pipe/export_test.go | 4 + pipe/nop_closer.go | 16 ++ pipe/pipe_matching_test.go | 376 +++++++++++++++++++++++++++++++++++++ 3 files changed, 396 insertions(+) create mode 100644 pipe/export_test.go create mode 100644 pipe/pipe_matching_test.go diff --git a/pipe/export_test.go b/pipe/export_test.go new file mode 100644 index 0000000..92862cc --- /dev/null +++ b/pipe/export_test.go @@ -0,0 +1,4 @@ +package pipe + +// This file exports a function to be used only for testing. +var UnwrapNopCloser = unwrapNopCloser diff --git a/pipe/nop_closer.go b/pipe/nop_closer.go index 8c66c72..18cf7a9 100644 --- a/pipe/nop_closer.go +++ b/pipe/nop_closer.go @@ -51,3 +51,19 @@ type writerNopCloser struct { func (w writerNopCloser) Close() error { return nil } + +// unwrapNopCloser unwraps the object if it is some kind of nop +// closer, and returns the underlying object. This function is used +// only for testing. +func unwrapNopCloser(obj any) (any, bool) { + switch obj := obj.(type) { + case readerNopCloser: + return obj.Reader, true + case readerWriterToNopCloser: + return obj.Reader, true + case writerNopCloser: + return obj.Writer, true + default: + return nil, false + } +} diff --git a/pipe/pipe_matching_test.go b/pipe/pipe_matching_test.go new file mode 100644 index 0000000..5528ba2 --- /dev/null +++ b/pipe/pipe_matching_test.go @@ -0,0 +1,376 @@ +package pipe_test + +import ( + "context" + "fmt" + "io" + "os" + "testing" + + "github.com/github/go-pipe/pipe" + "github.com/stretchr/testify/assert" +) + +// Tests that `Pipeline.Start()` uses the correct types of pipes in +// various situations. +// +// The type of pipe to use depends on both the source and the consumer +// of the data, including the overall pipeline's stdin and stdout. So +// there are a lot of possibilities to consider. + +// Additional values used for the expected types of stdin/stdout: +const ( + IOPreferenceUndefinedNopCloser pipe.IOPreference = iota + 100 + IOPreferenceFileNopCloser +) + +func file(t *testing.T) *os.File { + f, err := os.Open(os.DevNull) + assert.NoError(t, err) + return f +} + +func readCloser() io.ReadCloser { + r, w := io.Pipe() + w.Close() + return r +} + +func writeCloser() io.WriteCloser { + r, w := io.Pipe() + r.Close() + return w +} + +func newPipeSniffingStage( + stdinPreference, stdinExpectation pipe.IOPreference, + stdoutPreference, stdoutExpectation pipe.IOPreference, +) *pipeSniffingStage { + return &pipeSniffingStage{ + prefs: pipe.StagePreferences{ + StdinPreference: stdinPreference, + StdoutPreference: stdoutPreference, + }, + expect: pipe.StagePreferences{ + StdinPreference: stdinExpectation, + StdoutPreference: stdoutExpectation, + }, + } +} + +func newPipeSniffingFunc( + stdinExpectation, stdoutExpectation pipe.IOPreference, +) *pipeSniffingStage { + return newPipeSniffingStage( + pipe.IOPreferenceUndefined, stdinExpectation, + pipe.IOPreferenceUndefined, stdoutExpectation, + ) +} + +func newPipeSniffingCmd( + stdinExpectation, stdoutExpectation pipe.IOPreference, +) *pipeSniffingStage { + return newPipeSniffingStage( + pipe.IOPreferenceFile, stdinExpectation, + pipe.IOPreferenceFile, stdoutExpectation, + ) +} + +type pipeSniffingStage struct { + prefs pipe.StagePreferences + expect pipe.StagePreferences + stdin io.ReadCloser + stdout io.WriteCloser +} + +func (*pipeSniffingStage) Name() string { + return "pipe-sniffer" +} + +func (s *pipeSniffingStage) Preferences() pipe.StagePreferences { + return s.prefs +} + +func (s *pipeSniffingStage) Start( + _ context.Context, _ pipe.Env, stdin io.ReadCloser, stdout io.WriteCloser, +) error { + s.stdin = stdin + if stdin != nil { + _ = stdin.Close() + } + s.stdout = stdout + if stdout != nil { + _ = stdout.Close() + } + return nil +} + +func (s *pipeSniffingStage) check(t *testing.T, i int) { + t.Helper() + + checkStdinExpectation(t, i, s.expect.StdinPreference, s.stdin) + checkStdoutExpectation(t, i, s.expect.StdoutPreference, s.stdout) +} + +func (s *pipeSniffingStage) Wait() error { + return nil +} + +var _ pipe.Stage = (*pipeSniffingStage)(nil) + +func ioTypeString(f any) string { + if f == nil { + return "nil" + } + if f, ok := pipe.UnwrapNopCloser(f); ok { + return fmt.Sprintf("nopCloser(%s)", ioTypeString(f)) + } + switch f := f.(type) { + case *os.File: + return "*os.File" + case io.Reader: + return "other" + case io.Writer: + return "other" + default: + return fmt.Sprintf("%T", f) + } +} + +func prefString(pref pipe.IOPreference) string { + switch pref { + case pipe.IOPreferenceUndefined: + return "other" + case pipe.IOPreferenceFile: + return "*os.File" + case pipe.IOPreferenceNil: + return "nil" + case IOPreferenceUndefinedNopCloser: + return "nopCloser(other)" + case IOPreferenceFileNopCloser: + return "nopCloser(*os.File)" + default: + panic(fmt.Sprintf("invalid IOPreference: %d", pref)) + } +} + +func checkStdinExpectation(t *testing.T, i int, pref pipe.IOPreference, stdin io.ReadCloser) { + t.Helper() + + ioType := ioTypeString(stdin) + expType := prefString(pref) + assert.Equalf( + t, expType, ioType, + "stage %d stdin: expected %s, got %s (%T)", i, expType, ioType, stdin, + ) +} + +type WriterNopCloser interface { + NopCloserWriter() io.Writer +} + +func checkStdoutExpectation(t *testing.T, i int, pref pipe.IOPreference, stdout io.WriteCloser) { + t.Helper() + + ioType := ioTypeString(stdout) + expType := prefString(pref) + assert.Equalf( + t, expType, ioType, + "stage %d stdout: expected %s, got %s (%T)", i, expType, ioType, stdout, + ) +} + +type checker interface { + check(t *testing.T, i int) +} + +func TestPipeTypes(t *testing.T) { + ctx := context.Background() + + t.Parallel() + + for _, tc := range []struct { + name string + opts []pipe.Option + stages []pipe.Stage + stdin io.Reader + stdout io.Writer + }{ + { + name: "func", + opts: []pipe.Option{}, + stages: []pipe.Stage{ + newPipeSniffingFunc(pipe.IOPreferenceNil, pipe.IOPreferenceNil), + }, + }, + { + name: "func-file-stdin", + opts: []pipe.Option{ + pipe.WithStdin(file(t)), + }, + stages: []pipe.Stage{ + newPipeSniffingFunc(IOPreferenceFileNopCloser, pipe.IOPreferenceNil), + }, + }, + { + name: "func-file-stdout", + opts: []pipe.Option{ + pipe.WithStdout(file(t)), + }, + stages: []pipe.Stage{ + newPipeSniffingFunc(pipe.IOPreferenceNil, IOPreferenceFileNopCloser), + }, + }, + { + name: "func-file-stdout-closer", + opts: []pipe.Option{ + pipe.WithStdoutCloser(file(t)), + }, + stages: []pipe.Stage{ + newPipeSniffingFunc(pipe.IOPreferenceNil, pipe.IOPreferenceFile), + }, + }, + { + name: "func-file-stdin-other-stdout-closer-other", + opts: []pipe.Option{ + pipe.WithStdin(readCloser()), + pipe.WithStdoutCloser(writeCloser()), + }, + stages: []pipe.Stage{ + newPipeSniffingFunc(IOPreferenceUndefinedNopCloser, pipe.IOPreferenceUndefined), + }, + }, + { + name: "cmd", + opts: []pipe.Option{}, + stages: []pipe.Stage{ + newPipeSniffingCmd(pipe.IOPreferenceNil, pipe.IOPreferenceNil), + }, + }, + { + name: "cmd-file-stdin", + opts: []pipe.Option{ + pipe.WithStdin(file(t)), + }, + stages: []pipe.Stage{ + newPipeSniffingCmd(IOPreferenceFileNopCloser, pipe.IOPreferenceNil), + }, + }, + { + name: "cmd-file-stdout", + opts: []pipe.Option{ + pipe.WithStdout(file(t)), + }, + stages: []pipe.Stage{ + newPipeSniffingCmd(pipe.IOPreferenceNil, IOPreferenceFileNopCloser), + }, + }, + { + name: "cmd-file-stdout-closer", + opts: []pipe.Option{ + pipe.WithStdoutCloser(file(t)), + }, + stages: []pipe.Stage{ + newPipeSniffingCmd(pipe.IOPreferenceNil, pipe.IOPreferenceFile), + }, + }, + { + name: "cmd-file-stdin-other-stdout-closer-other", + opts: []pipe.Option{ + pipe.WithStdin(readCloser()), + pipe.WithStdoutCloser(writeCloser()), + }, + stages: []pipe.Stage{ + newPipeSniffingCmd(IOPreferenceUndefinedNopCloser, pipe.IOPreferenceUndefined), + }, + }, + { + name: "func-func", + opts: []pipe.Option{ + pipe.WithStdin(file(t)), + pipe.WithStdoutCloser(writeCloser()), + }, + stages: []pipe.Stage{ + newPipeSniffingFunc(IOPreferenceFileNopCloser, pipe.IOPreferenceUndefined), + newPipeSniffingFunc(pipe.IOPreferenceUndefined, pipe.IOPreferenceUndefined), + }, + }, + { + name: "func-cmd", + opts: []pipe.Option{ + pipe.WithStdout(file(t)), + }, + stages: []pipe.Stage{ + newPipeSniffingFunc(pipe.IOPreferenceNil, pipe.IOPreferenceFile), + newPipeSniffingCmd(pipe.IOPreferenceFile, IOPreferenceFileNopCloser), + }, + }, + { + name: "cmd-func", + opts: []pipe.Option{ + pipe.WithStdin(readCloser()), + }, + stages: []pipe.Stage{ + newPipeSniffingCmd(IOPreferenceUndefinedNopCloser, pipe.IOPreferenceFile), + newPipeSniffingFunc(pipe.IOPreferenceFile, pipe.IOPreferenceNil), + }, + }, + { + name: "cmd-cmd", + opts: []pipe.Option{}, + stages: []pipe.Stage{ + newPipeSniffingCmd(pipe.IOPreferenceNil, pipe.IOPreferenceFile), + newPipeSniffingCmd(pipe.IOPreferenceFile, pipe.IOPreferenceNil), + }, + }, + { + name: "hybrid1", + opts: []pipe.Option{}, + stages: []pipe.Stage{ + newPipeSniffingStage( + pipe.IOPreferenceUndefined, pipe.IOPreferenceNil, + pipe.IOPreferenceUndefined, pipe.IOPreferenceUndefined, + ), + newPipeSniffingStage( + pipe.IOPreferenceUndefined, pipe.IOPreferenceUndefined, + pipe.IOPreferenceFile, pipe.IOPreferenceFile, + ), + newPipeSniffingStage( + pipe.IOPreferenceUndefined, pipe.IOPreferenceFile, + pipe.IOPreferenceUndefined, pipe.IOPreferenceNil, + ), + }, + }, + { + name: "hybrid2", + opts: []pipe.Option{}, + stages: []pipe.Stage{ + newPipeSniffingStage( + pipe.IOPreferenceUndefined, pipe.IOPreferenceNil, + pipe.IOPreferenceUndefined, pipe.IOPreferenceFile, + ), + newPipeSniffingStage( + pipe.IOPreferenceFile, pipe.IOPreferenceFile, + pipe.IOPreferenceUndefined, pipe.IOPreferenceUndefined, + ), + newPipeSniffingStage( + pipe.IOPreferenceUndefined, pipe.IOPreferenceUndefined, + pipe.IOPreferenceUndefined, pipe.IOPreferenceNil, + ), + }, + }, + } { + tc := tc + + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + p := pipe.New(tc.opts...) + p.Add(tc.stages...) + assert.NoError(t, p.Run(ctx)) + for i, s := range tc.stages { + s.(checker).check(t, i) + } + }) + } +} From 6c6dfe53ab3731193f0d464051f95176d21f2107 Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Wed, 8 Apr 2026 19:45:51 +0200 Subject: [PATCH 07/23] Port MemoryLimitWithObserver to new Stage interface MemoryLimitWithObserver was added to main (PR #48) after the version-2 branch diverged. Port it to the new Stage interface and add test coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pipe/memorylimit.go | 93 ++++++++++++++++++++++++++++++++ pipe/memorylimit_test.go | 111 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 204 insertions(+) diff --git a/pipe/memorylimit.go b/pipe/memorylimit.go index 48e1bb6..66eac30 100644 --- a/pipe/memorylimit.go +++ b/pipe/memorylimit.go @@ -89,6 +89,99 @@ func killAtLimit(byteLimit uint64, eventHandler func(e *Event)) memoryWatchFunc } } +// MemoryLimitWithObserver combines MemoryLimit and MemoryObserver in +// one goroutine. It watches the memory usage of the stage, stops it +// if it exceeds the given limit, and logs the peak memory usage when +// the stage exits. +func MemoryLimitWithObserver(stage Stage, byteLimit uint64, eventHandler func(e *Event)) Stage { + limitableStage, ok := stage.(LimitableStage) + if !ok { + eventHandler(&Event{ + Command: stage.Name(), + Msg: "invalid pipe.MemoryLimitWithObserver usage", + Err: fmt.Errorf("invalid pipe.MemoryLimitWithObserver usage"), + }) + return stage + } + + return &memoryWatchStage{ + nameSuffix: " with memory limit", + stage: limitableStage, + watch: killAtLimitAndObserve(byteLimit, eventHandler), + } +} + +func killAtLimitAndObserve(byteLimit uint64, eventHandler func(e *Event)) memoryWatchFunc { + return func(ctx context.Context, stage LimitableStage) { + var ( + maxRSS uint64 + samples, errCount, consecutiveErrors int + killed bool + ) + + t := time.NewTicker(memoryPollInterval) + defer t.Stop() + + for { + select { + case <-ctx.Done(): + eventHandler(&Event{ + Command: stage.Name(), + Msg: "peak memory usage", + Context: map[string]interface{}{ + "max_rss_bytes": maxRSS, + "samples": samples, + "errors": errCount, + }, + }) + return + case <-t.C: + if killed { + continue + } + + rss, err := stage.GetRSSAnon(ctx) + if err != nil { + if !errors.Is(err, errProcessInfoMissing) { + errCount++ + consecutiveErrors++ + if consecutiveErrors == 2 { + eventHandler(&Event{ + Command: stage.Name(), + Msg: "error getting RSS", + Err: err, + }) + } + } else { + consecutiveErrors = 0 + } + continue + } + + consecutiveErrors = 0 + samples++ + if rss > maxRSS { + maxRSS = rss + } + + if rss >= byteLimit { + eventHandler(&Event{ + Command: stage.Name(), + Msg: "stage exceeded allowed memory use", + Err: fmt.Errorf("stage exceeded allowed memory use"), + Context: map[string]interface{}{ + "limit": byteLimit, + "used": rss, + }, + }) + stage.Kill(ErrMemoryLimitExceeded) + killed = true + } + } + } + } +} + // MemoryObserver watches memory use of the stage and logs the maximum // value when the stage exits. func MemoryObserver(stage Stage, eventHandler func(e *Event)) Stage { diff --git a/pipe/memorylimit_test.go b/pipe/memorylimit_test.go index e84f465..27a40cb 100644 --- a/pipe/memorylimit_test.go +++ b/pipe/memorylimit_test.go @@ -113,6 +113,84 @@ func TestMemoryLimitTreeMem(t *testing.T) { require.ErrorContains(t, err, "memory limit exceeded") } +func TestMemoryLimitWithObserverSimple(t *testing.T) { + t.Parallel() + msg, err := testMemoryLimitWithObserver(t, 400, 10_000_000, pipe.Command("less")) + assert.Contains(t, msg, "exceeded allowed memory") + assert.Contains(t, msg, "limit=10000000") + require.ErrorContains(t, err, "memory limit exceeded") +} + +func TestMemoryLimitWithObserverTreeMem(t *testing.T) { + t.Parallel() + msg, err := testMemoryLimitWithObserver(t, 400, 10_000_000, pipe.Command("sh", "-c", "less; :")) + assert.Contains(t, msg, "exceeded allowed memory") + assert.Contains(t, msg, "limit=10000000") + require.ErrorContains(t, err, "memory limit exceeded") +} + +func TestMemoryLimitWithObserverLogsPeakOnKill(t *testing.T) { + t.Parallel() + msg, err := testMemoryLimitWithObserver(t, 400, 10_000_000, pipe.Command("less")) + assert.Contains(t, msg, "exceeded allowed memory") + assert.Contains(t, msg, "peak memory usage") + require.ErrorContains(t, err, "memory limit exceeded") +} + +func TestMemoryLimitWithObserverBelowLimit(t *testing.T) { + t.Parallel() + rss := testMemoryLimitWithObserverBelowLimit(t, 400, pipe.Command("less")) + require.Greater(t, rss, 400_000_000) +} + +func TestMemoryLimitWithObserverBelowLimitTreeMem(t *testing.T) { + t.Parallel() + rss := testMemoryLimitWithObserverBelowLimit(t, 400, pipe.Command("sh", "-c", "less; :")) + require.Greater(t, rss, 400_000_000) +} + +// testMemoryLimitWithObserverBelowLimit exercises the observer half of +// `MemoryLimitWithObserver` when the memory limit is never hit: with a +// 100GiB limit, less should never be killed, but the wrapper should +// still poll RSS and emit a "peak memory usage" event when the stage +// exits normally. Mirrors `testMemoryObserver` in structure — we hold +// stdin open across at least one poll interval so RSS samples are +// guaranteed to be taken before the stage is allowed to exit. +func testMemoryLimitWithObserverBelowLimit(t *testing.T, mbs int, stage pipe.Stage) int { + ctx := context.Background() + + stdinReader, stdinWriter := io.Pipe() + + devNull, err := os.OpenFile("/dev/null", os.O_WRONLY, 0) + require.NoError(t, err) + + buf := &bytes.Buffer{} + logger := log.New(buf, "testMemoryLimitWithObserverBelowLimit", log.Ldate|log.Ltime) + + p := pipe.New(pipe.WithDir("/"), pipe.WithStdin(stdinReader), pipe.WithStdout(devNull)) + p.Add(pipe.MemoryLimitWithObserver(stage, 100*1024*1024*1024, LogEventHandler(logger))) + require.NoError(t, p.Start(ctx)) + + var bytes [1_000_000]byte + for i := 0; i < mbs; i++ { + n, err := stdinWriter.Write(bytes[:]) + require.NoError(t, err) + require.Equal(t, len(bytes), n) + } + + // Wrapper polls once per second; sleep long enough to guarantee at + // least one sample is taken before the stage is allowed to exit. + time.Sleep(2 * time.Second) + + require.NoError(t, stdinWriter.Close()) + require.NoError(t, p.Wait()) + + output := buf.String() + assert.Contains(t, output, "peak memory usage") + + return maxBytes(output) +} + func testMemoryLimit(t *testing.T, mbs int, limit uint64, stage pipe.Stage) (string, error) { ctx := context.Background() @@ -147,3 +225,36 @@ func testMemoryLimit(t *testing.T, mbs int, limit uint64, stage pipe.Stage) (str return buf.String(), err } + +func testMemoryLimitWithObserver(t *testing.T, mbs int, limit uint64, stage pipe.Stage) (string, error) { + ctx := context.Background() + + devNull, err := os.OpenFile("/dev/null", os.O_WRONLY, 0) + require.NoError(t, err) + + buf := &bytes.Buffer{} + logger := log.New(buf, "testMemoryLimitWithObserver", log.Ldate|log.Ltime) + + p := pipe.New(pipe.WithDir("/"), pipe.WithStdoutCloser(devNull)) + p.Add( + pipe.Function( + "write-to-less", + func(ctx context.Context, _ pipe.Env, _ io.Reader, stdout io.Writer) error { + var bytes [1_000_000]byte + for i := 0; i < mbs; i++ { + _, err := stdout.Write(bytes[:]) + if err != nil { + require.ErrorIs(t, err, syscall.EPIPE) + } + } + return nil + }, + ), + pipe.MemoryLimitWithObserver(stage, limit, LogEventHandler(logger)), + ) + require.NoError(t, p.Start(ctx)) + + err = p.Wait() + + return buf.String(), err +} From cd47557634d060123ec4d4fb07ae03994ef843bc Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Wed, 8 Apr 2026 19:47:04 +0200 Subject: [PATCH 08/23] Restore panic handler for Function stages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The version-2 Stage interface redesign dropped the panic recovery from goStage (or preceded its addition). Add/restore it: add SetPanicHandler(), recoverPanic(), and WithStagePanicHandler pipeline option. The goroutine uses stacked defers (close(done) → close stdin → close stdout → recoverPanic) so that when a Function panics, recoverPanic fires first (sets s.err), then cleanup runs, then done closes — allowing Wait() to return the caught panic error. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pipe/function.go | 53 +++++++++++++++++++++++++++++-------------- pipe/pipeline.go | 18 +++++++++++++++ pipe/pipeline_test.go | 21 +++++++++++++++++ 3 files changed, 75 insertions(+), 17 deletions(-) diff --git a/pipe/function.go b/pipe/function.go index fe8abd0..958ee4b 100644 --- a/pipe/function.go +++ b/pipe/function.go @@ -32,20 +32,26 @@ 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 + name string + f StageFunc + done chan struct{} + err error + panicHandler StagePanicHandler } var ( - _ Stage = (*goStage)(nil) + _ Stage = (*goStage)(nil) + _ StagePanicHandlerAware = (*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, @@ -67,21 +73,24 @@ func (s *goStage) Start( } go func() { - 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) + defer close(s.done) + defer func() { + 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) + } } - } - - 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) + }() + defer func() { + 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) + } } - } + }() + defer s.recoverPanic() - close(s.done) + s.err = s.f(ctx, env, r, w) }() return nil @@ -91,3 +100,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) + } +} diff --git a/pipe/pipeline.go b/pipe/pipeline.go index 4fdd4a5..2206206 100644 --- a/pipe/pipeline.go +++ b/pipe/pipeline.go @@ -65,6 +65,7 @@ type Pipeline struct { started uint32 eventHandler func(e *Event) + panicHandler StagePanicHandler } var emptyEventHandler = func(e *Event) {} @@ -216,6 +217,15 @@ func WithEventHandler(handler func(e *Event)) Option { } } +// WithStagePanicHandler sets a panic handler for the stages within a pipeline. +// When a pipeline stage panics, the provided handler will be invoked, allowing +// the client to handle the panic in whatever way they see fit. +func WithStagePanicHandler(ph StagePanicHandler) Option { + return func(p *Pipeline) { + p.panicHandler = ph + } +} + func (p *Pipeline) hasStarted() bool { return atomic.LoadUint32(&p.started) != 0 } @@ -315,6 +325,10 @@ func (p *Pipeline) Start(ctx context.Context) error { ss := &stageStarters[i] nextSS := &stageStarters[i+1] + if phs, ok := s.(StagePanicHandlerAware); ok && p.panicHandler != nil { + phs.SetPanicHandler(p.panicHandler) + } + // We need to generate a pipe pair for this stage to use // to communicate with its successor: if ss.prefs.StdoutPreference == IOPreferenceFile || @@ -343,6 +357,10 @@ func (p *Pipeline) Start(ctx context.Context) error { s := p.stages[i] ss := &stageStarters[i] + if phs, ok := s.(StagePanicHandlerAware); ok && p.panicHandler != nil { + phs.SetPanicHandler(p.panicHandler) + } + if err := s.Start(ctx, p.env, ss.stdin, ss.stdout); err != nil { return abort(i, err) } diff --git a/pipe/pipeline_test.go b/pipe/pipeline_test.go index 0a5d17d..e73884b 100644 --- a/pipe/pipeline_test.go +++ b/pipe/pipeline_test.go @@ -457,6 +457,27 @@ func TestFunction(t *testing.T) { assert.NoError(t, err) assert.EqualValues(t, "goodbye, cruel world", out) }) + + t.Run("panic with handler", func(t *testing.T) { + p := pipe.New( + pipe.WithStagePanicHandler(func(p any) error { + return fmt.Errorf("panic handled: %v", p) + }), + ) + p.Add( + pipe.Print("hello world"), + pipe.Function( + "farewell", + func(_ context.Context, _ pipe.Env, _ io.Reader, _ io.Writer) error { + panic("this is a panic") + }, + ), + ) + + out, err := p.Output(ctx) + assert.ErrorContains(t, err, "panic handled") + assert.Empty(t, out) + }) } func TestPipelineWithFunction(t *testing.T) { From 48dffbd2ac53f5e211a6e0d2bcad656fdd70fd22 Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Wed, 8 Apr 2026 19:48:29 +0200 Subject: [PATCH 09/23] Fix memoryWatchStage.Wait() to always call stopWatching() Wait() returned the inner stage's error immediately without calling stopWatching(), which meant the memory watcher goroutine was never cancelled when the stage exited with an error (e.g., from being killed due to memory limit). This prevented the observer from logging peak memory usage on kill. Fix: always call stopWatching() before returning, regardless of whether the inner stage returned an error. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pipe/memorylimit.go | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/pipe/memorylimit.go b/pipe/memorylimit.go index 66eac30..a863574 100644 --- a/pipe/memorylimit.go +++ b/pipe/memorylimit.go @@ -297,11 +297,8 @@ func (m *memoryWatchStage) monitor(ctx context.Context) { } func (m *memoryWatchStage) Wait() error { - if err := m.stage.Wait(); err != nil { - return err - } - m.stopWatching() - return nil + defer m.stopWatching() + return m.stage.Wait() } func (m *memoryWatchStage) GetRSSAnon(ctx context.Context) (uint64, error) { From 92ad57c05ce308c4766832d45f36465eaf769930 Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Thu, 9 Apr 2026 14:13:03 +0200 Subject: [PATCH 10/23] Fix lint errors - Remove Go 1.22+ unnecessary loop variable copy (copyloopvar) - Replace unused parameters with _ (revive) - Add nolint directive for FinishEarly naming (staticcheck ST1012) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pipe/memorylimit_test.go | 4 ++-- pipe/pipe_matching_test.go | 2 -- pipe/pipeline.go | 3 ++- pipe/pipeline_test.go | 8 ++++---- 4 files changed, 8 insertions(+), 9 deletions(-) diff --git a/pipe/memorylimit_test.go b/pipe/memorylimit_test.go index 27a40cb..17550e1 100644 --- a/pipe/memorylimit_test.go +++ b/pipe/memorylimit_test.go @@ -204,7 +204,7 @@ func testMemoryLimit(t *testing.T, mbs int, limit uint64, stage pipe.Stage) (str p.Add( pipe.Function( "write-to-less", - func(ctx context.Context, _ pipe.Env, _ io.Reader, stdout io.Writer) error { + func(_ context.Context, _ pipe.Env, _ io.Reader, stdout io.Writer) error { // Write some nonsense data to less. var bytes [1_000_000]byte for i := 0; i < mbs; i++ { @@ -239,7 +239,7 @@ func testMemoryLimitWithObserver(t *testing.T, mbs int, limit uint64, stage pipe p.Add( pipe.Function( "write-to-less", - func(ctx context.Context, _ pipe.Env, _ io.Reader, stdout io.Writer) error { + func(_ context.Context, _ pipe.Env, _ io.Reader, stdout io.Writer) error { var bytes [1_000_000]byte for i := 0; i < mbs; i++ { _, err := stdout.Write(bytes[:]) diff --git a/pipe/pipe_matching_test.go b/pipe/pipe_matching_test.go index 5528ba2..ace6103 100644 --- a/pipe/pipe_matching_test.go +++ b/pipe/pipe_matching_test.go @@ -360,8 +360,6 @@ func TestPipeTypes(t *testing.T) { }, }, } { - tc := tc - t.Run(tc.name, func(t *testing.T) { t.Parallel() diff --git a/pipe/pipeline.go b/pipe/pipeline.go index 2206206..9039232 100644 --- a/pipe/pipeline.go +++ b/pipe/pipeline.go @@ -29,6 +29,7 @@ type Env struct { // and is not reported to the caller. // //revive:disable:error-naming +//nolint:staticcheck // ST1012: FinishEarly is the intentional name for this sentinel error var FinishEarly = errors.New("finish stage early") //revive:enable:error-naming @@ -68,7 +69,7 @@ type Pipeline struct { panicHandler StagePanicHandler } -var emptyEventHandler = func(e *Event) {} +var emptyEventHandler = func(_ *Event) {} type NewPipeFn func(opts ...Option) *Pipeline diff --git a/pipe/pipeline_test.go b/pipe/pipeline_test.go index e73884b..f5eccfb 100644 --- a/pipe/pipeline_test.go +++ b/pipe/pipeline_test.go @@ -1001,7 +1001,7 @@ func BenchmarkMoreDataUnbuffered(b *testing.B) { p.Add( pipe.Function( "seq", - func(ctx context.Context, _ pipe.Env, stdin io.Reader, stdout io.Writer) error { + func(_ context.Context, _ pipe.Env, _ io.Reader, stdout io.Writer) error { for i := 1; i <= 100000; i++ { fmt.Fprintln(stdout, i) } @@ -1019,7 +1019,7 @@ func BenchmarkMoreDataUnbuffered(b *testing.B) { pipe.Command("cat"), pipe.LinewiseFunction( "count", - func(ctx context.Context, _ pipe.Env, line []byte, stdout *bufio.Writer) error { + func(_ context.Context, _ pipe.Env, _ []byte, _ *bufio.Writer) error { count++ return nil }, @@ -1046,7 +1046,7 @@ func BenchmarkMoreDataBuffered(b *testing.B) { p.Add( pipe.Function( "seq", - func(ctx context.Context, _ pipe.Env, stdin io.Reader, stdout io.Writer) error { + func(_ context.Context, _ pipe.Env, _ io.Reader, stdout io.Writer) error { out := bufio.NewWriter(stdout) for i := 1; i <= 1000000; i++ { fmt.Fprintln(out, i) @@ -1065,7 +1065,7 @@ func BenchmarkMoreDataBuffered(b *testing.B) { pipe.Command("cat"), pipe.LinewiseFunction( "count", - func(ctx context.Context, _ pipe.Env, line []byte, stdout *bufio.Writer) error { + func(_ context.Context, _ pipe.Env, _ []byte, _ *bufio.Writer) error { count++ return nil }, From 6a0abf3e32df35974dc7f0bcf7ae4a9506348966 Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Wed, 8 Apr 2026 19:47:44 +0200 Subject: [PATCH 11/23] Restore identity-copy behavior for empty pipelines Prior to the Stage interface redesign, an empty pipeline with a configured `stdout` ran a synthetic ioCopier that copied `p.stdin` (if any) to `p.stdout` and closed the destination if it came from `WithStdoutCloser()`. Restore that behavior by synthesizing an identity-copy Function stage when the pipeline has no stages but does have a configured output. The empty/no-output case remains a no-op as before. This affects callers like `pipe.New(WithStdin(r)).Output(ctx)`. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pipe/pipeline.go | 22 ++++++++++++++++++ pipe/pipeline_test.go | 54 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/pipe/pipeline.go b/pipe/pipeline.go index 9039232..d160857 100644 --- a/pipe/pipeline.go +++ b/pipe/pipeline.go @@ -269,6 +269,28 @@ func (p *Pipeline) Start(ctx context.Context) error { atomic.StoreUint32(&p.started, 1) ctx, p.cancel = context.WithCancel(ctx) + if len(p.stages) == 0 { + if p.stdout == nil { + // No stages and no destination: there is nothing to do + // and nowhere to put `p.stdin` even if it was set. + return nil + } + // No stages but a destination was configured: synthesize an + // identity-copy stage so that `WithStdin()` is drained into + // `WithStdout()`/`WithStdoutCloser()` and the destination + // closer (if any) is invoked. + p.stages = append(p.stages, Function( + "identity", + func(_ context.Context, _ Env, stdin io.Reader, stdout io.Writer) error { + if stdin == nil { + return nil + } + _, err := io.Copy(stdout, stdin) + return err + }, + )) + } + // We need to decide how to start the stages, especially what // pipes to use to connect adjacent stages (`os.Pipe()` vs. // `io.Pipe()`) based on the two stages' preferences. diff --git a/pipe/pipeline_test.go b/pipe/pipeline_test.go index f5eccfb..0184012 100644 --- a/pipe/pipeline_test.go +++ b/pipe/pipeline_test.go @@ -26,6 +26,60 @@ func TestMain(m *testing.M) { goleak.VerifyTestMain(m) } +func TestPipelineEmpty(t *testing.T) { + t.Parallel() + p := pipe.New() + assert.NoError(t, p.Run(context.Background())) +} + +func TestPipelineEmptyWithStdinAndStdout(t *testing.T) { + t.Parallel() + ctx := context.Background() + stdout := &bytes.Buffer{} + p := pipe.New( + pipe.WithStdin(strings.NewReader("hello world\n")), + pipe.WithStdout(stdout), + ) + if assert.NoError(t, p.Run(ctx)) { + assert.Equal(t, "hello world\n", stdout.String()) + } +} + +func TestPipelineEmptyOutput(t *testing.T) { + t.Parallel() + ctx := context.Background() + p := pipe.New(pipe.WithStdin(strings.NewReader("hello world\n"))) + out, err := p.Output(ctx) + if assert.NoError(t, err) { + assert.Equal(t, "hello world\n", string(out)) + } +} + +func TestPipelineEmptyWithStdoutCloser(t *testing.T) { + t.Parallel() + ctx := context.Background() + stdout := &closeTrackingWriter{} + p := pipe.New( + pipe.WithStdin(strings.NewReader("hello world\n")), + pipe.WithStdoutCloser(stdout), + ) + if assert.NoError(t, p.Run(ctx)) { + assert.Equal(t, "hello world\n", stdout.buf.String()) + assert.True(t, stdout.closed, "WithStdoutCloser destination should be closed") + } +} + +type closeTrackingWriter struct { + buf bytes.Buffer + closed bool +} + +func (w *closeTrackingWriter) Write(p []byte) (int, error) { return w.buf.Write(p) } +func (w *closeTrackingWriter) Close() error { + w.closed = true + return nil +} + func TestPipelineFirstStageFailsToStart(t *testing.T) { t.Parallel() ctx := context.Background() From 96063689bf579d6e42f2c86e388467a228569155 Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Thu, 28 May 2026 20:00:01 +0200 Subject: [PATCH 12/23] Add tests pinning command stdout fd-pass fast path One of the optimizations enabled by the Stage interface redesign in #21 is that when a commandStage at the end of a pipeline has its stdout pointed at an *os.File, exec.Cmd dup's that fd straight into the child process. The PR called out two benefits: - the Go-side copy stage between subprocess and Pipeline.stdout becomes unnecessary (wasteful goroutine + buffer); - the subprocess can detect when Pipeline.stdout is closed, which the old intermediate pipe hid from it. Until now nothing asserted this contract directly. The pipe-type negotiation tests in pipe_matching_test.go only check that the right kind of pipe was chosen between mock stages; the existing Command + WithStdoutCloser(*os.File) test would still pass even if a future refactor silently wrapped stdout in a Go-side io.Copy. Pin the contract with two tests (one direct, one through Pipeline) that assert s.cmd.Stdout == userProvidedFile after Start() for both WithStdoutCloser(*os.File) and WithStdout(*os.File) (writerNopCloser- wrapped). A regression would now fail loudly with a message that names the optimization. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pipe/command_stdout_fastpath_test.go | 118 +++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 pipe/command_stdout_fastpath_test.go diff --git a/pipe/command_stdout_fastpath_test.go b/pipe/command_stdout_fastpath_test.go new file mode 100644 index 0000000..66ce4fb --- /dev/null +++ b/pipe/command_stdout_fastpath_test.go @@ -0,0 +1,118 @@ +package pipe + +import ( + "context" + "io" + "os" + "os/exec" + "testing" +) + +// TestCommandStageStdoutFastPath asserts that when a commandStage's stdout is +// an `*os.File`, the file is set as `cmd.Stdout` so that `exec.Cmd` dup's the +// fd into the child process directly. This is one of the optimizations enabled +// by the Stage interface redesign in #21: the subprocess writes straight to +// the caller's destination fd with no Go-side copy stage in between, and the +// subprocess can detect when that fd is closed. +func TestCommandStageStdoutFastPath(t *testing.T) { + cases := []struct { + name string + wrap func(*os.File) io.WriteCloser + }{ + { + name: "raw *os.File via WithStdoutCloser", + wrap: func(f *os.File) io.WriteCloser { return f }, + }, + { + name: "writerNopCloser{*os.File} via WithStdout", + wrap: func(f *os.File) io.WriteCloser { return writerNopCloser{f} }, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + f, err := os.CreateTemp(t.TempDir(), "stdout") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = f.Close() }) + + cmd := exec.Command("true") + s := CommandStage("true", cmd).(*commandStage) + + if err := s.Start(ctx, Env{}, nil, tc.wrap(f)); err != nil { + t.Fatalf("Start: %v", err) + } + t.Cleanup(func() { _ = s.Wait() }) + + gotFile, ok := s.cmd.Stdout.(*os.File) + if !ok { + t.Fatalf("expected cmd.Stdout to be *os.File, got %T", s.cmd.Stdout) + } + if gotFile != f { + t.Errorf("expected cmd.Stdout to be the user-provided *os.File "+ + "(fd %d), got a different *os.File (fd %d). The fd-pass "+ + "fast path is broken; sendfile/zero-copy will not apply.", + f.Fd(), gotFile.Fd()) + } + }) + } +} + +// TestCommandStageStdoutFastPathThroughPipeline is the same assertion +// but driven end-to-end through `Pipeline.Start()`, so it also +// exercises the `Pipeline.stdout` plumbing that hands the writer to +// the last stage. +func TestCommandStageStdoutFastPathThroughPipeline(t *testing.T) { + cases := []struct { + name string + option func(*os.File) Option + }{ + { + name: "WithStdoutCloser(*os.File)", + option: func(f *os.File) Option { return WithStdoutCloser(f) }, + }, + { + name: "WithStdout(*os.File)", + option: func(f *os.File) Option { return WithStdout(f) }, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + f, err := os.CreateTemp(t.TempDir(), "stdout") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = f.Close() }) + + cmd := exec.Command("true") + s := CommandStage("true", cmd).(*commandStage) + + p := New(tc.option(f)) + p.Add(s) + if err := p.Start(ctx); err != nil { + t.Fatalf("Start: %v", err) + } + stdoutAfterStart := s.cmd.Stdout + t.Cleanup(func() { _ = p.Wait() }) + + gotFile, ok := stdoutAfterStart.(*os.File) + if !ok { + t.Fatalf("expected cmd.Stdout to be *os.File, got %T", stdoutAfterStart) + } + if gotFile != f { + t.Errorf("expected cmd.Stdout to be the user-provided *os.File "+ + "(fd %d), got a different *os.File (fd %d). The fd-pass "+ + "fast path is broken; sendfile/zero-copy will not apply.", + f.Fd(), gotFile.Fd()) + } + }) + } +} From 3a1f8d84ff1f8e4ad7a709a79f4b11907007dc26 Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Thu, 28 May 2026 19:09:35 +0200 Subject: [PATCH 13/23] Pool buffers for the non-*os.File command stdout copy path When a command stage's stdout is not an `*os.File` (e.g. a `bytes.Buffer` via `WithStdout()` or a custom writer via `WithStdoutCloser()`), the fd-pass fast path doesn't apply and the data has to flow through Go. Left to its own devices, `exec.Cmd` would set up an internal `os.Pipe()` and run `io.Copy` with a fresh 32KB buffer allocated per invocation. To reduce GC pressure, let's do that copy directly: create the `os.Pipe()` ourselves, set the write end as `cmd.Stdout` (so exec.Cmd still does an fd dup into the child), and run the copy from the read end to the user's writer in our own goroutine, drawing the 32KB buffer from a `sync.Pool` (`copy_pool.go`). Destinations that implement `io.ReaderFrom` (`*net.TCPConn`, `*os.File`) are still routed through `ReadFrom` so platform fast paths like splice continue to apply. The pure `*os.File` and `writerNopCloser{*os.File}` paths are unchanged: the fd is still passed directly to the child process. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pipe/command.go | 53 ++++++++++++++++++++++++++++++++++++++++++----- pipe/copy_pool.go | 40 +++++++++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 5 deletions(-) create mode 100644 pipe/copy_pool.go diff --git a/pipe/command.go b/pipe/command.go index 3d61782..e0e13b8 100644 --- a/pipe/command.go +++ b/pipe/command.go @@ -129,9 +129,19 @@ func (s *commandStage) Start( // 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 + // 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. + if f, ok := stdout.Writer.(*os.File); ok { + s.cmd.Stdout = f + } else { + ec, err := s.setupPooledStdout(stdout.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: @@ -139,8 +149,15 @@ func (s *commandStage) Start( 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 + // 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) } } @@ -295,3 +312,29 @@ func (s *commandStage) Wait() error { 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 +} diff --git a/pipe/copy_pool.go b/pipe/copy_pool.go new file mode 100644 index 0000000..69f8a21 --- /dev/null +++ b/pipe/copy_pool.go @@ -0,0 +1,40 @@ +package pipe + +import ( + "io" + "sync" +) + +// copyBufPool reuses 32KB buffers across `io.CopyBuffer` calls, +// avoiding a fresh heap allocation per copy. This matters in +// high-throughput pipelines where many command stages run +// concurrently and stdout is not an `*os.File` that can be passed +// directly through `exec.Cmd`. +var copyBufPool = sync.Pool{ + New: func() any { + b := make([]byte, 32*1024) + return &b + }, +} + +// readerOnly wraps an `io.Reader`, hiding any other interfaces (such +// as `io.WriterTo`) so that `io.CopyBuffer` is forced to use the +// provided buffer. Without this, `*os.File`'s `WriterTo` (added in +// Go 1.26) causes `CopyBuffer` to call `File.WriteTo`, which can +// fall back to `io.Copy` with a fresh allocation, bypassing the pool +// entirely. +type readerOnly struct{ io.Reader } + +// pooledCopy copies from `src` to `dst`. If `dst` implements +// `io.ReaderFrom` (e.g. `*net.TCPConn`, `*os.File`), it delegates to +// `ReadFrom` so platform fast paths like splice can be used. +// Otherwise it falls back to `io.CopyBuffer` with a pooled 32KB +// buffer. +func pooledCopy(dst io.Writer, src io.Reader) (int64, error) { + if rf, ok := dst.(io.ReaderFrom); ok { + return rf.ReadFrom(src) + } + bp := copyBufPool.Get().(*[]byte) + defer copyBufPool.Put(bp) + return io.CopyBuffer(dst, readerOnly{src}, *bp) +} From 46f19e4330624eaed90fac9e20dd1b95dac700e3 Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Thu, 28 May 2026 20:22:20 +0200 Subject: [PATCH 14/23] Forward panic handler through wrapper stages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `efStage` (the wrapper returned by `FilterError` and `IgnoreError`) embeds the `Stage` interface, which only exposes the four Stage methods. So when `Pipeline.Start()` checks `if phs, ok := s.(StagePanicHandlerAware); ok`, the assertion silently fails for any wrapped stage — even if the underlying stage is a `goStage` that implements `SetPanicHandler`. This means a configured `WithStagePanicHandler` is bypassed when the panicking Function is wrapped in `IgnoreError`. The goroutine inside `goStage.Start` sees `panicHandler == nil` and returns without calling `recover()`, letting the panic propagate up the runtime and crash the host process. memoryWatchStage had a similar issue. This was a pre-existing bug in main (not introduced by the version-2 Stage interface redesign), but we're already touching the panic-handler, so let's fix it here. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pipe/filter-error.go | 15 +++++++++++++++ pipe/memorylimit.go | 19 ++++++++++++++++++- pipe/pipeline_test.go | 29 +++++++++++++++++++++++++++++ 3 files changed, 62 insertions(+), 1 deletion(-) diff --git a/pipe/filter-error.go b/pipe/filter-error.go index 654796a..4181f47 100644 --- a/pipe/filter-error.go +++ b/pipe/filter-error.go @@ -26,6 +26,21 @@ func (s efStage) Wait() error { return s.filter(s.Stage.Wait()) } +// SetPanicHandler forwards the handler to the wrapped stage if it +// implements `StagePanicHandlerAware`. Without this, wrapping a +// panicking `Function` in `IgnoreError` / `FilterError` would +// silently bypass `WithStagePanicHandler` (the type assertion in +// `Pipeline.Start()` only sees this wrapper's methods, not the +// embedded stage's `SetPanicHandler`), letting the panic propagate +// out of the goroutine and crash the host process. +func (s efStage) SetPanicHandler(ph StagePanicHandler) { + if phs, ok := s.Stage.(StagePanicHandlerAware); ok { + phs.SetPanicHandler(ph) + } +} + +var _ StagePanicHandlerAware = efStage{} + // ErrorMatcher decides whether its argument matches some class of // errors (e.g., errors that we want to ignore). The function will // only be invoked for non-nil errors. diff --git a/pipe/memorylimit.go b/pipe/memorylimit.go index a863574..e548f2c 100644 --- a/pipe/memorylimit.go +++ b/pipe/memorylimit.go @@ -262,7 +262,10 @@ type memoryWatchStage struct { type memoryWatchFunc func(context.Context, LimitableStage) -var _ LimitableStage = (*memoryWatchStage)(nil) +var ( + _ LimitableStage = (*memoryWatchStage)(nil) + _ StagePanicHandlerAware = (*memoryWatchStage)(nil) +) func (m *memoryWatchStage) Name() string { return m.stage.Name() + m.nameSuffix @@ -272,6 +275,20 @@ func (m *memoryWatchStage) Preferences() StagePreferences { return m.stage.Preferences() } +// SetPanicHandler forwards the handler to the wrapped stage if it +// implements `StagePanicHandlerAware`. Without this, wrapping a +// panicking stage in `MemoryLimit` / `MemoryObserver` / +// `MemoryLimitWithObserver` would silently bypass +// `WithStagePanicHandler` (the type assertion in `Pipeline.Start()` +// only sees this wrapper's methods, not the wrapped stage's +// `SetPanicHandler`), letting the panic propagate out of the +// goroutine and crash the host process. +func (m *memoryWatchStage) SetPanicHandler(ph StagePanicHandler) { + if phs, ok := m.stage.(StagePanicHandlerAware); ok { + phs.SetPanicHandler(ph) + } +} + func (m *memoryWatchStage) Start( ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser, ) error { diff --git a/pipe/pipeline_test.go b/pipe/pipeline_test.go index 0184012..2a24109 100644 --- a/pipe/pipeline_test.go +++ b/pipe/pipeline_test.go @@ -532,6 +532,35 @@ func TestFunction(t *testing.T) { assert.ErrorContains(t, err, "panic handled") assert.Empty(t, out) }) + + t.Run("panic with handler through IgnoreError", func(t *testing.T) { + // Regression: efStage (the FilterError/IgnoreError wrapper) + // previously did not implement StagePanicHandlerAware, so + // the type assertion in Pipeline.Start() silently failed + // and the wrapped goStage never received the panic handler. + // The result was an unrecovered panic crashing the host. + p := pipe.New( + pipe.WithStagePanicHandler(func(p any) error { + return fmt.Errorf("panic handled: %v", p) + }), + ) + p.Add( + pipe.Print("hello world"), + pipe.IgnoreError( + pipe.Function( + "farewell", + func(_ context.Context, _ pipe.Env, _ io.Reader, _ io.Writer) error { + panic("this is a panic") + }, + ), + func(_ error) bool { return false }, + ), + ) + + out, err := p.Output(ctx) + assert.ErrorContains(t, err, "panic handled") + assert.Empty(t, out) + }) } func TestPipelineWithFunction(t *testing.T) { From 121ae945efc8efd77dde066f0b223b702bf209ed Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Fri, 29 May 2026 14:32:56 +0200 Subject: [PATCH 15/23] Avoid leaking pooled-stdout goroutine when cmd.Start() fails Ensure that in case of cmd.Start() failure for pipelines that use pooled buffers, we don't leak a pooled-buffer-copy goroutine, and all closers are closed. Could be triggered/detected by using a command stage that fails on command-not-found under the race detector. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pipe/command.go | 30 +++++++++++++--- pipe/command_starterror_test.go | 63 +++++++++++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 4 deletions(-) create mode 100644 pipe/command_starterror_test.go diff --git a/pipe/command.go b/pipe/command.go index e0e13b8..82a7904 100644 --- a/pipe/command.go +++ b/pipe/command.go @@ -162,6 +162,20 @@ func (s *commandStage) Start( } } + 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 // standard error into our `stderr` field: if s.cmd.Stderr == nil { @@ -171,6 +185,7 @@ func (s *commandStage) Start( // can be sure. p, err := s.cmd.StderrPipe() if err != nil { + cleanupOnStartFailure() return err } s.wg.Go(func() error { @@ -188,12 +203,11 @@ func (s *commandStage) Start( s.runInOwnProcessGroup() if err := s.cmd.Start(); err != nil { + cleanupOnStartFailure() return err } - for _, closer := range earlyClosers { - _ = closer.Close() - } + closeEarlyClosers() // Arrange for the process to be killed (gently) if the context // expires before the command exits normally: @@ -304,12 +318,20 @@ func (s *commandStage) Wait() error { err = wgErr } + 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 } } - return err } diff --git a/pipe/command_starterror_test.go b/pipe/command_starterror_test.go new file mode 100644 index 0000000..191370a --- /dev/null +++ b/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/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") + } +} From bf5820b55685caf18bfcd9a9d05f2d0f3d559a8f Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Fri, 29 May 2026 16:08:07 +0200 Subject: [PATCH 16/23] Bump module path to v2 for breaking Stage interface change The Stage interface in this series is not backwards compatible: - Start()'s signature changed from: Start(ctx, env, stdin io.ReadCloser) (io.ReadCloser, error) to: Start(ctx, env, stdin io.ReadCloser, stdout io.WriteCloser) error - Stage gained a new required method Preferences(). Callers that only construct stages via the package's exported constructors (Command(), CommandStage(), Function(), ...) are unaffected, but anyone implementing Stage themselves has to update their implementation. --- README.md | 4 ++-- go.mod | 2 +- internal/ptree/ptree_test.go | 2 +- pipe/command_linux.go | 2 +- pipe/command_starterror_test.go | 2 +- pipe/memorylimit_test.go | 2 +- pipe/pipe_matching_test.go | 2 +- pipe/pipeline_test.go | 2 +- 8 files changed, 9 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index fc61117..9896f66 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# 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 @@ -6,4 +6,4 @@ We have not thoroughly tested this package on OSs other than Linux, especially W # Links -* [Docs](https://pkg.go.dev/github.com/github/go-pipe) +* [Docs](https://pkg.go.dev/github.com/github/go-pipe/v2) diff --git a/go.mod b/go.mod index 6f69110..041809e 100644 --- a/go.mod +++ b/go.mod @@ -1,4 +1,4 @@ -module github.com/github/go-pipe +module github.com/github/go-pipe/v2 go 1.24.0 diff --git a/internal/ptree/ptree_test.go b/internal/ptree/ptree_test.go index 9c6d4e4..5c014c2 100644 --- a/internal/ptree/ptree_test.go +++ b/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" ) diff --git a/pipe/command_linux.go b/pipe/command_linux.go index 997d2cd..987bffb 100644 --- a/pipe/command_linux.go +++ b/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. diff --git a/pipe/command_starterror_test.go b/pipe/command_starterror_test.go index 191370a..6099d1d 100644 --- a/pipe/command_starterror_test.go +++ b/pipe/command_starterror_test.go @@ -7,7 +7,7 @@ import ( "sync/atomic" "testing" - "github.com/github/go-pipe/pipe" + "github.com/github/go-pipe/v2/pipe" ) // TestCommandStageStartFailureNoRace verifies that when `cmd.Start()` diff --git a/pipe/memorylimit_test.go b/pipe/memorylimit_test.go index 17550e1..b5943db 100644 --- a/pipe/memorylimit_test.go +++ b/pipe/memorylimit_test.go @@ -12,7 +12,7 @@ import ( "testing" "time" - "github.com/github/go-pipe/pipe" + "github.com/github/go-pipe/v2/pipe" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) diff --git a/pipe/pipe_matching_test.go b/pipe/pipe_matching_test.go index ace6103..5d0db15 100644 --- a/pipe/pipe_matching_test.go +++ b/pipe/pipe_matching_test.go @@ -7,7 +7,7 @@ import ( "os" "testing" - "github.com/github/go-pipe/pipe" + "github.com/github/go-pipe/v2/pipe" "github.com/stretchr/testify/assert" ) diff --git a/pipe/pipeline_test.go b/pipe/pipeline_test.go index 2a24109..bce81bf 100644 --- a/pipe/pipeline_test.go +++ b/pipe/pipeline_test.go @@ -18,7 +18,7 @@ import ( "github.com/stretchr/testify/require" "go.uber.org/goleak" - "github.com/github/go-pipe/pipe" + "github.com/github/go-pipe/v2/pipe" ) // Check whether this package's test suite leaks any goroutines: From 6102173e00c34c1fe64fa22a9fe2edb3f414a339 Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Fri, 29 May 2026 18:43:13 +0200 Subject: [PATCH 17/23] remove IOPreferenceNil It doesn't really mean much of anything to specify IOPreferenceNil. Nothing uses it, which isn't surprising because it's not clear what it would even mean anywhere other than the begin/end of a pipeline. --- pipe/command.go | 10 +--------- pipe/pipe_matching_test.go | 40 ++++++++++++++++++++++---------------- pipe/stage.go | 5 ----- 3 files changed, 24 insertions(+), 31 deletions(-) diff --git a/pipe/command.go b/pipe/command.go index 82a7904..cbe5521 100644 --- a/pipe/command.go +++ b/pipe/command.go @@ -70,18 +70,10 @@ func (s *commandStage) Name() string { } func (s *commandStage) Preferences() StagePreferences { - prefs := StagePreferences{ + return 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( diff --git a/pipe/pipe_matching_test.go b/pipe/pipe_matching_test.go index 5d0db15..fd07081 100644 --- a/pipe/pipe_matching_test.go +++ b/pipe/pipe_matching_test.go @@ -22,6 +22,12 @@ import ( const ( IOPreferenceUndefinedNopCloser pipe.IOPreference = iota + 100 IOPreferenceFileNopCloser + + // expectNil is a test-only expectation token meaning that the + // stage should be passed a `nil` stdin / stdout (which happens at + // the beginning / end of a pipeline when no overall stdin / stdout + // is configured). It is not a real `IOPreference`. + expectNil ) func file(t *testing.T) *os.File { @@ -143,7 +149,7 @@ func prefString(pref pipe.IOPreference) string { return "other" case pipe.IOPreferenceFile: return "*os.File" - case pipe.IOPreferenceNil: + case expectNil: return "nil" case IOPreferenceUndefinedNopCloser: return "nopCloser(other)" @@ -200,7 +206,7 @@ func TestPipeTypes(t *testing.T) { name: "func", opts: []pipe.Option{}, stages: []pipe.Stage{ - newPipeSniffingFunc(pipe.IOPreferenceNil, pipe.IOPreferenceNil), + newPipeSniffingFunc(expectNil, expectNil), }, }, { @@ -209,7 +215,7 @@ func TestPipeTypes(t *testing.T) { pipe.WithStdin(file(t)), }, stages: []pipe.Stage{ - newPipeSniffingFunc(IOPreferenceFileNopCloser, pipe.IOPreferenceNil), + newPipeSniffingFunc(IOPreferenceFileNopCloser, expectNil), }, }, { @@ -218,7 +224,7 @@ func TestPipeTypes(t *testing.T) { pipe.WithStdout(file(t)), }, stages: []pipe.Stage{ - newPipeSniffingFunc(pipe.IOPreferenceNil, IOPreferenceFileNopCloser), + newPipeSniffingFunc(expectNil, IOPreferenceFileNopCloser), }, }, { @@ -227,7 +233,7 @@ func TestPipeTypes(t *testing.T) { pipe.WithStdoutCloser(file(t)), }, stages: []pipe.Stage{ - newPipeSniffingFunc(pipe.IOPreferenceNil, pipe.IOPreferenceFile), + newPipeSniffingFunc(expectNil, pipe.IOPreferenceFile), }, }, { @@ -244,7 +250,7 @@ func TestPipeTypes(t *testing.T) { name: "cmd", opts: []pipe.Option{}, stages: []pipe.Stage{ - newPipeSniffingCmd(pipe.IOPreferenceNil, pipe.IOPreferenceNil), + newPipeSniffingCmd(expectNil, expectNil), }, }, { @@ -253,7 +259,7 @@ func TestPipeTypes(t *testing.T) { pipe.WithStdin(file(t)), }, stages: []pipe.Stage{ - newPipeSniffingCmd(IOPreferenceFileNopCloser, pipe.IOPreferenceNil), + newPipeSniffingCmd(IOPreferenceFileNopCloser, expectNil), }, }, { @@ -262,7 +268,7 @@ func TestPipeTypes(t *testing.T) { pipe.WithStdout(file(t)), }, stages: []pipe.Stage{ - newPipeSniffingCmd(pipe.IOPreferenceNil, IOPreferenceFileNopCloser), + newPipeSniffingCmd(expectNil, IOPreferenceFileNopCloser), }, }, { @@ -271,7 +277,7 @@ func TestPipeTypes(t *testing.T) { pipe.WithStdoutCloser(file(t)), }, stages: []pipe.Stage{ - newPipeSniffingCmd(pipe.IOPreferenceNil, pipe.IOPreferenceFile), + newPipeSniffingCmd(expectNil, pipe.IOPreferenceFile), }, }, { @@ -301,7 +307,7 @@ func TestPipeTypes(t *testing.T) { pipe.WithStdout(file(t)), }, stages: []pipe.Stage{ - newPipeSniffingFunc(pipe.IOPreferenceNil, pipe.IOPreferenceFile), + newPipeSniffingFunc(expectNil, pipe.IOPreferenceFile), newPipeSniffingCmd(pipe.IOPreferenceFile, IOPreferenceFileNopCloser), }, }, @@ -312,15 +318,15 @@ func TestPipeTypes(t *testing.T) { }, stages: []pipe.Stage{ newPipeSniffingCmd(IOPreferenceUndefinedNopCloser, pipe.IOPreferenceFile), - newPipeSniffingFunc(pipe.IOPreferenceFile, pipe.IOPreferenceNil), + newPipeSniffingFunc(pipe.IOPreferenceFile, expectNil), }, }, { name: "cmd-cmd", opts: []pipe.Option{}, stages: []pipe.Stage{ - newPipeSniffingCmd(pipe.IOPreferenceNil, pipe.IOPreferenceFile), - newPipeSniffingCmd(pipe.IOPreferenceFile, pipe.IOPreferenceNil), + newPipeSniffingCmd(expectNil, pipe.IOPreferenceFile), + newPipeSniffingCmd(pipe.IOPreferenceFile, expectNil), }, }, { @@ -328,7 +334,7 @@ func TestPipeTypes(t *testing.T) { opts: []pipe.Option{}, stages: []pipe.Stage{ newPipeSniffingStage( - pipe.IOPreferenceUndefined, pipe.IOPreferenceNil, + pipe.IOPreferenceUndefined, expectNil, pipe.IOPreferenceUndefined, pipe.IOPreferenceUndefined, ), newPipeSniffingStage( @@ -337,7 +343,7 @@ func TestPipeTypes(t *testing.T) { ), newPipeSniffingStage( pipe.IOPreferenceUndefined, pipe.IOPreferenceFile, - pipe.IOPreferenceUndefined, pipe.IOPreferenceNil, + pipe.IOPreferenceUndefined, expectNil, ), }, }, @@ -346,7 +352,7 @@ func TestPipeTypes(t *testing.T) { opts: []pipe.Option{}, stages: []pipe.Stage{ newPipeSniffingStage( - pipe.IOPreferenceUndefined, pipe.IOPreferenceNil, + pipe.IOPreferenceUndefined, expectNil, pipe.IOPreferenceUndefined, pipe.IOPreferenceFile, ), newPipeSniffingStage( @@ -355,7 +361,7 @@ func TestPipeTypes(t *testing.T) { ), newPipeSniffingStage( pipe.IOPreferenceUndefined, pipe.IOPreferenceUndefined, - pipe.IOPreferenceUndefined, pipe.IOPreferenceNil, + pipe.IOPreferenceUndefined, expectNil, ), }, }, diff --git a/pipe/stage.go b/pipe/stage.go index 8b6dff5..22a550e 100644 --- a/pipe/stage.go +++ b/pipe/stage.go @@ -138,9 +138,4 @@ const ( // IOPreferenceFile indicates that the stage would prefer for the // specified stdin / stdout to be an `*os.File`, to avoid copying. IOPreferenceFile - - // IOPreferenceNil indicates that the stage does not use the - // specified stdin / stdout, so `nil` should be passed in. This - // should only happen at the beginning / end of a pipeline. - IOPreferenceNil ) From 754784b58364869c527ca63508a4195e4fe6a96d Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Fri, 29 May 2026 19:50:40 +0200 Subject: [PATCH 18/23] properly handle nil input/output --- pipe/function.go | 9 +++++++++ pipe/memorylimit_test.go | 6 ++++-- pipe/pipeline_test.go | 20 ++++++++++++++++++++ 3 files changed, 33 insertions(+), 2 deletions(-) diff --git a/pipe/function.go b/pipe/function.go index 958ee4b..c9429fb 100644 --- a/pipe/function.go +++ b/pipe/function.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "io" + "strings" ) // StageFunc is a function that can be used to power a `goStage`. It @@ -66,11 +67,19 @@ func (s *goStage) Start( if stdin, ok := stdin.(readerNopCloser); ok { r = stdin.Reader } + if r == nil { + // treat nil as empty input. + r = strings.NewReader("") + } var w io.Writer = stdout if stdout, ok := stdout.(writerNopCloser); ok { w = stdout.Writer } + if w == nil { + // treat nil output as /dev/null + w = io.Discard + } go func() { defer close(s.done) diff --git a/pipe/memorylimit_test.go b/pipe/memorylimit_test.go index b5943db..582421d 100644 --- a/pipe/memorylimit_test.go +++ b/pipe/memorylimit_test.go @@ -210,7 +210,8 @@ func testMemoryLimit(t *testing.T, mbs int, limit uint64, stage pipe.Stage) (str for i := 0; i < mbs; i++ { _, err := stdout.Write(bytes[:]) if err != nil { - require.ErrorIs(t, err, syscall.EPIPE) + assert.ErrorIs(t, err, syscall.EPIPE) + return nil } } @@ -244,7 +245,8 @@ func testMemoryLimitWithObserver(t *testing.T, mbs int, limit uint64, stage pipe for i := 0; i < mbs; i++ { _, err := stdout.Write(bytes[:]) if err != nil { - require.ErrorIs(t, err, syscall.EPIPE) + assert.ErrorIs(t, err, syscall.EPIPE) + return nil } } return nil diff --git a/pipe/pipeline_test.go b/pipe/pipeline_test.go index bce81bf..9419c2e 100644 --- a/pipe/pipeline_test.go +++ b/pipe/pipeline_test.go @@ -798,6 +798,26 @@ func TestPrintf(t *testing.T) { } } +func TestPrintlnNoOutput(t *testing.T) { + t.Parallel() + ctx := context.Background() + p := pipe.New() + p.Add(pipe.Println("Look Ma, no output!")) + assert.NoError(t, p.Run(ctx)) +} + +func TestFunctionNoInput(t *testing.T) { + t.Parallel() + ctx := context.Background() + p := pipe.New() + p.Add(pipe.Function("read-all", func(_ context.Context, _ pipe.Env, stdin io.Reader, _ io.Writer) error { + n, err := io.Copy(io.Discard, stdin) + assert.Equal(t, int64(0), n) + return err + })) + assert.NoError(t, p.Run(ctx)) +} + func TestErrors(t *testing.T) { t.Parallel() ctx := context.Background() From 9dca23eae08a236897df53494484719755e71f96 Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Fri, 29 May 2026 22:07:01 +0200 Subject: [PATCH 19/23] allow earlier GC of lateClosers PR feedback: nil out lateClosers after use, so that closers can potentially be garbage collected in cases where the stage hangs around for a while. --- pipe/command.go | 1 + 1 file changed, 1 insertion(+) diff --git a/pipe/command.go b/pipe/command.go index cbe5521..a4f8916 100644 --- a/pipe/command.go +++ b/pipe/command.go @@ -324,6 +324,7 @@ func (s *commandStage) closeLateClosers() error { err = closeErr } } + s.lateClosers = nil return err } From 9ee30b1f23664cbbd9e7005e798f945bd5bd1c5d Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Sat, 30 May 2026 18:48:51 +0200 Subject: [PATCH 20/23] export Unwrap{Reader,Writer}; unwrap goStage stdin fully The internal noop-closer wrappers that go-pipe puts around the caller-owned stdin/stdout hid the underlying object's concrete type from stages. These were supposed to be unwrapped, but this wasn't done consistently (readerWriterToNopCloser), and external Stage implementations had no blessed way to unwrap at all. So, introduce exported UnwrapReader/UnwrapWriter as the sole way for recovering the underlying reader/writer, and delete readerWriterToNopCloser in favor of explicit unwrap. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pipe/command.go | 11 +++---- pipe/function.go | 10 ++---- pipe/nop_closer.go | 63 +++++++++++++++++++++---------------- pipe/nop_closer_test.go | 70 +++++++++++++++++++++++++++++++++++++++++ pipe/stage.go | 8 ++--- 5 files changed, 116 insertions(+), 46 deletions(-) create mode 100644 pipe/nop_closer_test.go diff --git a/pipe/command.go b/pipe/command.go index a4f8916..d3b7e8a 100644 --- a/pipe/command.go +++ b/pipe/command.go @@ -98,11 +98,7 @@ func (s *commandStage) Start( 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 + s.cmd.Stdin = UnwrapReader(stdin) case *os.File: // In this case, we can close stdin as soon as the command // has started: @@ -125,10 +121,11 @@ func (s *commandStage) Start( // 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. - if f, ok := stdout.Writer.(*os.File); ok { + writer := UnwrapWriter(stdout) + if f, ok := writer.(*os.File); ok { s.cmd.Stdout = f } else { - ec, err := s.setupPooledStdout(stdout.Writer) + ec, err := s.setupPooledStdout(writer) if err != nil { return err } diff --git a/pipe/function.go b/pipe/function.go index c9429fb..0ff48f3 100644 --- a/pipe/function.go +++ b/pipe/function.go @@ -63,19 +63,13 @@ func (s *goStage) Preferences() StagePreferences { 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 - } + r := UnwrapReader(stdin) if r == nil { // treat nil as empty input. r = strings.NewReader("") } - var w io.Writer = stdout - if stdout, ok := stdout.(writerNopCloser); ok { - w = stdout.Writer - } + w := UnwrapWriter(stdout) if w == nil { // treat nil output as /dev/null w = io.Discard diff --git a/pipe/nop_closer.go b/pipe/nop_closer.go index 18cf7a9..739497f 100644 --- a/pipe/nop_closer.go +++ b/pipe/nop_closer.go @@ -6,21 +6,17 @@ package pipe import "io" -// newReaderNopCloser returns a ReadCloser with a no-op Close method, -// wrapping the provided io.Reader `r`. If `r` implements -// `io.WriterTo`, the returned `io.ReadCloser` will also implement -// `io.WriterTo` by forwarding calls to `r`. +// newReaderNopCloser returns a ReadCloser with a no-op Close method, wrapping +// the provided io.Reader `r`. The wrapper deliberately hides `r`'s concrete +// type; use [UnwrapReader] to recover the underlying reader before use in +// situations where eg. use of io.WriterTo is important for performance. func newReaderNopCloser(r io.Reader) io.ReadCloser { - if _, ok := r.(io.WriterTo); ok { - return readerWriterToNopCloser{r} - } return readerNopCloser{r} } // readerNopCloser is a ReadCloser that wraps a provided `io.Reader`, -// but whose `Close()` method does nothing. We don't need to check -// whether the wrapped reader also implements `io.WriterTo`, since -// it's always unwrapped before use. +// but whose `Close()` method does nothing. It should be unwrapped (via +// [UnwrapReader]) before use. type readerNopCloser struct { io.Reader } @@ -29,21 +25,10 @@ func (readerNopCloser) Close() error { return nil } -// readerWriterToNopCloser is like `readerNopCloser` except that it -// also implements `io.WriterTo` by delegating `WriteTo()` to the -// wrapped `io.Reader` (which must also implement `io.WriterTo`). -type readerWriterToNopCloser struct { - io.Reader -} - -func (readerWriterToNopCloser) Close() error { return nil } - -func (r readerWriterToNopCloser) WriteTo(w io.Writer) (n int64, err error) { - return r.Reader.(io.WriterTo).WriteTo(w) -} - -// writerNopCloser is a WriteCloser that wraps a provided `io.Writer`, -// but whose `Close()` method does nothing. +// 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. type writerNopCloser struct { io.Writer } @@ -52,6 +37,32 @@ func (w writerNopCloser) Close() error { return nil } +// UnwrapReader returns the underlying [io.Reader] that go-pipe wrapped around +// the `stdin` it passes to a [Stage]'s `Start` method. [Stage] implementations +// should call this before reading from `stdin` so it sees the caller's +// concrete reader type, because that allows for use of fast-path interfaces +// (e.g. `io.WriterTo`) and identity (e.g. `*os.File`, for direct fd passing). +// +// If `r` is not a go-pipe wrapper (including nil), it is returned unchanged. +func UnwrapReader(r io.Reader) io.Reader { + if w, ok := r.(readerNopCloser); ok { + return w.Reader + } + return r +} + +// UnwrapWriter returns the underlying [io.Writer] that go-pipe wrapped around +// the `stdout` it passes to a [Stage]'s `Start` method. [Stage] +// implementations should call this before writing to `stdout` (see above). +// +// If `w` is not a go-pipe wrapper (including nil), it is returned unchanged. +func UnwrapWriter(w io.Writer) io.Writer { + if n, ok := w.(writerNopCloser); ok { + return n.Writer + } + return w +} + // unwrapNopCloser unwraps the object if it is some kind of nop // closer, and returns the underlying object. This function is used // only for testing. @@ -59,8 +70,6 @@ func unwrapNopCloser(obj any) (any, bool) { switch obj := obj.(type) { case readerNopCloser: return obj.Reader, true - case readerWriterToNopCloser: - return obj.Reader, true case writerNopCloser: return obj.Writer, true default: diff --git a/pipe/nop_closer_test.go b/pipe/nop_closer_test.go new file mode 100644 index 0000000..13f9935 --- /dev/null +++ b/pipe/nop_closer_test.go @@ -0,0 +1,70 @@ +package pipe + +import ( + "bytes" + "context" + "io" + "testing" +) + +func TestUnwrapReader(t *testing.T) { + src := bytes.NewReader([]byte("payload")) + + if got := UnwrapReader(newReaderNopCloser(src)); got != io.Reader(src) { + t.Errorf("UnwrapReader(wrapped) = %T %p, want %p", got, got, src) + } + + // A non-wrapped reader passes through unchanged. + if got := UnwrapReader(src); got != io.Reader(src) { + t.Errorf("UnwrapReader(plain) = %T %p, want %p", got, got, src) + } + + if got := UnwrapReader(nil); got != nil { + t.Errorf("UnwrapReader(nil) = %v, want nil", got) + } +} + +func TestUnwrapWriter(t *testing.T) { + dst := &bytes.Buffer{} + + if got := UnwrapWriter(writerNopCloser{dst}); got != io.Writer(dst) { + t.Errorf("UnwrapWriter(wrapped) = %T %p, want %p", got, got, dst) + } + + // A non-wrapped writer passes through unchanged. + if got := UnwrapWriter(dst); got != io.Writer(dst) { + t.Errorf("UnwrapWriter(plain) = %T %p, want %p", got, got, dst) + } + + if got := UnwrapWriter(nil); got != nil { + t.Errorf("UnwrapWriter(nil) = %v, want nil", got) + } +} + +// TestGoStageUnwrapsWriterToStdin verifies that a Function stage +// receives its stdin already unwrapped to the caller's concrete type, +// so fast-path interfaces such as io.WriterTo survive. This guards +// against the regression where goStage only unwrapped one of the +// internal nop-closer wrapper types. +func TestGoStageUnwrapsWriterToStdin(t *testing.T) { + src := bytes.NewReader([]byte("hello")) + + var got io.Reader + p := New(WithStdin(src), WithStdout(io.Discard)) + p.Add(Function("capture", func(_ context.Context, _ Env, stdin io.Reader, _ io.Writer) error { + got = stdin + _, err := io.Copy(io.Discard, stdin) + return err + })) + + if err := p.Run(context.Background()); err != nil { + t.Fatalf("Run: %v", err) + } + + if got != io.Reader(src) { + t.Fatalf("StageFunc stdin = %T %p, want unwrapped *bytes.Reader %p", got, got, src) + } + if _, ok := got.(io.WriterTo); !ok { + t.Fatalf("unwrapped stdin %T does not expose io.WriterTo fast path", got) + } +} diff --git a/pipe/stage.go b/pipe/stage.go index 22a550e..bc6866f 100644 --- a/pipe/stage.go +++ b/pipe/stage.go @@ -66,10 +66,10 @@ import ( // 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, it's -// advantageous for the first stage should ideally unwrap its stdin -// argument before actually using it. If the wrapped value is an -// `*os.File` and the stage is a command stage, then unwrapping is +// 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. // // For stdout, it depends on whether the user supplied it using From 75797b93d71c25b8baeb5130fec68376ad7e3e2e Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Fri, 29 May 2026 23:45:32 +0200 Subject: [PATCH 21/23] rewrite cleanup/unwind in forward order Based on PR feedback, rewrite the use of recover() for function stages so that the code flows in forward order. Inline recoverPanic() to keep the recover in the first stack frame. --- pipe/function.go | 27 +++++---------- pipe/function_panic_test.go | 68 +++++++++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 18 deletions(-) create mode 100644 pipe/function_panic_test.go diff --git a/pipe/function.go b/pipe/function.go index 0ff48f3..0f33599 100644 --- a/pipe/function.go +++ b/pipe/function.go @@ -76,23 +76,24 @@ func (s *goStage) Start( } go func() { - defer close(s.done) defer func() { - 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) + if s.panicHandler != nil { + if p := recover(); p != nil { + s.err = s.panicHandler(p) } } - }() - defer func() { 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) } } + 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) }() - defer s.recoverPanic() - s.err = s.f(ctx, env, r, w) }() @@ -103,13 +104,3 @@ 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) - } -} diff --git a/pipe/function_panic_test.go b/pipe/function_panic_test.go new file mode 100644 index 0000000..143abf3 --- /dev/null +++ b/pipe/function_panic_test.go @@ -0,0 +1,68 @@ +package pipe_test + +import ( + "context" + "io" + "os" + "os/exec" + "strings" + "testing" + "time" + + "github.com/github/go-pipe/v2/pipe" +) + +const panicChildEnv = "GO_PIPE_FUNCTION_PANIC_CHILD" +const panicSentinel = "function-panic-sentinel" + +// TestFunctionPanicWithoutHandlerPropagates verifies that when a +// `Function` stage panics and no panic handler is installed, the panic +// propagates (crashing the process) rather than being silently +// swallowed and reported as a successful run. Because a propagating +// panic would crash the test binary itself, the actual pipeline is run +// in a re-exec'd subprocess and this test asserts on its outcome. +func TestFunctionPanicWithoutHandlerPropagates(t *testing.T) { + if os.Getenv(panicChildEnv) == "1" { + runPanicChild() + return + } + + cmd := exec.Command(os.Args[0], "-test.run=^TestFunctionPanicWithoutHandlerPropagates$", "-test.v") //nolint:gosec // re-exec of this test binary with constant arguments. + cmd.Env = append(os.Environ(), panicChildEnv+"=1") + out, err := cmd.CombinedOutput() + output := string(out) + + if err == nil { + t.Fatalf("expected subprocess to crash from a propagated panic, but it exited 0\noutput:\n%s", output) + } + if strings.Contains(output, "SURVIVED") { + t.Fatalf("panic was swallowed: Run returned instead of propagating\noutput:\n%s", output) + } + if !strings.Contains(output, "panic:") || !strings.Contains(output, panicSentinel) { + t.Fatalf("expected a propagated panic mentioning %q, got:\n%s", panicSentinel, output) + } +} + +// runPanicChild runs a pipeline whose only stage is a `Function` that panics, +// with no panic handler configured. The panic, being unhandled, should crash +// the process before the sleep elapses; if it is swallowed (the regression), +// Run returns and we print SURVIVED so the parent can detect the failure. +func runPanicChild() { + p := pipe.New(pipe.WithStdout(io.Discard)) + p.Add(pipe.Function("boom", func(_ context.Context, _ pipe.Env, _ io.Reader, _ io.Writer) error { + panic(panicSentinel) + })) + + err := p.Run(context.Background()) + + // reaching this point at all indicates the panic was swallowed. + time.Sleep(2 * time.Second) + os.Stdout.WriteString("SURVIVED: Run returned err=") + if err != nil { + os.Stdout.WriteString(err.Error()) + } else { + os.Stdout.WriteString("") + } + os.Stdout.WriteString("\n") + os.Exit(0) +} From 512c3452d14784b00fd7608b5dc650e9fe0474b8 Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Sat, 30 May 2026 00:27:24 +0200 Subject: [PATCH 22/23] Thread panic handler through Start; drop StagePanicHandlerAware Of all the different types of pipe.Stage, most don't need to have a panic handler, because most are not running user functions. Yet we were paying the price of having panic forwarding as part of the interface, which was awkward and error-prone for the rest of the stages to implement cleanly. Instead, we can just pass the panic handler through Start instead. We use a trailing StartOptions struct carrying PanicHandler in Stage.Start. The StagePanicHandlerAware interface and its panic.go file are removed (StagePanicHandler moves next to StartOptions in stage.go). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pipe/command.go | 2 +- pipe/command_nil_panic_test.go | 2 +- pipe/command_stdout_fastpath_test.go | 2 +- pipe/filter-error.go | 15 --------------- pipe/function.go | 24 ++++++++---------------- pipe/memorylimit.go | 23 +++-------------------- pipe/panic.go | 12 ------------ pipe/pipe_matching_test.go | 2 +- pipe/pipeline.go | 12 ++---------- pipe/pipeline_test.go | 7 +------ pipe/stage.go | 16 +++++++++++++++- 11 files changed, 33 insertions(+), 84 deletions(-) delete mode 100644 pipe/panic.go diff --git a/pipe/command.go b/pipe/command.go index d3b7e8a..0daa91e 100644 --- a/pipe/command.go +++ b/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, + ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser, _ StartOptions, ) error { if s.cmd.Dir == "" { s.cmd.Dir = env.Dir diff --git a/pipe/command_nil_panic_test.go b/pipe/command_nil_panic_test.go index d19ce44..54c1508 100644 --- a/pipe/command_nil_panic_test.go +++ b/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, nil) + err := stage.Start(ctx, Env{}, nil, nil, StartOptions{}) if err == nil { t.Fatal("Expected start to fail, but it succeeded") } diff --git a/pipe/command_stdout_fastpath_test.go b/pipe/command_stdout_fastpath_test.go index 66ce4fb..42f581a 100644 --- a/pipe/command_stdout_fastpath_test.go +++ b/pipe/command_stdout_fastpath_test.go @@ -43,7 +43,7 @@ func TestCommandStageStdoutFastPath(t *testing.T) { cmd := exec.Command("true") s := CommandStage("true", cmd).(*commandStage) - if err := s.Start(ctx, Env{}, nil, tc.wrap(f)); err != nil { + if err := s.Start(ctx, Env{}, nil, tc.wrap(f), StartOptions{}); err != nil { t.Fatalf("Start: %v", err) } t.Cleanup(func() { _ = s.Wait() }) diff --git a/pipe/filter-error.go b/pipe/filter-error.go index 4181f47..654796a 100644 --- a/pipe/filter-error.go +++ b/pipe/filter-error.go @@ -26,21 +26,6 @@ func (s efStage) Wait() error { return s.filter(s.Stage.Wait()) } -// SetPanicHandler forwards the handler to the wrapped stage if it -// implements `StagePanicHandlerAware`. Without this, wrapping a -// panicking `Function` in `IgnoreError` / `FilterError` would -// silently bypass `WithStagePanicHandler` (the type assertion in -// `Pipeline.Start()` only sees this wrapper's methods, not the -// embedded stage's `SetPanicHandler`), letting the panic propagate -// out of the goroutine and crash the host process. -func (s efStage) SetPanicHandler(ph StagePanicHandler) { - if phs, ok := s.Stage.(StagePanicHandlerAware); ok { - phs.SetPanicHandler(ph) - } -} - -var _ StagePanicHandlerAware = efStage{} - // ErrorMatcher decides whether its argument matches some class of // errors (e.g., errors that we want to ignore). The function will // only be invoked for non-nil errors. diff --git a/pipe/function.go b/pipe/function.go index 0f33599..00ca595 100644 --- a/pipe/function.go +++ b/pipe/function.go @@ -33,26 +33,18 @@ 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) - _ StagePanicHandlerAware = (*goStage)(nil) -) +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, @@ -61,7 +53,7 @@ func (s *goStage) Preferences() StagePreferences { } func (s *goStage) Start( - ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser, + ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser, opts StartOptions, ) error { r := UnwrapReader(stdin) if r == nil { @@ -77,9 +69,9 @@ func (s *goStage) Start( go func() { defer func() { - if s.panicHandler != nil { + if opts.PanicHandler != nil { if p := recover(); p != nil { - s.err = s.panicHandler(p) + s.err = opts.PanicHandler(p) } } if stdout != nil { diff --git a/pipe/memorylimit.go b/pipe/memorylimit.go index e548f2c..408f44d 100644 --- a/pipe/memorylimit.go +++ b/pipe/memorylimit.go @@ -262,10 +262,7 @@ type memoryWatchStage struct { type memoryWatchFunc func(context.Context, LimitableStage) -var ( - _ LimitableStage = (*memoryWatchStage)(nil) - _ StagePanicHandlerAware = (*memoryWatchStage)(nil) -) +var _ LimitableStage = (*memoryWatchStage)(nil) func (m *memoryWatchStage) Name() string { return m.stage.Name() + m.nameSuffix @@ -275,24 +272,10 @@ func (m *memoryWatchStage) Preferences() StagePreferences { return m.stage.Preferences() } -// SetPanicHandler forwards the handler to the wrapped stage if it -// implements `StagePanicHandlerAware`. Without this, wrapping a -// panicking stage in `MemoryLimit` / `MemoryObserver` / -// `MemoryLimitWithObserver` would silently bypass -// `WithStagePanicHandler` (the type assertion in `Pipeline.Start()` -// only sees this wrapper's methods, not the wrapped stage's -// `SetPanicHandler`), letting the panic propagate out of the -// goroutine and crash the host process. -func (m *memoryWatchStage) SetPanicHandler(ph StagePanicHandler) { - if phs, ok := m.stage.(StagePanicHandlerAware); ok { - phs.SetPanicHandler(ph) - } -} - func (m *memoryWatchStage) Start( - ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser, + ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser, opts StartOptions, ) error { - if err := m.stage.Start(ctx, env, stdin, stdout); err != nil { + if err := m.stage.Start(ctx, env, stdin, stdout, opts); err != nil { return err } diff --git a/pipe/panic.go b/pipe/panic.go deleted file mode 100644 index e0ca600..0000000 --- a/pipe/panic.go +++ /dev/null @@ -1,12 +0,0 @@ -package pipe - -// 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 -// within that goroutine. -type StagePanicHandlerAware interface { - SetPanicHandler(StagePanicHandler) -} - -// StagePanicHandler is a function that handles panics in a pipeline's stages. -type StagePanicHandler func(p any) error diff --git a/pipe/pipe_matching_test.go b/pipe/pipe_matching_test.go index fd07081..4ac206c 100644 --- a/pipe/pipe_matching_test.go +++ b/pipe/pipe_matching_test.go @@ -98,7 +98,7 @@ func (s *pipeSniffingStage) Preferences() pipe.StagePreferences { } func (s *pipeSniffingStage) Start( - _ context.Context, _ pipe.Env, stdin io.ReadCloser, stdout io.WriteCloser, + _ context.Context, _ pipe.Env, stdin io.ReadCloser, stdout io.WriteCloser, _ pipe.StartOptions, ) error { s.stdin = stdin if stdin != nil { diff --git a/pipe/pipeline.go b/pipe/pipeline.go index d160857..2ec44b5 100644 --- a/pipe/pipeline.go +++ b/pipe/pipeline.go @@ -348,10 +348,6 @@ func (p *Pipeline) Start(ctx context.Context) error { ss := &stageStarters[i] nextSS := &stageStarters[i+1] - if phs, ok := s.(StagePanicHandlerAware); ok && p.panicHandler != nil { - phs.SetPanicHandler(p.panicHandler) - } - // We need to generate a pipe pair for this stage to use // to communicate with its successor: if ss.prefs.StdoutPreference == IOPreferenceFile || @@ -365,7 +361,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); err != nil { + if err := s.Start(ctx, p.env, ss.stdin, ss.stdout, StartOptions{PanicHandler: p.panicHandler}); err != nil { nextSS.stdin.Close() ss.stdout.Close() return abort(i, err) @@ -380,11 +376,7 @@ func (p *Pipeline) Start(ctx context.Context) error { s := p.stages[i] ss := &stageStarters[i] - if phs, ok := s.(StagePanicHandlerAware); ok && p.panicHandler != nil { - phs.SetPanicHandler(p.panicHandler) - } - - if err := s.Start(ctx, p.env, ss.stdin, ss.stdout); err != nil { + if err := s.Start(ctx, p.env, ss.stdin, ss.stdout, StartOptions{PanicHandler: p.panicHandler}); err != nil { return abort(i, err) } } diff --git a/pipe/pipeline_test.go b/pipe/pipeline_test.go index 9419c2e..f8712b4 100644 --- a/pipe/pipeline_test.go +++ b/pipe/pipeline_test.go @@ -534,11 +534,6 @@ func TestFunction(t *testing.T) { }) t.Run("panic with handler through IgnoreError", func(t *testing.T) { - // Regression: efStage (the FilterError/IgnoreError wrapper) - // previously did not implement StagePanicHandlerAware, so - // the type assertion in Pipeline.Start() silently failed - // and the wrapped goStage never received the panic handler. - // The result was an unrecovered panic crashing the host. p := pipe.New( pipe.WithStagePanicHandler(func(p any) error { return fmt.Errorf("panic handled: %v", p) @@ -607,7 +602,7 @@ func (s ErrorStartingStage) Preferences() pipe.StagePreferences { } func (s ErrorStartingStage) Start( - _ context.Context, _ pipe.Env, stdin io.ReadCloser, stdout io.WriteCloser, + _ context.Context, _ pipe.Env, stdin io.ReadCloser, stdout io.WriteCloser, _ pipe.StartOptions, ) error { if stdin != nil { _ = stdin.Close() diff --git a/pipe/stage.go b/pipe/stage.go index bc6866f..4ebbaa3 100644 --- a/pipe/stage.go +++ b/pipe/stage.go @@ -100,7 +100,7 @@ type Stage interface { // // If `Start()` returns without an error, `Wait()` must also be // called, to allow all resources to be freed. - Start(ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser) error + Start(ctx context.Context, env Env, stdin io.ReadCloser, stdout io.WriteCloser, opts StartOptions) error // Wait waits for the stage to be done, either because it has // finished or because it has been killed due to the expiration of @@ -108,6 +108,20 @@ type Stage interface { Wait() error } +// StartOptions carries run-scoped options passed to `Stage.Start`. +// It is a struct (rather than positional parameters) so that future +// options can be added without breaking the `Stage` interface. +type StartOptions struct { + // PanicHandler, if non-nil, is invoked to recover a panic that + // escapes a Function stage's goroutine, converting it into an + // error. Stage types that don't run user code in a + // library-spawned goroutine ignore it. + PanicHandler StagePanicHandler +} + +// StagePanicHandler is a function that handles panics in a pipeline's stages. +type StagePanicHandler func(p any) error + // StagePreferences is the way that a `Stage` indicates its // preferences about how it is run. This is used within // `pipe.Pipeline` to decide when to use `os.Pipe()` vs. `io.Pipe()` From 4278c36aa605100186b6f14f9c5ad8de00b80413 Mon Sep 17 00:00:00 2001 From: Jason Lunz Date: Sat, 30 May 2026 18:49:04 +0200 Subject: [PATCH 23/23] Recover panics escaping the memory-watch goroutine Despite allowing a panic handler to be set, any panic in the user-supplied event handler run by memoryWatchStage in a library-spawned goroutine was not recovered. To cover that gap, pass opts.PanicHandler into monitor() and recover around the watch() call. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- pipe/memorylimit.go | 85 +++++++++++----- pipe/memorylimit_panic_test.go | 172 +++++++++++++++++++++++++++++++++ pipe/stage.go | 9 +- 3 files changed, 236 insertions(+), 30 deletions(-) create mode 100644 pipe/memorylimit_panic_test.go diff --git a/pipe/memorylimit.go b/pipe/memorylimit.go index 408f44d..df362c6 100644 --- a/pipe/memorylimit.go +++ b/pipe/memorylimit.go @@ -26,6 +26,11 @@ type LimitableStage interface { // MemoryLimit watches the memory usage of the stage and stops it if it // exceeds the given limit. +// +// If the event handler panics while reporting the over-limit event, the +// stage is still killed. A panic in any other event-handler call (an +// RSS-read error) is recovered via StartOptions.PanicHandler and the +// stage keeps running unmonitored; see StartOptions.PanicHandler. func MemoryLimit(stage Stage, byteLimit uint64, eventHandler func(e *Event)) Stage { limitableStage, ok := stage.(LimitableStage) @@ -73,16 +78,20 @@ func killAtLimit(byteLimit uint64, eventHandler func(e *Event)) memoryWatchFunc if rss < byteLimit { continue } - eventHandler(&Event{ - Command: stage.Name(), - Msg: "stage exceeded allowed memory use", - Err: fmt.Errorf("stage exceeded allowed memory use"), - Context: map[string]interface{}{ - "limit": byteLimit, - "used": rss, - }, - }) - stage.Kill(ErrMemoryLimitExceeded) + func() { + // Guarantee the over-limit stage is killed even if + // the user's event handler panics. + defer stage.Kill(ErrMemoryLimitExceeded) + eventHandler(&Event{ + Command: stage.Name(), + Msg: "stage exceeded allowed memory use", + Err: fmt.Errorf("stage exceeded allowed memory use"), + Context: map[string]interface{}{ + "limit": byteLimit, + "used": rss, + }, + }) + }() return } } @@ -93,6 +102,11 @@ func killAtLimit(byteLimit uint64, eventHandler func(e *Event)) memoryWatchFunc // one goroutine. It watches the memory usage of the stage, stops it // if it exceeds the given limit, and logs the peak memory usage when // the stage exits. +// +// Its event-handler panic behavior matches MemoryLimit: the over-limit +// kill always happens, while a panic in the RSS-error or peak-usage +// handler is recovered via StartOptions.PanicHandler and the stage keeps +// running unmonitored. See StartOptions.PanicHandler. func MemoryLimitWithObserver(stage Stage, byteLimit uint64, eventHandler func(e *Event)) Stage { limitableStage, ok := stage.(LimitableStage) if !ok { @@ -165,16 +179,20 @@ func killAtLimitAndObserve(byteLimit uint64, eventHandler func(e *Event)) memory } if rss >= byteLimit { - eventHandler(&Event{ - Command: stage.Name(), - Msg: "stage exceeded allowed memory use", - Err: fmt.Errorf("stage exceeded allowed memory use"), - Context: map[string]interface{}{ - "limit": byteLimit, - "used": rss, - }, - }) - stage.Kill(ErrMemoryLimitExceeded) + func() { + // Guarantee the over-limit stage is killed even if + // the user's event handler panics. + defer stage.Kill(ErrMemoryLimitExceeded) + eventHandler(&Event{ + Command: stage.Name(), + Msg: "stage exceeded allowed memory use", + Err: fmt.Errorf("stage exceeded allowed memory use"), + Context: map[string]interface{}{ + "limit": byteLimit, + "used": rss, + }, + }) + }() killed = true } } @@ -258,6 +276,7 @@ type memoryWatchStage struct { watch memoryWatchFunc cancel context.CancelFunc wg sync.WaitGroup + watchErr error } type memoryWatchFunc func(context.Context, LimitableStage) @@ -279,26 +298,40 @@ func (m *memoryWatchStage) Start( return err } - m.monitor(ctx) + m.monitor(ctx, opts.PanicHandler) return nil } -// monitor starts up a goroutine that monitors the memory of `m`. -func (m *memoryWatchStage) monitor(ctx context.Context) { +// monitor starts up a goroutine that monitors the memory of `m`. If +// panicHandler is set, any panic that escapes the user-supplied event handler +// (via m.watch) is recovered. +func (m *memoryWatchStage) monitor(ctx context.Context, panicHandler StagePanicHandler) { ctx, cancel := context.WithCancel(ctx) m.cancel = cancel m.wg.Add(1) go func() { + defer m.wg.Done() + defer func() { + if p := recover(); p != nil { + if panicHandler == nil { + panic(p) + } + m.watchErr = panicHandler(p) + } + }() m.watch(ctx, m.stage) - m.wg.Done() }() } func (m *memoryWatchStage) Wait() error { - defer m.stopWatching() - return m.stage.Wait() + err := m.stage.Wait() + m.stopWatching() + if err == nil { + err = m.watchErr // non-nil if panicHandler() returned anything + } + return err } func (m *memoryWatchStage) GetRSSAnon(ctx context.Context) (uint64, error) { diff --git a/pipe/memorylimit_panic_test.go b/pipe/memorylimit_panic_test.go new file mode 100644 index 0000000..da9977e --- /dev/null +++ b/pipe/memorylimit_panic_test.go @@ -0,0 +1,172 @@ +package pipe + +import ( + "context" + "fmt" + "io" + "os" + "os/exec" + "strings" + "testing" + "time" +) + +const memWatchPanicSentinel = "memwatch-panic-sentinel" +const memWatchPanicChildEnv = "GO_PIPE_MEMWATCH_PANIC_CHILD" + +// fakeLimitableStage is a minimal LimitableStage whose Wait returns +// immediately, letting a memoryWatchStage test exercise its watch +// goroutine in isolation. +type fakeLimitableStage struct{} + +func (fakeLimitableStage) Name() string { return "fake" } +func (fakeLimitableStage) Preferences() StagePreferences { return StagePreferences{} } +func (fakeLimitableStage) Start( + context.Context, Env, io.ReadCloser, io.WriteCloser, StartOptions, +) error { + return nil +} +func (fakeLimitableStage) Wait() error { return nil } +func (fakeLimitableStage) GetRSSAnon(context.Context) (uint64, error) { return 0, nil } +func (fakeLimitableStage) Kill(error) {} + +func panickingWatchStage() *memoryWatchStage { + return &memoryWatchStage{ + stage: fakeLimitableStage{}, + watch: func(context.Context, LimitableStage) { panic(memWatchPanicSentinel) }, + } +} + +// TestMemoryWatchStagePanicWithHandlerSurfaced verifies that a panic +// escaping the memory-watch goroutine (where the user-supplied event +// handler runs) is recovered via the configured panic handler and +// surfaced as the stage's Wait error. +func TestMemoryWatchStagePanicWithHandlerSurfaced(t *testing.T) { + ms := panickingWatchStage() + opts := StartOptions{ + PanicHandler: func(p any) error { return fmt.Errorf("recovered: %v", p) }, + } + + if err := ms.Start(context.Background(), Env{}, nil, nil, opts); err != nil { + t.Fatalf("Start returned unexpected error: %v", err) + } + + err := ms.Wait() + if err == nil { + t.Fatal("expected Wait to surface the recovered panic, got nil") + } + if !strings.Contains(err.Error(), memWatchPanicSentinel) { + t.Fatalf("expected error to mention %q, got: %v", memWatchPanicSentinel, err) + } +} + +// TestMemoryWatchStagePanicWithoutHandlerPropagates verifies that when +// the memory-watch goroutine panics and no panic handler is installed, +// the panic propagates (crashing the process) rather than being +// silently swallowed. Because that would crash the test binary, the +// scenario runs in a re-exec'd subprocess. +func TestMemoryWatchStagePanicWithoutHandlerPropagates(t *testing.T) { + if os.Getenv(memWatchPanicChildEnv) == "1" { + runMemWatchPanicChild() + return + } + + cmd := exec.Command(os.Args[0], "-test.run=^TestMemoryWatchStagePanicWithoutHandlerPropagates$", "-test.v") //nolint:gosec // re-exec of this test binary with constant arguments. + cmd.Env = append(os.Environ(), memWatchPanicChildEnv+"=1") + out, err := cmd.CombinedOutput() + output := string(out) + + if err == nil { + t.Fatalf("expected subprocess to crash from a propagated panic, but it exited 0\noutput:\n%s", output) + } + if strings.Contains(output, "SURVIVED") { + t.Fatalf("panic was swallowed: Wait returned instead of propagating\noutput:\n%s", output) + } + if !strings.Contains(output, "panic:") || !strings.Contains(output, memWatchPanicSentinel) { + t.Fatalf("expected a propagated panic mentioning %q, got:\n%s", memWatchPanicSentinel, output) + } +} + +func runMemWatchPanicChild() { + ms := panickingWatchStage() + + if err := ms.Start(context.Background(), Env{}, nil, nil, StartOptions{}); err != nil { + os.Stdout.WriteString("SURVIVED: Start returned err=" + err.Error() + "\n") + os.Exit(0) + } + + _ = ms.Wait() + + // Reaching this point at all indicates the panic was swallowed. + time.Sleep(2 * time.Second) + os.Stdout.WriteString("SURVIVED: Wait returned\n") + os.Exit(0) +} + +// killTrackingStage is a LimitableStage that reports an over-limit RSS +// and blocks in Wait until it is killed, recording that the kill +// happened. It lets a test assert that the memory limit is enforced. +type killTrackingStage struct { + killed chan struct{} + done chan struct{} +} + +func newKillTrackingStage() *killTrackingStage { + return &killTrackingStage{ + killed: make(chan struct{}), + done: make(chan struct{}), + } +} + +func (*killTrackingStage) Name() string { return "kill-tracking" } +func (*killTrackingStage) Preferences() StagePreferences { return StagePreferences{} } +func (*killTrackingStage) Start( + context.Context, Env, io.ReadCloser, io.WriteCloser, StartOptions, +) error { + return nil +} +func (s *killTrackingStage) Wait() error { <-s.done; return ErrMemoryLimitExceeded } +func (*killTrackingStage) GetRSSAnon(context.Context) (uint64, error) { + return 1 << 30, nil +} + +func (s *killTrackingStage) Kill(error) { + select { + case <-s.killed: + // already killed + default: + close(s.killed) + close(s.done) + } +} + +// TestMemoryLimitKillsEvenIfEventHandlerPanics verifies that an over-limit +// stage is still killed (the limit enforced) even when the user's event +// handler panics and that panic is recovered by the configured handler. +// Without the kill being guaranteed during unwinding, the runaway stage +// would never be killed and Wait would hang. +func TestMemoryLimitKillsEvenIfEventHandlerPanics(t *testing.T) { + stage := newKillTrackingStage() + ms := &memoryWatchStage{ + stage: stage, + watch: killAtLimit(1, func(*Event) { panic(memWatchPanicSentinel) }), + } + opts := StartOptions{ + PanicHandler: func(p any) error { return fmt.Errorf("recovered: %v", p) }, + } + + if err := ms.Start(context.Background(), Env{}, nil, nil, opts); err != nil { + t.Fatalf("Start returned unexpected error: %v", err) + } + + select { + case <-stage.killed: + // expected: the limit was enforced despite the handler panic. + case <-time.After(5 * time.Second): + t.Fatal("over-limit stage was not killed after the event handler panicked") + } + + if err := ms.Wait(); err != ErrMemoryLimitExceeded { + t.Fatalf("Wait = %v, want %v", err, ErrMemoryLimitExceeded) + } +} diff --git a/pipe/stage.go b/pipe/stage.go index 4ebbaa3..611c9dd 100644 --- a/pipe/stage.go +++ b/pipe/stage.go @@ -112,10 +112,11 @@ type Stage interface { // It is a struct (rather than positional parameters) so that future // options can be added without breaking the `Stage` interface. type StartOptions struct { - // PanicHandler, if non-nil, is invoked to recover a panic that - // escapes a Function stage's goroutine, converting it into an - // error. Stage types that don't run user code in a - // library-spawned goroutine ignore it. + // PanicHandler, if non-nil, is invoked to recover a panic that escapes + // user code that a stage runs in a library-spawned goroutine (a + // Function stage's StageFunc, or a memory-limit stage's event + // handler), converting it into an error. Stage types that don't run + // user code in a library-spawned goroutine ignore it. PanicHandler StagePanicHandler }