Add MemoryLimitWithObserver to combine limit and observer in one goroutine

MemoryLimitWithObserver provides the same functionality as composing
MemoryLimit(MemoryObserver(stage, h), limit, h) but uses a single
goroutine instead of two. The combined watcher polls GetRSSAnon once
per second, enforces the byte limit, tracks peak RSS, and logs peak
memory usage on exit.

In gitrpcd, each upload-pack request wraps its git subprocess with
MemoryLimit(MemoryObserver(...)), spawning two goroutines that both
poll the same PID's RSS every second. During a goroutine spike on
github-dfs-dfee0d0 (2026-03-20 ~08:04 UTC), this contributed ~415
goroutines out of 4,430 total. Using the combined API eliminates
~200 of those goroutines.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: torstenwalter

pipe/memorylimit.go

@@ -89,6 +89,102 @@ 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 {
@@ -194,11 +290,9 @@ func (m *memoryWatchStage) Start(ctx context.Context, env Env, stdin io.ReadClos
 }
 
 func (m *memoryWatchStage) Wait() error {
-	if err := m.stage.Wait(); err != nil {
-		return err
-	}
+	err := m.stage.Wait()
 	m.stopWatching()
-	return nil
+	return err
 }
 
 func (m *memoryWatchStage) GetRSSAnon(ctx context.Context) (uint64, error) {

pipe/memorylimit_test.go

@@ -121,6 +121,117 @@ 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()