executor: filter command output lines instead of buffering all of them (#365)

Replace the full-output bytes.Buffer capture in Local.Run and Remote.sshRun with a streaming line writer, add RunOpts.KeepLine so each caller retains only the lines it reads, drop the dead RunOpts.Verbose field, and compile the secret masking patterns once in MakeLogs.

Resolves #363.

Author: paskal

pkg/executor/dry.go

@@ -1,7 +1,6 @@
 package executor
 
 import (
-	"bytes"
 	"context"
 	"io"
 	"log"
@@ -21,12 +20,12 @@ func NewDry(logs Logs) *Dry {
 }
 
 // Run shows the command content, doesn't execute it
-func (ex *Dry) Run(_ context.Context, cmd string, _ *RunOpts) (out []string, err error) {
+func (ex *Dry) Run(_ context.Context, cmd string, opts *RunOpts) (out []string, err error) {
 	log.Printf("[DEBUG] run %s", cmd)
-	var stdoutBuf bytes.Buffer
-	mwr := io.MultiWriter(ex.logs.Out, &stdoutBuf)
+	capture := newLineCapture(opts)
+	mwr := io.MultiWriter(ex.logs.Out, capture)
 	mwr.Write([]byte(cmd)) // nolint
-	return splitOutputLines(stdoutBuf.String()), nil
+	return capture.result(), nil
 }
 
 // Upload doesn't actually upload, just prints the command

pkg/executor/dry_test.go

@@ -18,14 +18,14 @@ func TestDry_Run(t *testing.T) {
 	dry := NewDry(MakeLogs(true, false, nil))
 
 	t.Run("single line", func(t *testing.T) {
-		res, err := dry.Run(ctx, "ls -la /srv", &RunOpts{Verbose: true})
+		res, err := dry.Run(ctx, "ls -la /srv", nil)
 		require.NoError(t, err)
 		require.Len(t, res, 1)
 		require.Equal(t, "ls -la /srv", res[0])
 	})
 
 	t.Run("multi line with blank line", func(t *testing.T) {
-		res, err := dry.Run(ctx, "ls -la /srv\n\ndf -h\n", &RunOpts{Verbose: true})
+		res, err := dry.Run(ctx, "ls -la /srv\n\ndf -h\n", nil)
 		require.NoError(t, err)
 		require.Equal(t, []string{"ls -la /srv", "", "df -h"}, res)
 	})
@@ -143,7 +143,7 @@ func TestDry_RunLineOverScannerLimit(t *testing.T) {
 	logs.Out = logs.Out.WithWriter(&buf)
 
 	cmd := strings.Repeat("x", 100000)
-	res, err := NewDry(logs).Run(context.Background(), cmd, &RunOpts{Verbose: true})
+	res, err := NewDry(logs).Run(context.Background(), cmd, nil)
 	require.NoError(t, err)
 	require.Len(t, res, 1)
 	assert.Equal(t, cmd, res[0])

pkg/executor/executor.go

@@ -3,6 +3,7 @@
 package executor
 
 import (
+	"bytes"
 	"context"
 	"path"
 	"strings"
@@ -21,8 +22,12 @@ type Interface interface {
 }
 
 // RunOpts is a struct for run options.
+// Note it is not comparable, since KeepLine is a func.
 type RunOpts struct {
-	Verbose bool // print more info to primary stdout
+	// KeepLine decides which stdout lines Run returns; a nil predicate keeps every line.
+	// Every line still reaches the log in full, so this bounds only what the caller holds:
+	// a command printing a large log costs nothing extra when its output is not read.
+	KeepLine func(line string) bool
 }
 
 // UpDownOpts is a struct for upload and download options.
@@ -64,6 +69,73 @@ func splitOutputLines(s string) []string {
 	return res
 }
 
+// lineCapture splits everything written to it into lines exactly as splitOutputLines does, keeping
+// only the ones keep accepts. Executors pass it alongside the log writer, so retention scales with
+// what the caller reads rather than with what the command printed. A line of any length is handled,
+// there is no scanner token limit, and an unterminated final line is returned like a terminated one.
+type lineCapture struct {
+	keep    func(line string) bool
+	partial []byte
+	lines   []string
+}
+
+func (lc *lineCapture) Write(p []byte) (n int, err error) {
+	n = len(p)
+	for {
+		i := bytes.IndexByte(p, '\n')
+		if i < 0 {
+			lc.partial = append(lc.partial, p...)
+			return n, nil
+		}
+		if len(lc.partial) == 0 {
+			lc.take(p[:i]) // whole line arrived in this write, no need to stage it
+		} else {
+			lc.partial = append(lc.partial, p[:i]...)
+			lc.take(lc.partial)
+			lc.reset()
+		}
+		p = p[i+1:]
+	}
+}
+
+// reset readies the staging buffer for the next line. A buffer that stayed small is kept, sparing an
+// allocation per line, but a large one is released: holding it would pin the longest line for the rest
+// of the command, which is the retention this type exists to avoid.
+func (lc *lineCapture) reset() {
+	const maxStageReuse = 64 << 10
+	if cap(lc.partial) > maxStageReuse {
+		lc.partial = nil
+		return
+	}
+	lc.partial = lc.partial[:0]
+}
+
+// take converts one line and keeps it if the predicate accepts. The string conversion copies, so
+// the retained line never pins the writer's buffer, and a rejected line is garbage at once.
+func (lc *lineCapture) take(b []byte) {
+	line := string(bytes.TrimSuffix(b, []byte("\r")))
+	if lc.keep == nil || lc.keep(line) {
+		lc.lines = append(lc.lines, line)
+	}
+}
+
+// result flushes an unterminated final line and returns the kept lines, nil when nothing was kept.
+func (lc *lineCapture) result() []string {
+	if len(lc.partial) > 0 {
+		lc.take(lc.partial)
+		lc.partial = nil
+	}
+	return lc.lines
+}
+
+// newLineCapture builds a capture honoring opts, which may be nil.
+func newLineCapture(opts *RunOpts) *lineCapture {
+	if opts == nil {
+		return &lineCapture{}
+	}
+	return &lineCapture{keep: opts.KeepLine}
+}
+
 // isExcluded reports whether fpath matches any of the exclude patterns. A pattern ending in "/*"
 // also matches the directory it names, so the whole subtree is protected, but this directory match
 // only applies when fpath is itself a directory. This prevents a pattern like "dir*/*" from

pkg/executor/executor_test.go

@@ -2,8 +2,10 @@ package executor
 
 import (
 	"bytes"
+	"fmt"
 	"io"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
@@ -145,3 +147,92 @@ func TestSplitOutputLines(t *testing.T) {
 		})
 	}
 }
+
+func TestLineCapture(t *testing.T) {
+	tbl := []struct {
+		name string
+		in   string
+		res  []string
+	}{
+		{"empty", "", nil},
+		{"single line, no trailing newline", "hello", []string{"hello"}},
+		{"single line with trailing newline", "hello\n", []string{"hello"}},
+		{"multiple lines", "line1\nline2\nline3\n", []string{"line1", "line2", "line3"}},
+		{"blank line in the middle", "line1\n\nline2\n", []string{"line1", "", "line2"}},
+		{"single newline", "\n", []string{""}},
+		{"trailing blank line", "line1\n\n", []string{"line1", ""}},
+		{"crlf line endings", "line1\r\nline2\r\n", []string{"line1", "line2"}},
+	}
+
+	// a nil predicate reproduces splitOutputLines whatever the write boundaries are, including a
+	// chunk size that splits a line, a crlf pair or a run of newlines
+	for _, tt := range tbl {
+		for _, chunk := range []int{0, 1, 2, 3, 7} {
+			t.Run(fmt.Sprintf("%s/chunk %d", tt.name, chunk), func(t *testing.T) {
+				lc := &lineCapture{}
+				writeInChunks(t, lc, tt.in, chunk)
+				assert.Equal(t, tt.res, lc.result())
+				assert.Equal(t, splitOutputLines(tt.in), lc.result(), "result is idempotent and matches the batch split")
+			})
+		}
+	}
+}
+
+func TestLineCaptureKeepLine(t *testing.T) {
+	t.Run("keeps only accepted lines", func(t *testing.T) {
+		lc := &lineCapture{keep: func(line string) bool { return strings.HasPrefix(line, "setvar ") }}
+		writeInChunks(t, lc, "noise\nsetvar a=1\nmore noise\nsetvar b=2\n", 3)
+		assert.Equal(t, []string{"setvar a=1", "setvar b=2"}, lc.result())
+	})
+
+	t.Run("rejecting everything retains nothing", func(t *testing.T) {
+		lc := &lineCapture{keep: func(string) bool { return false }}
+		writeInChunks(t, lc, "one\ntwo\nthree", 0)
+		assert.Nil(t, lc.result())
+	})
+
+	t.Run("predicate sees an unterminated final line", func(t *testing.T) {
+		var seen []string
+		lc := &lineCapture{keep: func(line string) bool { seen = append(seen, line); return true }}
+		writeInChunks(t, lc, "first\nlast-no-newline", 4)
+		assert.Equal(t, []string{"first"}, seen, "the final line is only offered once result is called")
+		assert.Equal(t, []string{"first", "last-no-newline"}, lc.result())
+	})
+
+	t.Run("line longer than a scanner token limit survives", func(t *testing.T) {
+		long := strings.Repeat("x", 1<<20)
+		lc := &lineCapture{}
+		writeInChunks(t, lc, long+"\nshort\n", 4096)
+		require.Len(t, lc.result(), 2)
+		assert.Equal(t, long, lc.result()[0])
+		assert.Equal(t, "short", lc.result()[1])
+	})
+}
+
+// writeInChunks feeds s to w in fixed-size pieces, or in one write when size is 0, so a test can
+// pin behavior across the write boundaries a real command produces.
+func writeInChunks(t *testing.T, w io.Writer, s string, size int) {
+	t.Helper()
+	if size <= 0 {
+		_, err := w.Write([]byte(s))
+		require.NoError(t, err)
+		return
+	}
+	for i := 0; i < len(s); i += size {
+		end := min(i+size, len(s))
+		_, err := w.Write([]byte(s[i:end]))
+		require.NoError(t, err)
+	}
+}
+
+func TestLineCaptureReleasesLargeStagingBuffer(t *testing.T) {
+	// a long line arriving in pieces has to be staged, but once it is consumed the buffer must go:
+	// keeping it would pin the longest line for as long as the command runs, which is the retention
+	// this type exists to remove
+	lc := &lineCapture{keep: func(string) bool { return false }}
+	writeInChunks(t, lc, strings.Repeat("x", 4<<20)+"\n", 4096)
+	assert.LessOrEqual(t, cap(lc.partial), 64<<10, "the staging buffer is released once the line is consumed")
+
+	writeInChunks(t, lc, "short\n", 2) // still usable for the output that follows
+	assert.Nil(t, lc.result())
+}

pkg/executor/local.go

@@ -1,7 +1,6 @@
 package executor
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -27,7 +26,7 @@ func NewLocal(logs Logs) *Local {
 }
 
 // Run executes command on local hostAddr, inside the shell
-func (l *Local) Run(ctx context.Context, cmd string, _ *RunOpts) (out []string, err error) {
+func (l *Local) Run(ctx context.Context, cmd string, opts *RunOpts) (out []string, err error) {
 	shell := func() string {
 		if strings.HasPrefix(cmd, "sh -c") {
 			return "sh" // command has sh -c prefix, so use sh
@@ -51,15 +50,15 @@ func (l *Local) Run(ctx context.Context, cmd string, _ *RunOpts) (out []string,
 	errLog := l.logs.Err.WithHost("localhost", "")
 	outLog.Write([]byte(cmd)) // nolint
 
-	var stdoutBuf bytes.Buffer
-	mwr := io.MultiWriter(outLog, &stdoutBuf)
+	capture := newLineCapture(opts)
+	mwr := io.MultiWriter(outLog, capture)
 	command.Stdout, command.Stderr = mwr, errLog
 	err = command.Run()
 	if err != nil {
 		return nil, err
 	}
 
-	return splitOutputLines(stdoutBuf.String()), nil
+	return capture.result(), nil
 }
 
 // Upload just copy file from one place to another

pkg/executor/local_test.go

@@ -21,19 +21,19 @@ func TestRun(t *testing.T) {
 	l := NewLocal(MakeLogs(true, false, nil))
 
 	t.Run("single line out, success", func(t *testing.T) {
-		out, e := l.Run(ctx, "echo 'hello world'", &RunOpts{Verbose: true})
+		out, e := l.Run(ctx, "echo 'hello world'", nil)
 		require.NoError(t, e)
 		assert.Equal(t, []string{"hello world"}, out)
 	})
 
 	t.Run("single line with sh -c, success", func(t *testing.T) {
-		out, e := l.Run(ctx, "sh -c echo hello world", &RunOpts{Verbose: true})
+		out, e := l.Run(ctx, "sh -c echo hello world", nil)
 		require.NoError(t, e)
 		assert.Equal(t, []string{"hello world"}, out)
 	})
 
 	t.Run("single line with sh -c with single quotes, success", func(t *testing.T) {
-		out, e := l.Run(ctx, "sh -c 'echo hello world'", &RunOpts{Verbose: true})
+		out, e := l.Run(ctx, "sh -c 'echo hello world'", nil)
 		require.NoError(t, e)
 		assert.Equal(t, []string{"hello world"}, out)
 	})
@@ -45,11 +45,11 @@ func TestRun(t *testing.T) {
 
 	t.Run("multi line out success", func(t *testing.T) {
 		// prepare the test environment
-		_, err := l.Run(ctx, "mkdir -p /tmp/st", &RunOpts{Verbose: true})
+		_, err := l.Run(ctx, "mkdir -p /tmp/st", nil)
 		require.NoError(t, err)
-		_, err = l.Run(ctx, "cp testdata/data1.txt /tmp/st/data1.txt", &RunOpts{Verbose: true})
+		_, err = l.Run(ctx, "cp testdata/data1.txt /tmp/st/data1.txt", nil)
 		require.NoError(t, err)
-		_, err = l.Run(ctx, "cp testdata/data2.txt /tmp/st/data2.txt", &RunOpts{Verbose: true})
+		_, err = l.Run(ctx, "cp testdata/data2.txt /tmp/st/data2.txt", nil)
 		require.NoError(t, err)
 
 		out, err := l.Run(ctx, "ls -1 /tmp/st", nil)
@@ -65,7 +65,7 @@ func TestRun(t *testing.T) {
 	})
 
 	t.Run("find out", func(t *testing.T) {
-		out, e := l.Run(ctx, "find /tmp/st -type f", &RunOpts{Verbose: true})
+		out, e := l.Run(ctx, "find /tmp/st -type f", nil)
 		require.NoError(t, e)
 		slices.Sort(out)
 		assert.Contains(t, out, "/tmp/st/data1.txt")
@@ -86,7 +86,7 @@ func TestRun(t *testing.T) {
 	t.Run("with secrets", func(t *testing.T) {
 		stdout := captureStdOut(t, func() {
 			l := NewLocal(MakeLogs(true, false, []string{"data2"}))
-			out, e := l.Run(ctx, "find /tmp/st -type f", &RunOpts{Verbose: true})
+			out, e := l.Run(ctx, "find /tmp/st -type f", nil)
 			require.NoError(t, e)
 			slices.Sort(out)
 			assert.Equal(t, []string{"/tmp/st/data1.txt", "/tmp/st/data2.txt"}, out)