diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index fa6d1b2398..aae2ce2aa3 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -5782,15 +5782,15 @@ func (q *querier) UpdateChatByID(ctx context.Context, arg database.UpdateChatByI return q.db.UpdateChatByID(ctx, arg) } -func (q *querier) UpdateChatHeartbeat(ctx context.Context, arg database.UpdateChatHeartbeatParams) (int64, error) { - chat, err := q.db.GetChatByID(ctx, arg.ID) - if err != nil { - return 0, err +func (q *querier) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) { + // The batch heartbeat is a system-level operation filtered by + // worker_id. Authorization is enforced by the AsChatd context + // at the call site rather than per-row, because checking each + // row individually would defeat the purpose of batching. + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil { + return nil, err } - if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { - return 0, err - } - return q.db.UpdateChatHeartbeat(ctx, arg) + return q.db.UpdateChatHeartbeats(ctx, arg) } func (q *querier) UpdateChatLabelsByID(ctx context.Context, arg database.UpdateChatLabelsByIDParams) (database.Chat, error) { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 22648aa6a5..01449b4f06 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -842,15 +842,15 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().UpdateChatStatusPreserveUpdatedAt(gomock.Any(), arg).Return(chat, nil).AnyTimes() check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(chat) })) - s.Run("UpdateChatHeartbeat", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { - chat := testutil.Fake(s.T(), faker, database.Chat{}) - arg := database.UpdateChatHeartbeatParams{ - ID: chat.ID, + s.Run("UpdateChatHeartbeats", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + resultID := uuid.New() + arg := database.UpdateChatHeartbeatsParams{ + IDs: []uuid.UUID{resultID}, WorkerID: uuid.New(), + Now: time.Now(), } - dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() - dbm.EXPECT().UpdateChatHeartbeat(gomock.Any(), arg).Return(int64(1), nil).AnyTimes() - check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(int64(1)) + dbm.EXPECT().UpdateChatHeartbeats(gomock.Any(), arg).Return([]uuid.UUID{resultID}, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]uuid.UUID{resultID}) })) s.Run("UpdateChatMessageByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 1c4e5955b0..f91b6f6efa 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -4136,11 +4136,11 @@ func (m queryMetricsStore) UpdateChatByID(ctx context.Context, arg database.Upda return r0, r1 } -func (m queryMetricsStore) UpdateChatHeartbeat(ctx context.Context, arg database.UpdateChatHeartbeatParams) (int64, error) { +func (m queryMetricsStore) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) { start := time.Now() - r0, r1 := m.s.UpdateChatHeartbeat(ctx, arg) - m.queryLatencies.WithLabelValues("UpdateChatHeartbeat").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpdateChatHeartbeat").Inc() + r0, r1 := m.s.UpdateChatHeartbeats(ctx, arg) + m.queryLatencies.WithLabelValues("UpdateChatHeartbeats").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpdateChatHeartbeats").Inc() return r0, r1 } diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index dcc142214e..9aa4a08bb4 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -7835,19 +7835,19 @@ func (mr *MockStoreMockRecorder) UpdateChatByID(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatByID", reflect.TypeOf((*MockStore)(nil).UpdateChatByID), ctx, arg) } -// UpdateChatHeartbeat mocks base method. -func (m *MockStore) UpdateChatHeartbeat(ctx context.Context, arg database.UpdateChatHeartbeatParams) (int64, error) { +// UpdateChatHeartbeats mocks base method. +func (m *MockStore) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpdateChatHeartbeat", ctx, arg) - ret0, _ := ret[0].(int64) + ret := m.ctrl.Call(m, "UpdateChatHeartbeats", ctx, arg) + ret0, _ := ret[0].([]uuid.UUID) ret1, _ := ret[1].(error) return ret0, ret1 } -// UpdateChatHeartbeat indicates an expected call of UpdateChatHeartbeat. -func (mr *MockStoreMockRecorder) UpdateChatHeartbeat(ctx, arg any) *gomock.Call { +// UpdateChatHeartbeats indicates an expected call of UpdateChatHeartbeats. +func (mr *MockStoreMockRecorder) UpdateChatHeartbeats(ctx, arg any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatHeartbeat", reflect.TypeOf((*MockStore)(nil).UpdateChatHeartbeat), ctx, arg) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatHeartbeats", reflect.TypeOf((*MockStore)(nil).UpdateChatHeartbeats), ctx, arg) } // UpdateChatLabelsByID mocks base method. diff --git a/coderd/database/querier.go b/coderd/database/querier.go index fc442a56f3..6719cb8862 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -870,9 +870,11 @@ type sqlcQuerier interface { UpdateAPIKeyByID(ctx context.Context, arg UpdateAPIKeyByIDParams) error UpdateChatBuildAgentBinding(ctx context.Context, arg UpdateChatBuildAgentBindingParams) (Chat, error) UpdateChatByID(ctx context.Context, arg UpdateChatByIDParams) (Chat, error) - // Bumps the heartbeat timestamp for a running chat so that other - // replicas know the worker is still alive. - UpdateChatHeartbeat(ctx context.Context, arg UpdateChatHeartbeatParams) (int64, error) + // Bumps the heartbeat timestamp for the given set of chat IDs, + // provided they are still running and owned by the specified + // worker. Returns the IDs that were actually updated so the + // caller can detect stolen or completed chats via set-difference. + UpdateChatHeartbeats(ctx context.Context, arg UpdateChatHeartbeatsParams) ([]uuid.UUID, error) UpdateChatLabelsByID(ctx context.Context, arg UpdateChatLabelsByIDParams) (Chat, error) // Updates the cached injected context parts (AGENTS.md + // skills) on the chat row. Called only when context changes diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 7e1c7f0009..4549830418 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6615,30 +6615,49 @@ func (q *sqlQuerier) UpdateChatByID(ctx context.Context, arg UpdateChatByIDParam return i, err } -const updateChatHeartbeat = `-- name: UpdateChatHeartbeat :execrows +const updateChatHeartbeats = `-- name: UpdateChatHeartbeats :many UPDATE chats SET - heartbeat_at = NOW() + heartbeat_at = $1::timestamptz WHERE - id = $1::uuid - AND worker_id = $2::uuid + id = ANY($2::uuid[]) + AND worker_id = $3::uuid AND status = 'running'::chat_status +RETURNING id ` -type UpdateChatHeartbeatParams struct { - ID uuid.UUID `db:"id" json:"id"` - WorkerID uuid.UUID `db:"worker_id" json:"worker_id"` +type UpdateChatHeartbeatsParams struct { + Now time.Time `db:"now" json:"now"` + IDs []uuid.UUID `db:"ids" json:"ids"` + WorkerID uuid.UUID `db:"worker_id" json:"worker_id"` } -// Bumps the heartbeat timestamp for a running chat so that other -// replicas know the worker is still alive. -func (q *sqlQuerier) UpdateChatHeartbeat(ctx context.Context, arg UpdateChatHeartbeatParams) (int64, error) { - result, err := q.db.ExecContext(ctx, updateChatHeartbeat, arg.ID, arg.WorkerID) +// Bumps the heartbeat timestamp for the given set of chat IDs, +// provided they are still running and owned by the specified +// worker. Returns the IDs that were actually updated so the +// caller can detect stolen or completed chats via set-difference. +func (q *sqlQuerier) UpdateChatHeartbeats(ctx context.Context, arg UpdateChatHeartbeatsParams) ([]uuid.UUID, error) { + rows, err := q.db.QueryContext(ctx, updateChatHeartbeats, arg.Now, pq.Array(arg.IDs), arg.WorkerID) if err != nil { - return 0, err + return nil, err } - return result.RowsAffected() + defer rows.Close() + var items []uuid.UUID + for rows.Next() { + var id uuid.UUID + if err := rows.Scan(&id); err != nil { + return nil, err + } + items = append(items, id) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil } const updateChatLabelsByID = `-- name: UpdateChatLabelsByID :one diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index b29d766b07..31d606d70c 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -674,17 +674,20 @@ WHERE status = 'running'::chat_status AND heartbeat_at < @stale_threshold::timestamptz; --- name: UpdateChatHeartbeat :execrows --- Bumps the heartbeat timestamp for a running chat so that other --- replicas know the worker is still alive. +-- name: UpdateChatHeartbeats :many +-- Bumps the heartbeat timestamp for the given set of chat IDs, +-- provided they are still running and owned by the specified +-- worker. Returns the IDs that were actually updated so the +-- caller can detect stolen or completed chats via set-difference. UPDATE chats SET - heartbeat_at = NOW() + heartbeat_at = @now::timestamptz WHERE - id = @id::uuid + id = ANY(@ids::uuid[]) AND worker_id = @worker_id::uuid - AND status = 'running'::chat_status; + AND status = 'running'::chat_status +RETURNING id; -- name: GetChatDiffStatusByChatID :one SELECT diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 0219cad2a7..45a0fdd46b 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -7,6 +7,7 @@ import ( "encoding/json" "errors" "fmt" + "maps" "net/http" "slices" "strconv" @@ -151,6 +152,12 @@ type Server struct { inFlightChatStaleAfter time.Duration chatHeartbeatInterval time.Duration + // heartbeatMu guards heartbeatRegistry. + heartbeatMu sync.Mutex + // heartbeatRegistry maps chat IDs to their cancel functions + // and workspace state for the centralized heartbeat loop. + heartbeatRegistry map[uuid.UUID]*heartbeatEntry + // wakeCh is signaled by SendMessage, EditMessage, CreateChat, // and PromoteQueued so the run loop calls processOnce // immediately instead of waiting for the next ticker. @@ -706,6 +713,17 @@ type chatStreamState struct { bufferRetainedAt time.Time } +// heartbeatEntry tracks a single chat's cancel function and workspace +// state for the centralized heartbeat loop. Instead of spawning a +// per-chat goroutine, processChat registers an entry here and the +// single heartbeatLoop goroutine handles all chats. +type heartbeatEntry struct { + cancelWithCause context.CancelCauseFunc + chatID uuid.UUID + workspaceID uuid.NullUUID + logger slog.Logger +} + // resetDropCounters zeroes the rate-limiting state for both buffer // and subscriber drop warnings. The caller must hold s.mu. func (s *chatStreamState) resetDropCounters() { @@ -2420,8 +2438,8 @@ func New(cfg Config) *Server { clock: clk, recordingSem: make(chan struct{}, maxConcurrentRecordingUploads), wakeCh: make(chan struct{}, 1), + heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), } - //nolint:gocritic // The chat processor uses a scoped chatd context. ctx = dbauthz.AsChatd(ctx) @@ -2461,6 +2479,9 @@ func (p *Server) start(ctx context.Context) { // to handle chats orphaned by crashed or redeployed workers. p.recoverStaleChats(ctx) + // Single heartbeat loop for all chats on this replica. + go p.heartbeatLoop(ctx) + acquireTicker := p.clock.NewTicker( p.pendingChatAcquireInterval, "chatd", @@ -2730,6 +2751,97 @@ func (p *Server) cleanupStreamIfIdle(chatID uuid.UUID, state *chatStreamState) { p.workspaceMCPToolsCache.Delete(chatID) } +// registerHeartbeat enrolls a chat in the centralized batch +// heartbeat loop. Must be called after chatCtx is created. +func (p *Server) registerHeartbeat(entry *heartbeatEntry) { + p.heartbeatMu.Lock() + defer p.heartbeatMu.Unlock() + if _, exists := p.heartbeatRegistry[entry.chatID]; exists { + p.logger.Warn(context.Background(), + "duplicate heartbeat registration, skipping", + slog.F("chat_id", entry.chatID)) + return + } + p.heartbeatRegistry[entry.chatID] = entry +} + +// unregisterHeartbeat removes a chat from the centralized +// heartbeat loop when chat processing finishes. +func (p *Server) unregisterHeartbeat(chatID uuid.UUID) { + p.heartbeatMu.Lock() + defer p.heartbeatMu.Unlock() + delete(p.heartbeatRegistry, chatID) +} + +// heartbeatLoop runs in a single goroutine, issuing one batch +// heartbeat query per interval for all registered chats. +func (p *Server) heartbeatLoop(ctx context.Context) { + ticker := p.clock.NewTicker(p.chatHeartbeatInterval, "chatd", "batch-heartbeat") + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + p.heartbeatTick(ctx) + } + } +} + +// heartbeatTick issues a single batch UPDATE for all running chats +// owned by this worker. Chats missing from the result set are +// interrupted (stolen by another replica or already completed). +func (p *Server) heartbeatTick(ctx context.Context) { + // Snapshot the registry under the lock. + p.heartbeatMu.Lock() + snapshot := maps.Clone(p.heartbeatRegistry) + p.heartbeatMu.Unlock() + + if len(snapshot) == 0 { + return + } + + // Collect the IDs we believe we own. + ids := slices.Collect(maps.Keys(snapshot)) + + //nolint:gocritic // AsChatd provides narrowly-scoped daemon + // access for batch-updating heartbeats. + chatdCtx := dbauthz.AsChatd(ctx) + updatedIDs, err := p.db.UpdateChatHeartbeats(chatdCtx, database.UpdateChatHeartbeatsParams{ + IDs: ids, + WorkerID: p.workerID, + Now: p.clock.Now(), + }) + if err != nil { + p.logger.Error(ctx, "batch heartbeat failed", slog.Error(err)) + return + } + + // Build a set of IDs that were successfully updated. + updated := make(map[uuid.UUID]struct{}, len(updatedIDs)) + for _, id := range updatedIDs { + updated[id] = struct{}{} + } + + // Interrupt registered chats that were not in the result + // (stolen by another replica or already completed). + for id, entry := range snapshot { + if _, ok := updated[id]; !ok { + entry.logger.Warn(ctx, "chat not in batch heartbeat result, interrupting") + entry.cancelWithCause(chatloop.ErrInterrupted) + continue + } + // Bump workspace usage for surviving chats. + newWsID := p.trackWorkspaceUsage(ctx, entry.chatID, entry.workspaceID, entry.logger) + // Update workspace ID in the registry for next tick. + p.heartbeatMu.Lock() + if current, exists := p.heartbeatRegistry[id]; exists { + current.workspaceID = newWsID + } + p.heartbeatMu.Unlock() + } +} + func (p *Server) Subscribe( ctx context.Context, chatID uuid.UUID, @@ -3575,33 +3687,17 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { } }() - // Periodically update the heartbeat so other replicas know this - // worker is still alive. The goroutine stops when chatCtx is - // canceled (either by completion or interruption). - go func() { - ticker := p.clock.NewTicker(p.chatHeartbeatInterval, "chatd", "heartbeat") - defer ticker.Stop() - for { - select { - case <-chatCtx.Done(): - return - case <-ticker.C: - rows, err := p.db.UpdateChatHeartbeat(chatCtx, database.UpdateChatHeartbeatParams{ - ID: chat.ID, - WorkerID: p.workerID, - }) - if err != nil { - logger.Warn(chatCtx, "failed to update chat heartbeat", slog.Error(err)) - continue - } - if rows == 0 { - cancel(chatloop.ErrInterrupted) - return - } - chat.WorkspaceID = p.trackWorkspaceUsage(chatCtx, chat.ID, chat.WorkspaceID, logger) - } - } - }() + // Register with the centralized heartbeat loop instead of + // running a per-chat goroutine. The loop issues a single batch + // UPDATE for all chats on this worker and detects stolen chats + // via set-difference. + p.registerHeartbeat(&heartbeatEntry{ + cancelWithCause: cancel, + chatID: chat.ID, + workspaceID: chat.WorkspaceID, + logger: logger, + }) + defer p.unregisterHeartbeat(chat.ID) // Start buffering stream events BEFORE publishing the running // status. This closes a race where a subscriber sees diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 2f9e90dabe..a256200d3f 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -21,6 +21,7 @@ import ( dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/x/chatd/chaterror" + "github.com/coder/coder/v2/coderd/x/chatd/chatloop" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" "github.com/coder/coder/v2/coderd/x/chatd/chattest" "github.com/coder/coder/v2/coderd/x/chatd/chattool" @@ -2071,6 +2072,7 @@ func TestProcessChat_IgnoresStaleControlNotification(t *testing.T) { workerID: workerID, chatHeartbeatInterval: time.Minute, configCache: newChatConfigCache(ctx, db, clock), + heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), } // Publish a stale "pending" notification on the control channel @@ -2133,3 +2135,130 @@ func TestProcessChat_IgnoresStaleControlNotification(t *testing.T) { require.Equal(t, database.ChatStatusError, finalStatus, "processChat should have reached runChat (error), not been interrupted (waiting)") } + +// TestHeartbeatTick_StolenChatIsInterrupted verifies that when the +// batch heartbeat UPDATE does not return a registered chat's ID +// (because another replica stole it or it was completed), the +// heartbeat tick cancels that chat's context with ErrInterrupted +// while leaving surviving chats untouched. +func TestHeartbeatTick_StolenChatIsInterrupted(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + clock := quartz.NewMock(t) + + workerID := uuid.New() + + server := &Server{ + db: db, + logger: logger, + clock: clock, + workerID: workerID, + chatHeartbeatInterval: time.Minute, + heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), + } + + // Create three chats with independent cancel functions. + chat1 := uuid.New() + chat2 := uuid.New() + chat3 := uuid.New() + + _, cancel1 := context.WithCancelCause(ctx) + _, cancel2 := context.WithCancelCause(ctx) + ctx3, cancel3 := context.WithCancelCause(ctx) + + server.registerHeartbeat(&heartbeatEntry{ + cancelWithCause: cancel1, + chatID: chat1, + logger: logger, + }) + server.registerHeartbeat(&heartbeatEntry{ + cancelWithCause: cancel2, + chatID: chat2, + logger: logger, + }) + server.registerHeartbeat(&heartbeatEntry{ + cancelWithCause: cancel3, + chatID: chat3, + logger: logger, + }) + + // The batch UPDATE returns only chat1 and chat2 — + // chat3 was "stolen" by another replica. + db.EXPECT().UpdateChatHeartbeats(gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, params database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) { + require.Equal(t, workerID, params.WorkerID) + require.Len(t, params.IDs, 3) + // Return only chat1 and chat2 as surviving. + return []uuid.UUID{chat1, chat2}, nil + }, + ) + + server.heartbeatTick(ctx) + + // chat3's context should be canceled with ErrInterrupted. + require.ErrorIs(t, context.Cause(ctx3), chatloop.ErrInterrupted, + "stolen chat should be interrupted") + + // chat3 should have been removed from the registry by + // unregister (in production this happens via defer in + // processChat). The heartbeat tick itself does not + // unregister — it only cancels. Verify the entry is + // still present (processChat's defer would clean it up). + server.heartbeatMu.Lock() + _, chat1Exists := server.heartbeatRegistry[chat1] + _, chat2Exists := server.heartbeatRegistry[chat2] + _, chat3Exists := server.heartbeatRegistry[chat3] + server.heartbeatMu.Unlock() + + require.True(t, chat1Exists, "surviving chat1 should remain registered") + require.True(t, chat2Exists, "surviving chat2 should remain registered") + require.True(t, chat3Exists, + "stolen chat3 should still be in registry (processChat defer removes it)") +} + +// TestHeartbeatTick_DBErrorDoesNotInterruptChats verifies that a +// transient database failure causes the tick to log and return +// without canceling any registered chats. +func TestHeartbeatTick_DBErrorDoesNotInterruptChats(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + clock := quartz.NewMock(t) + + server := &Server{ + db: db, + logger: logger, + clock: clock, + workerID: uuid.New(), + chatHeartbeatInterval: time.Minute, + heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), + } + + chatID := uuid.New() + chatCtx, cancel := context.WithCancelCause(ctx) + + server.registerHeartbeat(&heartbeatEntry{ + cancelWithCause: cancel, + chatID: chatID, + logger: logger, + }) + + // Simulate a transient DB error. + db.EXPECT().UpdateChatHeartbeats(gomock.Any(), gomock.Any()).Return( + nil, xerrors.New("connection reset"), + ) + + server.heartbeatTick(ctx) + + // Chat should NOT be interrupted — the tick logged and + // returned early. + require.NoError(t, chatCtx.Err(), + "chat context should not be canceled on transient DB error") +} diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index fec70a60c8..c8388a7054 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -474,7 +474,7 @@ func TestArchiveChatInterruptsActiveProcessing(t *testing.T) { require.Equal(t, 1, userMessages, "expected queued message to stay queued after archive") } -func TestUpdateChatHeartbeatRequiresOwnership(t *testing.T) { +func TestUpdateChatHeartbeatsRequiresOwnership(t *testing.T) { t.Parallel() db, ps := dbtestutil.NewDB(t) @@ -501,19 +501,24 @@ func TestUpdateChatHeartbeatRequiresOwnership(t *testing.T) { }) require.NoError(t, err) - rows, err := db.UpdateChatHeartbeat(ctx, database.UpdateChatHeartbeatParams{ - ID: chat.ID, + // Wrong worker_id should return no IDs. + ids, err := db.UpdateChatHeartbeats(ctx, database.UpdateChatHeartbeatsParams{ + IDs: []uuid.UUID{chat.ID}, WorkerID: uuid.New(), + Now: time.Now(), }) require.NoError(t, err) - require.Equal(t, int64(0), rows) + require.Empty(t, ids) - rows, err = db.UpdateChatHeartbeat(ctx, database.UpdateChatHeartbeatParams{ - ID: chat.ID, + // Correct worker_id should return the chat's ID. + ids, err = db.UpdateChatHeartbeats(ctx, database.UpdateChatHeartbeatsParams{ + IDs: []uuid.UUID{chat.ID}, WorkerID: workerID, + Now: time.Now(), }) require.NoError(t, err) - require.Equal(t, int64(1), rows) + require.Len(t, ids, 1) + require.Equal(t, chat.ID, ids[0]) } func TestSendMessageQueueBehaviorQueuesWhenBusy(t *testing.T) {