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>

Author: znull

pipe/memorylimit.go

@@ -297,11 +297,9 @@ func (m *memoryWatchStage) monitor(ctx context.Context) {
 }
 
 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

@@ -129,6 +129,14 @@ func TestMemoryLimitWithObserverTreeMem(t *testing.T) {
 	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 testMemoryLimit(t *testing.T, mbs int, limit uint64, stage pipe.Stage) (string, error) {
 	ctx := context.Background()