diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index 313709681f..513bbba1c8 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -41,7 +41,7 @@ const ( homeInstructionLookupTimeout = 5 * time.Second instructionCacheTTL = 5 * time.Minute - chatHeartbeatInterval = 30 * time.Second + chatHeartbeatInterval = 60 * time.Second maxChatSteps = 1200 // maxStreamBufferSize caps the number of events buffered // per chat during a single LLM step. When exceeded the @@ -53,6 +53,12 @@ const ( // of 5 means recovery runs at 1/5 of the stale-after duration. staleRecoveryIntervalDivisor = 5 + // maxChatsPerAcquire is the maximum number of chats to + // acquire in a single processOnce call. Batching avoids + // waiting a full polling interval between acquisitions + // when many chats are pending. + maxChatsPerAcquire int32 = 10 + defaultSubagentInstruction = "You are running as a delegated sub-agent chat. Complete the delegated task and provide clear, concise assistant responses for the parent agent." ) @@ -75,10 +81,10 @@ type Server struct { webpushDispatcher webpush.Dispatcher providerAPIKeys chatprovider.ProviderAPIKeys - // streamMu guards chatStreams which tracks in-flight chat - // stream state for broadcasting ephemeral events. - streamMu sync.Mutex - chatStreams map[uuid.UUID]*chatStreamState + // chatStreams stores per-chat stream state. Using sync.Map + // gives each chat independent locking — concurrent chats + // never contend with each other. + chatStreams sync.Map // uuid.UUID -> *chatStreamState // instructionCache caches home instruction file contents by // workspace agent ID so we don't re-dial on every chat turn. @@ -137,6 +143,7 @@ type SubscribeFnParams struct { } type chatStreamState struct { + mu sync.Mutex buffer []codersdk.ChatStreamEvent buffering bool subscribers map[uuid.UUID]chan codersdk.ChatStreamEvent @@ -989,7 +996,6 @@ func New(cfg Config) *Server { pubsub: cfg.Pubsub, webpushDispatcher: cfg.WebpushDispatcher, providerAPIKeys: cfg.ProviderAPIKeys, - chatStreams: make(map[uuid.UUID]*chatStreamState), instructionCache: make(map[uuid.UUID]cachedInstruction), pendingChatAcquireInterval: pendingChatAcquireInterval, inFlightChatStaleAfter: inFlightChatStaleAfter, @@ -1029,79 +1035,73 @@ func (p *Server) start(ctx context.Context) { } func (p *Server) processOnce(ctx context.Context) { - // Bail out early if the server is shutting down. The main - // loop's select can randomly pick the ticker over ctx.Done(), - // so we must guard against acquiring a chat we cannot process. if ctx.Err() != nil { return } - // Try to acquire a pending chat. We detach from the server - // lifetime to prevent a phantom-acquire race: when the server - // context is canceled, the pq driver's watchCancel goroutine - // races with the actual query on the wire. The UPDATE can - // commit in Postgres (setting the chat to "running") before - // the cancel request arrives via a second TCP connection, yet - // the Go driver still returns context.Canceled to the caller - // because the awaitDone goroutine in database/sql closes the - // Rows before Scan reads them. This leaves the chat stuck as - // "running" with no goroutine to process it. Using a context - // that cannot be canceled ensures the driver sees the query - // result if Postgres executed it. + // We detach from the server lifetime to prevent a + // phantom-acquire race: when the server context is + // canceled, the pq driver's watchCancel goroutine + // races with the actual query on the wire. Using a + // context that cannot be canceled ensures the driver + // sees the query result if Postgres executed it. acquireCtx, acquireCancel := context.WithTimeout( context.WithoutCancel(ctx), 10*time.Second, ) - defer acquireCancel() - chat, err := p.db.AcquireChat(acquireCtx, database.AcquireChatParams{ + chats, err := p.db.AcquireChats(acquireCtx, database.AcquireChatsParams{ StartedAt: time.Now(), WorkerID: p.workerID, + NumChats: maxChatsPerAcquire, }) + acquireCancel() if err != nil { - if !xerrors.Is(err, sql.ErrNoRows) { - p.logger.Error(ctx, "failed to acquire chat", slog.Error(err)) - } - // No pending chats or error. + p.logger.Error(ctx, "failed to acquire chats", slog.Error(err)) + return + } + if len(chats) == 0 { return } - // If the server context was canceled while we were acquiring, - // release the chat back to pending immediately so another - // replica can pick it up. + // If the server context was canceled while we were + // acquiring, release the chats back to pending. if ctx.Err() != nil { releaseCtx, releaseCancel := context.WithTimeout( context.WithoutCancel(ctx), 10*time.Second, ) - defer releaseCancel() - _, updateErr := p.db.UpdateChatStatus(releaseCtx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusPending, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: sql.NullString{}, - }) - if updateErr != nil { - p.logger.Error(ctx, "failed to release chat acquired during shutdown", - slog.F("chat_id", chat.ID), slog.Error(updateErr)) + for _, chat := range chats { + _, updateErr := p.db.UpdateChatStatus(releaseCtx, database.UpdateChatStatusParams{ + ID: chat.ID, + Status: database.ChatStatusPending, + WorkerID: uuid.NullUUID{}, + StartedAt: sql.NullTime{}, + HeartbeatAt: sql.NullTime{}, + LastError: sql.NullString{}, + }) + if updateErr != nil { + p.logger.Error(ctx, "failed to release chat acquired during shutdown", + slog.F("chat_id", chat.ID), slog.Error(updateErr)) + } } + releaseCancel() return } - // Process the chat (don't block the main loop). - p.inflight.Add(1) - go func() { - defer p.inflight.Done() - p.processChat(ctx, chat) - }() + for _, chat := range chats { + p.inflight.Add(1) + go func() { + defer p.inflight.Done() + p.processChat(ctx, chat) + }() + } } func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEvent) { - p.streamMu.Lock() - state := p.streamStateLocked(chatID) + state := p.getOrCreateStreamState(chatID) + state.mu.Lock() if event.Type == codersdk.ChatStreamEventTypeMessagePart { if !state.buffering { - p.cleanupStreamIfIdleLocked(chatID, state) - p.streamMu.Unlock() + p.cleanupStreamIfIdle(chatID, state) + state.mu.Unlock() return } if len(state.buffer) >= maxStreamBufferSize { @@ -1115,7 +1115,7 @@ func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEven for _, ch := range state.subscribers { subscribers = append(subscribers, ch) } - p.streamMu.Unlock() + state.mu.Unlock() for _, ch := range subscribers { select { @@ -1127,13 +1127,11 @@ func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEven } // Clean up the stream entry if it was created by - // streamStateLocked but has no subscribers and is not + // getOrCreateStreamState but has no subscribers and is not // actively buffering (e.g. publish with no watchers). - p.streamMu.Lock() - if cur, ok := p.chatStreams[chatID]; ok { - p.cleanupStreamIfIdleLocked(chatID, cur) - } - p.streamMu.Unlock() + state.mu.Lock() + p.cleanupStreamIfIdle(chatID, state) + state.mu.Unlock() } func (p *Server) subscribeToStream(chatID uuid.UUID) ( @@ -1141,48 +1139,52 @@ func (p *Server) subscribeToStream(chatID uuid.UUID) ( <-chan codersdk.ChatStreamEvent, func(), ) { - p.streamMu.Lock() - state := p.streamStateLocked(chatID) + state := p.getOrCreateStreamState(chatID) + state.mu.Lock() snapshot := append([]codersdk.ChatStreamEvent(nil), state.buffer...) id := uuid.New() ch := make(chan codersdk.ChatStreamEvent, 128) state.subscribers[id] = ch - p.streamMu.Unlock() + state.mu.Unlock() cancel := func() { - p.streamMu.Lock() - state, ok := p.chatStreams[chatID] - if ok { - // Remove the subscriber but do not close the channel. - // publishToStream copies subscriber references under - // streamMu then sends outside the lock; closing here - // races with that send and can panic. The channel - // becomes unreachable once removed and will be GC'd. - delete(state.subscribers, id) - p.cleanupStreamIfIdleLocked(chatID, state) - } - p.streamMu.Unlock() + state.mu.Lock() + // Remove the subscriber but do not close the channel. + // publishToStream copies subscriber references under + // the per-chat lock then sends outside; closing here + // races with that send and can panic. The channel + // becomes unreachable once removed and will be GC'd. + delete(state.subscribers, id) + p.cleanupStreamIfIdle(chatID, state) + state.mu.Unlock() } return snapshot, ch, cancel } -// cleanupStreamIfIdleLocked removes the chat entry when there -// are no subscribers and the stream is not buffering. The -// caller must hold p.streamMu. -func (p *Server) cleanupStreamIfIdleLocked(chatID uuid.UUID, state *chatStreamState) { - if !state.buffering && len(state.subscribers) == 0 { - delete(p.chatStreams, chatID) +// getOrCreateStreamState returns the per-chat stream state, +// creating one atomically if it doesn't exist. The returned +// state has its own mutex — callers must lock state.mu for +// access. +func (p *Server) getOrCreateStreamState(chatID uuid.UUID) *chatStreamState { + if val, ok := p.chatStreams.Load(chatID); ok { + state, _ := val.(*chatStreamState) + return state } + val, _ := p.chatStreams.LoadOrStore(chatID, &chatStreamState{ + subscribers: make(map[uuid.UUID]chan codersdk.ChatStreamEvent), + }) + state, _ := val.(*chatStreamState) + return state } -func (p *Server) streamStateLocked(chatID uuid.UUID) *chatStreamState { - state, ok := p.chatStreams[chatID] - if !ok { - state = &chatStreamState{subscribers: make(map[uuid.UUID]chan codersdk.ChatStreamEvent)} - p.chatStreams[chatID] = state +// cleanupStreamIfIdle removes the chat entry from the sync.Map +// when there are no subscribers and the stream is not buffering. +// The caller must hold state.mu. +func (p *Server) cleanupStreamIfIdle(chatID uuid.UUID, state *chatStreamState) { + if !state.buffering && len(state.subscribers) == 0 { + p.chatStreams.Delete(chatID) } - return state } func (p *Server) Subscribe( @@ -1876,19 +1878,17 @@ func (p *Server) processChat(ctx context.Context, chat database.Chat) { // buffering hasn't started yet — the subscriber gets an empty // snapshot and publishToStream drops message_parts while // buffering is false. - p.streamMu.Lock() - startState := p.streamStateLocked(chat.ID) - startState.buffer = nil - startState.buffering = true - p.streamMu.Unlock() + streamState := p.getOrCreateStreamState(chat.ID) + streamState.mu.Lock() + streamState.buffer = nil + streamState.buffering = true + streamState.mu.Unlock() defer func() { - p.streamMu.Lock() - if stopState, ok := p.chatStreams[chat.ID]; ok { - stopState.buffer = nil - stopState.buffering = false - p.cleanupStreamIfIdleLocked(chat.ID, stopState) - } - p.streamMu.Unlock() + streamState.mu.Lock() + streamState.buffer = nil + streamState.buffering = false + p.cleanupStreamIfIdle(chat.ID, streamState) + streamState.mu.Unlock() }() p.publishStatus(chat.ID, database.ChatStatusRunning, uuid.NullUUID{ @@ -2373,11 +2373,13 @@ func (p *Server) runChat( // Clear the stream buffer now that the step is // persisted. Late-joining subscribers will load // these messages from the database instead. - p.streamMu.Lock() - if state, ok := p.chatStreams[chat.ID]; ok { - state.buffer = nil + if val, ok := p.chatStreams.Load(chat.ID); ok { + if ss, ok := val.(*chatStreamState); ok { + ss.mu.Lock() + ss.buffer = nil + ss.mu.Unlock() + } } - p.streamMu.Unlock() return nil } @@ -2531,12 +2533,13 @@ func (p *Server) runChat( }, OnRetry: func(attempt int, retryErr error, delay time.Duration) { - p.streamMu.Lock() - if state, ok := p.chatStreams[chat.ID]; ok { - state.buffer = nil + if val, ok := p.chatStreams.Load(chat.ID); ok { + if rs, ok := val.(*chatStreamState); ok { + rs.mu.Lock() + rs.buffer = nil + rs.mu.Unlock() + } } - p.streamMu.Unlock() - logger.Warn(ctx, "retrying LLM stream", slog.F("attempt", attempt), slog.F("delay", delay.String()), diff --git a/coderd/chatd/chatloop/chatloop.go b/coderd/chatd/chatloop/chatloop.go index aa4373775b..1beec43917 100644 --- a/coderd/chatd/chatloop/chatloop.go +++ b/coderd/chatd/chatloop/chatloop.go @@ -8,6 +8,7 @@ import ( "slices" "strconv" "strings" + "sync" "time" "charm.land/fantasy" @@ -575,9 +576,10 @@ func processStepStream( return result, nil } -// executeTools runs each tool call sequentially after the stream -// completes. Results are published via onResult as each tool -// finishes. +// executeTools runs all tool calls concurrently after the stream +// completes. Results are published via onResult in the original +// tool-call order after all tools finish, preserving deterministic +// event ordering for SSE subscribers. func executeTools( ctx context.Context, allTools []fantasy.AgentTool, @@ -593,11 +595,32 @@ func executeTools( toolMap[t.Info().Name] = t } - results := make([]fantasy.ToolResultContent, 0, len(toolCalls)) - for _, tc := range toolCalls { - tr := executeSingleTool(ctx, toolMap, tc) - results = append(results, tr) - if onResult != nil { + results := make([]fantasy.ToolResultContent, len(toolCalls)) + var wg sync.WaitGroup + wg.Add(len(toolCalls)) + for i, tc := range toolCalls { + go func(i int, tc fantasy.ToolCallContent) { + defer wg.Done() + defer func() { + if r := recover(); r != nil { + results[i] = fantasy.ToolResultContent{ + ToolCallID: tc.ToolCallID, + ToolName: tc.ToolName, + Result: fantasy.ToolResultOutputContentError{ + Error: xerrors.Errorf("tool panicked: %v", r), + }, + } + } + }() + results[i] = executeSingleTool(ctx, toolMap, tc) + }(i, tc) + } + wg.Wait() + + // Publish results in the original tool-call order so SSE + // subscribers see a deterministic event sequence. + if onResult != nil { + for _, tr := range results { onResult(tr) } } diff --git a/coderd/chatd/subagent.go b/coderd/chatd/subagent.go index 9c4f2ec99f..cbc390c94e 100644 --- a/coderd/chatd/subagent.go +++ b/coderd/chatd/subagent.go @@ -2,6 +2,7 @@ package chatd import ( "context" + "database/sql" "encoding/json" "sort" "strings" @@ -13,12 +14,14 @@ import ( "github.com/coder/coder/v2/coderd/chatd/chatprompt" "github.com/coder/coder/v2/coderd/database" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" ) var ErrSubagentNotDescendant = xerrors.New("target chat is not a descendant of current chat") const ( subagentAwaitPollInterval = 200 * time.Millisecond + subagentAwaitFallbackPoll = 5 * time.Second defaultSubagentWaitTimeout = 5 * time.Minute ) @@ -322,41 +325,90 @@ func (p *Server) awaitSubagentCompletion( return database.Chat{}, "", ErrSubagentNotDescendant } + // Check immediately before entering the poll loop. + targetChat, report, done, checkErr := p.checkSubagentCompletion(ctx, targetChatID) + if checkErr != nil { + return database.Chat{}, "", checkErr + } + if done { + return handleSubagentDone(targetChat, report) + } + if timeout <= 0 { timeout = defaultSubagentWaitTimeout } timer := time.NewTimer(timeout) defer timer.Stop() - ticker := time.NewTicker(subagentAwaitPollInterval) + // When pubsub is available, subscribe for fast status + // notifications and use a less aggressive fallback poll. + // Without pubsub (single-instance / in-memory) fall back + // to the original 200ms polling. + pollInterval := subagentAwaitPollInterval + var notifyCh <-chan struct{} + if p.pubsub != nil { + pollInterval = subagentAwaitFallbackPoll + ch := make(chan struct{}, 1) + notifyCh = ch + cancel, subErr := p.pubsub.SubscribeWithErr( + coderdpubsub.ChatStreamNotifyChannel(targetChatID), + func(_ context.Context, _ []byte, _ error) { + // Non-blocking send so we never stall the + // pubsub dispatch goroutine. + select { + case ch <- struct{}{}: + default: + } + }, + ) + if subErr == nil { + defer cancel() + } else { + // Subscription failed; fall back to fast polling. + pollInterval = subagentAwaitPollInterval + notifyCh = nil + } + } + + ticker := time.NewTicker(pollInterval) defer ticker.Stop() for { - targetChat, report, done, checkErr := p.checkSubagentCompletion(ctx, targetChatID) - if checkErr != nil { - return database.Chat{}, "", checkErr - } - if done { - if targetChat.Status == database.ChatStatusError { - reason := strings.TrimSpace(report) - if reason == "" { - reason = "agent reached error status" - } - return database.Chat{}, "", xerrors.New(reason) - } - return targetChat, report, nil - } - select { + case <-notifyCh: case <-ticker.C: case <-timer.C: return database.Chat{}, "", xerrors.New("timed out waiting for delegated subagent completion") case <-ctx.Done(): return database.Chat{}, "", ctx.Err() } + + targetChat, report, done, checkErr = p.checkSubagentCompletion(ctx, targetChatID) + if checkErr != nil { + return database.Chat{}, "", checkErr + } + if done { + return handleSubagentDone(targetChat, report) + } } } +// handleSubagentDone translates a completed subagent check into the +// appropriate return value, surfacing error-status chats as errors. +func handleSubagentDone( + chat database.Chat, + report string, +) (database.Chat, string, error) { + if chat.Status == database.ChatStatusError { + reason := strings.TrimSpace(report) + if reason == "" { + reason = "agent reached error status" + } + return database.Chat{}, "", xerrors.New(reason) + } + return chat, report, nil +} + func (p *Server) closeSubagent( ctx context.Context, parentChatID uuid.UUID, @@ -448,6 +500,9 @@ func latestSubagentAssistantMessage( return "", nil } +// isSubagentDescendant reports whether targetChatID is a descendant +// of ancestorChatID by walking up the parent chain from the target. +// This is O(depth) DB queries instead of O(nodes) BFS. func isSubagentDescendant( ctx context.Context, store database.Store, @@ -458,47 +513,29 @@ func isSubagentDescendant( return false, nil } - descendants, err := listSubagentDescendants(ctx, store, ancestorChatID) - if err != nil { - return false, err - } - for _, descendant := range descendants { - if descendant.ID == targetChatID { + currentID := targetChatID + visited := map[uuid.UUID]struct{}{} // cycle protection + for { + if _, seen := visited[currentID]; seen { + return false, nil + } + visited[currentID] = struct{}{} + + chat, err := store.GetChatByID(ctx, currentID) + if err != nil { + if xerrors.Is(err, sql.ErrNoRows) { + return false, nil // chain broken; not a confirmed descendant + } + return false, xerrors.Errorf("get chat %s: %w", currentID, err) + } + if !chat.ParentChatID.Valid { + return false, nil // reached root without finding ancestor + } + if chat.ParentChatID.UUID == ancestorChatID { return true, nil } + currentID = chat.ParentChatID.UUID } - return false, nil -} - -func listSubagentDescendants( - ctx context.Context, - store database.Store, - chatID uuid.UUID, -) ([]database.Chat, error) { - queue := []uuid.UUID{chatID} - visited := map[uuid.UUID]struct{}{chatID: {}} - - out := make([]database.Chat, 0) - for len(queue) > 0 { - parentChatID := queue[0] - queue = queue[1:] - - children, err := store.ListChildChatsByParentID(ctx, parentChatID) - if err != nil { - return nil, xerrors.Errorf("list child chats for %s: %w", parentChatID, err) - } - - for _, child := range children { - if _, ok := visited[child.ID]; ok { - continue - } - visited[child.ID] = struct{}{} - out = append(out, child) - queue = append(queue, child.ID) - } - } - - return out, nil } func subagentFallbackChatTitle(message string) string { diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index e48a455348..36f963abdb 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -1512,13 +1512,13 @@ func (q *querier) authorizeProvisionerJob(ctx context.Context, job database.Prov return nil } -func (q *querier) AcquireChat(ctx context.Context, arg database.AcquireChatParams) (database.Chat, error) { - // AcquireChat is a system-level operation used by the chat processor. +func (q *querier) AcquireChats(ctx context.Context, arg database.AcquireChatsParams) ([]database.Chat, error) { + // AcquireChats is a system-level operation used by the chat processor. // Authorization is done at the system level, not per-user. if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil { - return database.Chat{}, err + return nil, err } - return q.db.AcquireChat(ctx, arg) + return q.db.AcquireChats(ctx, arg) } func (q *querier) AcquireLock(ctx context.Context, id int64) error { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 6af3477f60..5a7e47c2af 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -373,14 +373,15 @@ func (s *MethodTestSuite) TestConnectionLogs() { } func (s *MethodTestSuite) TestChats() { - s.Run("AcquireChat", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { - arg := database.AcquireChatParams{ + s.Run("AcquireChats", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + arg := database.AcquireChatsParams{ StartedAt: dbtime.Now(), WorkerID: uuid.New(), + NumChats: 1, } chat := testutil.Fake(s.T(), faker, database.Chat{}) - dbm.EXPECT().AcquireChat(gomock.Any(), arg).Return(chat, nil).AnyTimes() - check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns(chat) + dbm.EXPECT().AcquireChats(gomock.Any(), arg).Return([]database.Chat{chat}, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.Chat{chat}) })) s.Run("DeleteAllChatQueuedMessages", 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 46adaeec89..27c0d88181 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -104,11 +104,11 @@ func (m queryMetricsStore) DeleteOrganization(ctx context.Context, id uuid.UUID) return r0 } -func (m queryMetricsStore) AcquireChat(ctx context.Context, arg database.AcquireChatParams) (database.Chat, error) { +func (m queryMetricsStore) AcquireChats(ctx context.Context, arg database.AcquireChatsParams) ([]database.Chat, error) { start := time.Now() - r0, r1 := m.s.AcquireChat(ctx, arg) - m.queryLatencies.WithLabelValues("AcquireChat").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "AcquireChat").Inc() + r0, r1 := m.s.AcquireChats(ctx, arg) + m.queryLatencies.WithLabelValues("AcquireChats").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "AcquireChats").Inc() return r0, r1 } diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index cd3ad0d843..618c63a770 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -44,19 +44,19 @@ func (m *MockStore) EXPECT() *MockStoreMockRecorder { return m.recorder } -// AcquireChat mocks base method. -func (m *MockStore) AcquireChat(ctx context.Context, arg database.AcquireChatParams) (database.Chat, error) { +// AcquireChats mocks base method. +func (m *MockStore) AcquireChats(ctx context.Context, arg database.AcquireChatsParams) ([]database.Chat, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "AcquireChat", ctx, arg) - ret0, _ := ret[0].(database.Chat) + ret := m.ctrl.Call(m, "AcquireChats", ctx, arg) + ret0, _ := ret[0].([]database.Chat) ret1, _ := ret[1].(error) return ret0, ret1 } -// AcquireChat indicates an expected call of AcquireChat. -func (mr *MockStoreMockRecorder) AcquireChat(ctx, arg any) *gomock.Call { +// AcquireChats indicates an expected call of AcquireChats. +func (mr *MockStoreMockRecorder) AcquireChats(ctx, arg any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AcquireChat", reflect.TypeOf((*MockStore)(nil).AcquireChat), ctx, arg) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "AcquireChats", reflect.TypeOf((*MockStore)(nil).AcquireChats), ctx, arg) } // AcquireLock mocks base method. diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 98964b54a5..32c883b0a6 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -12,9 +12,9 @@ import ( ) type sqlcQuerier interface { - // Acquires a pending chat for processing. Uses SKIP LOCKED to prevent - // multiple replicas from acquiring the same chat. - AcquireChat(ctx context.Context, arg AcquireChatParams) (Chat, error) + // Acquires up to @num_chats pending chats for processing. Uses SKIP LOCKED + // to prevent multiple replicas from acquiring the same chat. + AcquireChats(ctx context.Context, arg AcquireChatsParams) ([]Chat, error) // Blocks until the lock is acquired. // // This must be called from within a transaction. The lock will be automatically diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index a9f6269dc3..874e61ff5d 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2968,7 +2968,7 @@ func (q *sqlQuerier) UpdateChatProvider(ctx context.Context, arg UpdateChatProvi return i, err } -const acquireChat = `-- name: AcquireChat :one +const acquireChats = `-- name: AcquireChats :many UPDATE chats SET @@ -2978,7 +2978,7 @@ SET updated_at = $1::timestamptz, worker_id = $2::uuid WHERE - id = ( + id = ANY( SELECT id FROM @@ -2990,40 +2990,57 @@ WHERE FOR UPDATE SKIP LOCKED LIMIT - 1 + $3::int ) RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error ` -type AcquireChatParams struct { +type AcquireChatsParams struct { StartedAt time.Time `db:"started_at" json:"started_at"` WorkerID uuid.UUID `db:"worker_id" json:"worker_id"` + NumChats int32 `db:"num_chats" json:"num_chats"` } -// Acquires a pending chat for processing. Uses SKIP LOCKED to prevent -// multiple replicas from acquiring the same chat. -func (q *sqlQuerier) AcquireChat(ctx context.Context, arg AcquireChatParams) (Chat, error) { - row := q.db.QueryRowContext(ctx, acquireChat, arg.StartedAt, arg.WorkerID) - var i Chat - err := row.Scan( - &i.ID, - &i.OwnerID, - &i.WorkspaceID, - &i.Title, - &i.Status, - &i.WorkerID, - &i.StartedAt, - &i.HeartbeatAt, - &i.CreatedAt, - &i.UpdatedAt, - &i.ParentChatID, - &i.RootChatID, - &i.LastModelConfigID, - &i.Archived, - &i.LastError, - ) - return i, err +// Acquires up to @num_chats pending chats for processing. Uses SKIP LOCKED +// to prevent multiple replicas from acquiring the same chat. +func (q *sqlQuerier) AcquireChats(ctx context.Context, arg AcquireChatsParams) ([]Chat, error) { + rows, err := q.db.QueryContext(ctx, acquireChats, arg.StartedAt, arg.WorkerID, arg.NumChats) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Chat + for rows.Next() { + var i Chat + if err := rows.Scan( + &i.ID, + &i.OwnerID, + &i.WorkspaceID, + &i.Title, + &i.Status, + &i.WorkerID, + &i.StartedAt, + &i.HeartbeatAt, + &i.CreatedAt, + &i.UpdatedAt, + &i.ParentChatID, + &i.RootChatID, + &i.LastModelConfigID, + &i.Archived, + &i.LastError, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil } const acquireStaleChatDiffStatuses = `-- name: AcquireStaleChatDiffStatuses :many diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 49059bc2c2..dcb3b8e124 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -257,9 +257,9 @@ WHERE RETURNING *; --- name: AcquireChat :one --- Acquires a pending chat for processing. Uses SKIP LOCKED to prevent --- multiple replicas from acquiring the same chat. +-- name: AcquireChats :many +-- Acquires up to @num_chats pending chats for processing. Uses SKIP LOCKED +-- to prevent multiple replicas from acquiring the same chat. UPDATE chats SET @@ -269,7 +269,7 @@ SET updated_at = @started_at::timestamptz, worker_id = @worker_id::uuid WHERE - id = ( + id = ANY( SELECT id FROM @@ -281,7 +281,7 @@ WHERE FOR UPDATE SKIP LOCKED LIMIT - 1 + @num_chats::int ) RETURNING *;