diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 904ca1cc79..3b658f433a 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -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 diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index bf088e118c..0a7620d8b1 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -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") +} diff --git a/enterprise/coderd/x/chatd/chatd_test.go b/enterprise/coderd/x/chatd/chatd_test.go index d8b9595830..b959d7002f 100644 --- a/enterprise/coderd/x/chatd/chatd_test.go +++ b/enterprise/coderd/x/chatd/chatd_test.go @@ -1425,6 +1425,148 @@ func TestSubscribeRelayDialCanceledOnFastCompletion(t *testing.T) { "worker completes before the relay is established") } +// TestSubscribeRelayDrainWithinGraceLeavesBufferRetained characterizes +// the multi-replica trigger for the retained-buffer leak: an enterprise +// relay drain (relayDrainTimeout = 200ms) always fires inside the +// worker's 5s grace window, so the worker-side subscriber-detach hits +// cleanupStreamIfIdle's early-return and the buffer stays mapped. +// streamJanitorLoop is the timer-driven backstop. +// +// The assertion is behavioral (a fresh worker.Subscribe sees the +// retained message_parts) rather than a chatStreams-size check because +// _test.go identifiers in coderd/x/chatd do not link into the +// enterprise test binary, and adding a production accessor for this +// isn't justified. The matching reap assertion lives in the OSS unit +// tests in coderd/x/chatd/chatd_internal_test.go. +func TestSubscribeRelayDrainWithinGraceLeavesBufferRetained(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + workerID := uuid.New() + subscriberID := uuid.New() + + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("relay-drain-characterization") + } + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks("hello ", "from ", "worker")..., + ) + }) + + workerLogger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + // Freeze the worker's clock so streamJanitorLoop cannot race the + // buffer-retained assertion on slow CI. + workerClock := quartz.NewMock(t) + worker := osschatd.New(osschatd.Config{ + Logger: workerLogger, + Database: db, + ReplicaID: workerID, + Pubsub: ps, + PendingChatAcquireInterval: time.Hour, + InFlightChatStaleAfter: testutil.WaitSuperLong, + Clock: workerClock, + }) + t.Cleanup(func() { + require.NoError(t, worker.Close()) + }) + + // Subscriber dials through to the worker. On cancel the relay + // drain fires well inside the worker's 5s grace, exercising the + // cleanupStreamIfIdle early-return path. + subscriber := newTestServer(t, db, ps, subscriberID, func( + ctx context.Context, + chatID uuid.UUID, + targetWorkerID uuid.UUID, + requestHeader http.Header, + ) ( + []codersdk.ChatStreamEvent, + <-chan codersdk.ChatStreamEvent, + func(), + error, + ) { + snapshot, relayEvents, cancel, ok := worker.Subscribe(ctx, chatID, requestHeader, math.MaxInt64) + if !ok { + return nil, nil, nil, xerrors.New("worker subscribe failed") + } + return snapshot, relayEvents, cancel, nil + }, nil) + + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(ctx, t, db) + setOpenAIProviderBaseURL(ctx, t, db, openAIURL) + + chat := seedWaitingChat(ctx, t, db, org.ID, user, model, "relay-drain-characterization") + + // Attach before processing so the relay opens as soon as + // status=running arrives. + _, events, subCancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) + require.True(t, ok) + + _, err := worker.SendMessage(ctx, osschatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, + }) + require.NoError(t, err) + + // Drain events until processing has clearly completed: we need + // the assistant message and at least one message_part so we know + // processChat's defer has flipped buffering=false and populated + // bufferRetainedAt before the subscriber detaches. + var committedAssistantMsgs int + var messagePartsSeen int + testutil.Eventually(ctx, t, func(context.Context) bool { + select { + case event := <-events: + switch event.Type { + case codersdk.ChatStreamEventTypeMessagePart: + messagePartsSeen++ + case codersdk.ChatStreamEventTypeMessage: + if event.Message != nil && event.Message.Role == codersdk.ChatMessageRoleAssistant { + committedAssistantMsgs++ + } + } + return committedAssistantMsgs > 0 && messagePartsSeen > 0 + default: + return false + } + }, testutil.IntervalFast) + + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + fromDB, dbErr := db.GetChatByID(ctx, chat.ID) + if dbErr != nil { + return false + } + return fromDB.Status == database.ChatStatusWaiting + }, testutil.IntervalFast) + + // Tear the subscriber down inside the worker's grace window. + subCancel() + + // A fresh worker.Subscribe still sees the retained + // message_parts: the buffer was not reaped when the relay + // drained. Eventually absorbs the short window before the + // worker observes the teardown. The retry itself re-enters + // cleanupStreamIfIdle via its own cancel defer but still + // early-returns because grace is still open. + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + snap, _, snapCancel, ok := worker.Subscribe(ctx, chat.ID, nil, math.MaxInt64) + if !ok { + return false + } + defer snapCancel() + for _, e := range snap { + if e.Type == codersdk.ChatStreamEventTypeMessagePart { + return true + } + } + return false + }, testutil.IntervalFast, + "retained buffer must still contain message_parts after the "+ + "relay drains within grace") +} + // TestSubscribeRelayEstablishedMidStream demonstrates that when the // relay is established while the worker is still streaming, the // subscriber receives buffered parts via the relay snapshot and live