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>

Author: znull

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 {

pipe/memorylimit_test.go

@@ -113,6 +113,22 @@ 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 testMemoryLimit(t *testing.T, mbs int, limit uint64, stage pipe.Stage) (string, error) {
 	ctx := context.Background()
 
@@ -147,3 +163,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
+}