fix: reap idle chatd stream states on a timer (#24476)

* Adds `streamJanitorLoop` to clean up stale streams every 30s * zeroes dropped slots to aid in gc-eligibliity * Adds regression tests in coderd/x/chatd and enterprise/coderd/x/chatd > 🤖
2026-09-24 15:04:27 +08:00 · 2026-04-17 19:22:00 +01:00
parent ee563636ed
commit 3f6b40a833
3 changed files with 554 additions and 10 deletions
@@ -14,6 +14,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"

 	"charm.land/fantasy"
@@ -100,6 +101,11 @@ const (
 	// snapshot the buffer before it is garbage-collected.
 	bufferRetainGracePeriod = 5 * time.Second

+	// streamJanitorInterval is how often sweepIdleStreams runs.
+	// Worst-case retention is bufferRetainGracePeriod +
+	// streamJanitorInterval.
+	streamJanitorInterval = 30 * time.Second
+
 	// DefaultMaxChatsPerAcquire is the maximum number of chats to
 	// acquire in a single processOnce call. Batching avoids
 	// waiting a full polling interval between acquisitions
@@ -2902,6 +2908,8 @@ func (p *Server) start(ctx context.Context) {
 	// Single heartbeat loop for all chats on this replica.
 	go p.heartbeatLoop(ctx)

+	go p.streamJanitorLoop(ctx)
+
 	acquireTicker := p.clock.NewTicker(
 		p.pendingChatAcquireInterval,
 		"chatd",
@@ -3025,6 +3033,10 @@ func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEven
 				state.bufferDropCount = 0
 				state.bufferLastWarnAt = now
 			}
+			// Zero the dropped slot so its *ChatStreamMessagePart is
+			// GC-eligible; the later append reuses this slot in place
+			// whenever cap > len.
+			state.buffer[0] = codersdk.ChatStreamEvent{}
 			state.buffer = state.buffer[1:]
 		}
 		state.buffer = append(state.buffer, event)
@@ -3078,6 +3090,9 @@ func (p *Server) cacheDurableMessage(chatID uuid.UUID, event codersdk.ChatStream
 		if evicted := state.durableMessages[0]; evicted.Message != nil {
 			state.durableEvictedBefore = evicted.Message.ID
 		}
+		// Zero the dropped slot so the evicted *ChatMessage is
+		// GC-eligible; see publishToStream for the same pattern.
+		state.durableMessages[0] = codersdk.ChatStreamEvent{}
 		state.durableMessages = state.durableMessages[1:]
 	}
 	state.durableMessages = append(state.durableMessages, event)
@@ -3151,25 +3166,96 @@ func (p *Server) getOrCreateStreamState(chatID uuid.UUID) *chatStreamState {
 	return state
 }

-// cleanupStreamIfIdle removes the chat entry from the sync.Map
-// when there are no subscribers and the stream is not buffering.
-// When bufferRetainedAt is set, cleanup is deferred until the
-// grace period expires so cross-replica relay subscribers can
-// still snapshot the buffer.
-// The caller must hold state.mu.
-func (p *Server) cleanupStreamIfIdle(chatID uuid.UUID, state *chatStreamState) {
+// cleanupStreamIfIdle removes the chat entry from the sync.Map when
+// there are no subscribers, the stream is not buffering, and any
+// grace period for late-connecting relay subscribers has elapsed. If
+// the grace window is still open it returns without rescheduling.
+// streamJanitorLoop is the backstop that re-checks on a timer.
+//
+// The caller must hold state.mu. The state pointer may have been
+// captured outside this lock (sync.Map.Load or Range); we use
+// CompareAndDelete so a stale pointer cannot evict a fresh entry
+// installed by a racing getOrCreateStreamState. Returns true
+// if the state was deleted, false otherwise.
+func (p *Server) cleanupStreamIfIdle(chatID uuid.UUID, state *chatStreamState) bool {
 	if state.buffering || len(state.subscribers) > 0 {
-		return
+		return false
 	}
 	// Keep stream state alive during the grace period so
 	// late-connecting relay subscribers can snapshot the
 	// buffer after the worker finishes processing.
 	if !state.bufferRetainedAt.IsZero() &&
 		p.clock.Now().Before(state.bufferRetainedAt.Add(bufferRetainGracePeriod)) {
-		return
+		return false
+	}
+	if !p.chatStreams.CompareAndDelete(chatID, state) {
+		return false
 	}
-	p.chatStreams.Delete(chatID)
 	p.workspaceMCPToolsCache.Delete(chatID)
+	return true
+}
+
+// streamJanitorLoop periodically reaps idle chat stream states whose
+// grace period has expired. It is the backstop for the grace-window
+// early-return in cleanupStreamIfIdle; without it, a subscriber that
+// detaches inside grace (the common enterprise relay-drain case,
+// relayDrainTimeout = 200ms vs. 5s grace) pins the state forever.
+func (p *Server) streamJanitorLoop(ctx context.Context) {
+	ticker := p.clock.NewTicker(streamJanitorInterval, "chatd", "stream-janitor")
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			p.safeSweepIdleStreams(ctx)
+		}
+	}
+}
+
+// safeSweepIdleStreams runs sweepIdleStreams under a panic recovery
+// so an unexpected panic in the sweep cannot kill the janitor
+// goroutine and silently reintroduce the very leak it exists to
+// prevent. The next tick retries.
+func (p *Server) safeSweepIdleStreams(ctx context.Context) {
+	defer func() {
+		if r := recover(); r != nil {
+			p.logger.Error(ctx, "stream janitor sweep panicked, will retry next tick",
+				slog.F("panic", r))
+		}
+	}()
+	p.sweepIdleStreams()
+}
+
+// sweepIdleStreams iterates chatStreams once and delegates each entry
+// to cleanupStreamIfIdle. Range may skip entries that become reapable
+// concurrently. Any such entry is reaped on the next tick.
+func (p *Server) sweepIdleStreams() {
+	var reaped atomic.Int64
+	defer func() {
+		if count := reaped.Load(); count > 0 {
+			p.logger.Info(context.Background(), "reaped idle chat streams", slog.F("count", count))
+		}
+	}()
+	p.chatStreams.Range(func(key, value any) bool {
+		chatID, ok := key.(uuid.UUID)
+		if !ok {
+			return true
+		}
+		state, ok := value.(*chatStreamState)
+		if !ok {
+			return true
+		}
+		// guard against any panic from cleanupStreamIfIdle locking state.mu for all time
+		func() {
+			state.mu.Lock()
+			defer state.mu.Unlock()
+			if p.cleanupStreamIfIdle(chatID, state) {
+				reaped.Add(1)
+			}
+		}()
+		return true
+	})
 }

 // registerHeartbeat enrolls a chat in the centralized batch
@@ -3046,3 +3046,319 @@ func TestHeartbeatTick_DBErrorDoesNotInterruptChats(t *testing.T) {
 	require.NoError(t, chatCtx.Err(),
 		"chat context should not be canceled on transient DB error")
 }
+
+// TestSubscribeCancelDuringGrace_ReapedBySweep verifies that a
+// subscriber detach inside bufferRetainGracePeriod (the OSS trigger
+// for the retained-buffer leak) leaves the state mapped, and the
+// next sweep past the grace window reaps it.
+func TestSubscribeCancelDuringGrace_ReapedBySweep(t *testing.T) {
+	t.Parallel()
+
+	logger := slogtest.Make(t, nil)
+	mClock := quartz.NewMock(t)
+
+	server := &Server{
+		logger: logger,
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+	start := mClock.Now()
+
+	// Just-finished chat: processing done, buffer retained for
+	// late-connecting relay subscribers.
+	state := &chatStreamState{
+		buffering:        false,
+		bufferRetainedAt: start,
+		subscribers:      map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+		buffer: []codersdk.ChatStreamEvent{{
+			Type: codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{
+				Role: codersdk.ChatMessageRoleAssistant,
+			},
+		}},
+	}
+	server.chatStreams.Store(chatID, state)
+
+	// Real subscribeToStream cancel path: the WS subscriber detach
+	// that leaks in prod.
+	_, _, cancelSub := server.subscribeToStream(chatID)
+
+	mClock.Advance(bufferRetainGracePeriod / 2)
+	cancelSub()
+
+	_, ok := server.chatStreams.Load(chatID)
+	require.True(t, ok,
+		"entry should remain during grace window after subscriber detach")
+
+	mClock.Advance(bufferRetainGracePeriod)
+	server.sweepIdleStreams()
+
+	_, ok = server.chatStreams.Load(chatID)
+	require.False(t, ok,
+		"entry should be reaped after grace period expires and sweep runs")
+}
+
+// TestSweepIdleStreams_ReapsStaleRetainedBuffer: grace expired, no
+// subscribers, not buffering -> reaped.
+func TestSweepIdleStreams_ReapsStaleRetainedBuffer(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+	state := &chatStreamState{
+		buffering:        false,
+		bufferRetainedAt: mClock.Now(),
+		subscribers:      map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+		buffer: []codersdk.ChatStreamEvent{{
+			Type:        codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{},
+		}},
+	}
+	server.chatStreams.Store(chatID, state)
+
+	mClock.Advance(bufferRetainGracePeriod + time.Second)
+	server.sweepIdleStreams()
+
+	_, ok := server.chatStreams.Load(chatID)
+	require.False(t, ok, "stale retained state should be reaped")
+}
+
+// TestSweepIdleStreams_DoesNotReapActiveBuffering: buffering=true
+// blocks reap even long after any grace would have expired.
+func TestSweepIdleStreams_DoesNotReapActiveBuffering(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+	state := &chatStreamState{
+		buffering:   true,
+		subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+		buffer: []codersdk.ChatStreamEvent{{
+			Type:        codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{},
+		}},
+	}
+	server.chatStreams.Store(chatID, state)
+
+	mClock.Advance(time.Hour)
+	server.sweepIdleStreams()
+
+	_, ok := server.chatStreams.Load(chatID)
+	require.True(t, ok, "actively-buffering state must not be reaped")
+}
+
+// TestSweepIdleStreams_DoesNotReapWithSubscribers: attached
+// subscribers block reap even when grace has expired.
+func TestSweepIdleStreams_DoesNotReapWithSubscribers(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+	state := &chatStreamState{
+		buffering:        false,
+		bufferRetainedAt: mClock.Now(),
+		subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{
+			uuid.New(): make(chan codersdk.ChatStreamEvent, 1),
+		},
+		buffer: []codersdk.ChatStreamEvent{{
+			Type:        codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{},
+		}},
+	}
+	server.chatStreams.Store(chatID, state)
+
+	mClock.Advance(bufferRetainGracePeriod + time.Second)
+	server.sweepIdleStreams()
+
+	_, ok := server.chatStreams.Load(chatID)
+	require.True(t, ok, "state with subscribers must not be reaped")
+}
+
+// TestSweepIdleStreams_DefersDuringGracePeriod: sweep inside grace
+// is a no-op; the next sweep past grace reaps.
+func TestSweepIdleStreams_DefersDuringGracePeriod(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+	start := mClock.Now()
+	state := &chatStreamState{
+		buffering:        false,
+		bufferRetainedAt: start,
+		subscribers:      map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+		buffer: []codersdk.ChatStreamEvent{{
+			Type:        codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{},
+		}},
+	}
+	server.chatStreams.Store(chatID, state)
+
+	mClock.Advance(bufferRetainGracePeriod / 2)
+	server.sweepIdleStreams()
+
+	_, ok := server.chatStreams.Load(chatID)
+	require.True(t, ok, "sweep inside grace window must not reap")
+
+	mClock.Advance(bufferRetainGracePeriod)
+	server.sweepIdleStreams()
+
+	_, ok = server.chatStreams.Load(chatID)
+	require.False(t, ok, "sweep after grace window must reap")
+}
+
+// TestPublishToStream_DropZeroesBackingSlot verifies that evicting
+// the oldest buffered event at capacity zeroes the dropped slot so
+// its *ChatStreamMessagePart becomes GC-eligible immediately.
+func TestPublishToStream_DropZeroesBackingSlot(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+
+	// Over-allocate by one so the post-drop append fits in place and
+	// exercises the backing-array reuse this test is checking.
+	buf := make([]codersdk.ChatStreamEvent, maxStreamBufferSize, maxStreamBufferSize+1)
+	for i := range buf {
+		buf[i] = codersdk.ChatStreamEvent{
+			Type:        codersdk.ChatStreamEventTypeMessagePart,
+			MessagePart: &codersdk.ChatStreamMessagePart{},
+		}
+	}
+	// Sentinel in slot 0 distinguishes "slot was zeroed" from "slot
+	// was overwritten by a later append".
+	sentinel := &codersdk.ChatStreamMessagePart{
+		Role: codersdk.ChatMessageRoleAssistant,
+	}
+	buf[0] = codersdk.ChatStreamEvent{
+		Type:        codersdk.ChatStreamEventTypeMessagePart,
+		MessagePart: sentinel,
+	}
+	// Alias over the full backing array so we can still observe slot
+	// 0 after publishToStream reslices state.buffer forward.
+	origBacking := buf[:cap(buf)]
+
+	state := &chatStreamState{
+		buffering:   true,
+		buffer:      buf,
+		subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+	}
+	server.chatStreams.Store(chatID, state)
+
+	newPart := &codersdk.ChatStreamMessagePart{
+		Role: codersdk.ChatMessageRoleAssistant,
+	}
+	server.publishToStream(chatID, codersdk.ChatStreamEvent{
+		Type:        codersdk.ChatStreamEventTypeMessagePart,
+		MessagePart: newPart,
+	})
+
+	require.Equal(t, codersdk.ChatStreamEvent{}, origBacking[0],
+		"dropped slot must be zero-valued so its *ChatStreamMessagePart "+
+			"is eligible for GC; got %+v", origBacking[0])
+
+	// Sanity-check the in-place append path the fix targets: if Go's
+	// growth policy ever makes this append reallocate, this fails
+	// loudly so the test author revisits the setup.
+	require.Same(t, newPart, origBacking[len(origBacking)-1].MessagePart,
+		"append must have landed in the original backing array; the "+
+			"zero-out invariant only matters when cap > len")
+}
+
+// TestCleanupStreamIfIdle_StalePointerDoesNotDeleteFreshEntry covers
+// the race where a caller holds a pointer to a no-longer-mapped
+// state (e.g. a janitor Range callback racing a fresh
+// getOrCreateStreamState) and would otherwise evict the fresh entry.
+// With CompareAndDelete in cleanupStreamIfIdle the stale delete is
+// a no-op.
+func TestCleanupStreamIfIdle_StalePointerDoesNotDeleteFreshEntry(t *testing.T) {
+	t.Parallel()
+
+	mClock := quartz.NewMock(t)
+	server := &Server{
+		logger: slogtest.Make(t, nil),
+		clock:  mClock,
+	}
+
+	chatID := uuid.New()
+
+	// Stale pointer: reapable (not buffering, no subscribers, grace
+	// expired) but no longer the map's live entry.
+	stale := &chatStreamState{
+		buffering:        false,
+		bufferRetainedAt: mClock.Now(),
+		subscribers:      map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+	}
+
+	// Fresh entry: the state getOrCreateStreamState would install
+	// after a racing processChat run. Actively buffering, so not
+	// reapable. Only this state is in the map.
+	fresh := &chatStreamState{
+		buffering:   true,
+		subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{},
+	}
+	server.chatStreams.Store(chatID, fresh)
+
+	mClock.Advance(bufferRetainGracePeriod + time.Second)
+
+	// Stale caller mirrors the janitor Range callback after the map
+	// entry has already been replaced.
+	stale.mu.Lock()
+	server.cleanupStreamIfIdle(chatID, stale)
+	stale.mu.Unlock()
+
+	got, ok := server.chatStreams.Load(chatID)
+	require.True(t, ok,
+		"fresh entry must remain mapped when cleanup is called with a stale pointer")
+	require.Same(t, fresh, got,
+		"cleanup must not replace the fresh entry with the stale one")
+}
+
+// TestSafeSweepIdleStreams_RecoversFromPanic verifies that an
+// unexpected panic inside sweepIdleStreams is recovered rather than
+// killing the janitor goroutine. Without this guard, a panic would
+// silently reintroduce the very leak the janitor exists to prevent.
+func TestSafeSweepIdleStreams_RecoversFromPanic(t *testing.T) {
+	t.Parallel()
+
+	server := &Server{
+		logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}),
+		clock:  quartz.NewMock(t),
+	}
+
+	chatID := uuid.New()
+	// A nil *chatStreamState passes the type assertion in sweepIdleStreams
+	// but panics on state.mu.Lock with a nil-pointer deref. Any future
+	// panic source in the sweep would trigger the same recovery path.
+	var nilState *chatStreamState
+	server.chatStreams.Store(chatID, nilState)
+
+	require.NotPanics(t, func() {
+		server.safeSweepIdleStreams(context.Background())
+	}, "safeSweepIdleStreams must recover panics so the janitor loop keeps running")
+}