diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index b86285bc49a..a83a6895983 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -655,6 +655,14 @@ For every matching chat, it locks it, checks if the chat still meets the aforeme When a chat is successfully acquired, the acquisition loop requests the [Runner manager](#runner-manager) to spawn a chat runner for it. +### Capacity wait + +The [Concurrent agent limiter](#concurrent-agent-limiter) can refuse an otherwise acquirable chat when its pool is full. The acquisition loop remembers, per chat, when this worker first saw the chat refused for capacity. When the same worker later acquires the chat, it records a `capacity_wait` lifecycle stage from that first refusal to the acquisition (see [Lifecycle tracing](#lifecycle-tracing)). Chats admitted on their first attempt record nothing. + +The bookkeeping is local to the worker. The wait start is dropped when the worker skips the chat for a reason other than capacity (it is owned by a live runner, archived, or no longer runnable), and entries for chats that have left the candidate set are pruned only when the candidate batch is shorter than its limit, since a chat missing from a truncated batch may still be waiting. The map is touched only by the acquisition goroutine. + +Because the capacity limit is deployment-wide but the refusal history is per worker, the recorded wait is a lower bound. A chat refused on one replica and acquired by another is measured from the acquiring replica's first refusal, or not at all if that replica admitted it on its first attempt. A replica restart also discards its history. + ### Load balancing The design doesn't attempt to distribute load between workers fairly. Whenever a chat needs an owner, all replicas race to acquire it. If there's a coder replica that has a lower latency to the database, it'll tend to acquire chats more frequently than other replicas. @@ -802,9 +810,9 @@ Work detached from the turn, such as title, summary, and status label generation #### Stages -Every stage carries `scope` (`turn` or `background`) and `chat_kind` (`root` or `subagent`), and once the model is resolved, `model`. Spans additionally carry `reasoning_effort`; it is not a metric label because it multiplies series per model. Stages that are not tied to a model call (`acquisition`, `queue_wait`, `mcp_connect`, `retry_backoff`, `commit`) carry an empty `model` label. `prepare` is stamped with the model once preparation resolves it. +Every stage carries `scope` (`turn` or `background`) and `chat_kind` (`root` or `subagent`), and once the model is resolved, `model`. Spans additionally carry `reasoning_effort`; it is not a metric label because it multiplies series per model. Stages that are not tied to a model call (`acquisition`, `queue_wait`, `capacity_wait`, `mcp_connect`, `retry_backoff`, `commit`) carry an empty `model` label. `prepare` is stamped with the model once preparation resolves it. -At `--chat-stage-metrics=basic` the histogram observes `chat_turn`, `queue_wait`, `acquisition`, `mcp_connect`, `stream`, `time_to_first_token`, `provider_attempt`, `tool_call`, `commit`, and `retry_backoff`. `generation_step`, `prepare`, `thinking`, and `compaction` are span-only at that level. +At `--chat-stage-metrics=basic` the histogram observes `chat_turn`, `queue_wait`, `capacity_wait`, `acquisition`, `mcp_connect`, `stream`, `time_to_first_token`, `provider_attempt`, `tool_call`, `commit`, and `retry_backoff`. `generation_step`, `prepare`, `thinking`, and `compaction` are span-only at that level. Live stages wrap a section of code and end when it returns: diff --git a/coderd/x/chatd/capacity.go b/coderd/x/chatd/capacity.go index 1ebcf8669a7..b619178d8cd 100644 --- a/coderd/x/chatd/capacity.go +++ b/coderd/x/chatd/capacity.go @@ -5,8 +5,10 @@ import ( "github.com/google/uuid" "github.com/prometheus/client_golang/prometheus" + "go.opentelemetry.io/otel/attribute" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatloop" ) type capacityMetrics struct { @@ -46,6 +48,62 @@ func (w *chatWorker) capacityMetricsLoop(ctx context.Context) { } } +// noteCapacityRefused remembers when a chat was first refused a +// capacity slot. Only the acquisition loop touches the map, so it +// needs no lock. +func (w *chatWorker) noteCapacityRefused(chatID uuid.UUID) { + if _, ok := w.capacityWaitSince[chatID]; ok { + return + } + w.capacityWaitSince[chatID] = w.opts.Clock.Now() +} + +// recordCapacityWait emits the capacity_wait stage for a chat that is +// being acquired after at least one capacity refusal, measured from +// the first refusal this worker saw. Chats admitted on their first +// attempt record nothing. No turn span exists at this point, so the +// turn scope and the chat kind are stated explicitly. +func (w *chatWorker) recordCapacityWait(ctx context.Context, chat database.Chat) { + since, waited := w.capacityWaitSince[chat.ID] + if !waited { + return + } + delete(w.capacityWaitSince, chat.ID) + ctx = chatloop.ContextWithChatKind(ctx, chatKindAttr(chat)) + w.server.stages.RecordAs(ctx, chatloop.StageCapacityWait, chatloop.ScopeTurn, chatloop.StageModel{}, + since, w.opts.Clock.Now(), nil, + attribute.String(chatloop.AttrChatID, chat.ID.String()), + ) +} + +// forgetCapacityWait drops the wait start of a chat this worker will +// not acquire on the current pass for a reason other than capacity: it +// is owned by a live runner, archived, or no longer runnable. A wait +// that resumes later starts from the next refusal. +func (w *chatWorker) forgetCapacityWait(chatID uuid.UUID) { + delete(w.capacityWaitSince, chatID) +} + +// pruneCapacityWaits drops wait starts for chats that are no longer +// acquisition candidates, which happens when they are archived, +// deleted, or picked up by another worker. candidates must be the +// complete candidate set: a chat missing from a truncated batch is +// still waiting, and dropping it would restart its clock. +func (w *chatWorker) pruneCapacityWaits(candidates []database.GetChatWorkerAcquisitionCandidatesRow) { + if len(w.capacityWaitSince) == 0 { + return + } + stillCandidate := make(map[uuid.UUID]struct{}, len(candidates)) + for _, row := range candidates { + stillCandidate[row.ID] = struct{}{} + } + for chatID := range w.capacityWaitSince { + if _, ok := stillCandidate[chatID]; !ok { + delete(w.capacityWaitSince, chatID) + } + } +} + func (w *chatWorker) refreshCapacityMetrics(ctx context.Context) { active, err := w.opts.Store.CountChatCapacityActiveByPool(ctx, database.CountChatCapacityActiveByPoolParams{ ExcludeChatID: uuid.Nil, diff --git a/coderd/x/chatd/capacity_internal_test.go b/coderd/x/chatd/capacity_internal_test.go new file mode 100644 index 00000000000..7c6e072ae35 --- /dev/null +++ b/coderd/x/chatd/capacity_internal_test.go @@ -0,0 +1,80 @@ +package chatd //nolint:testpackage // Tests the acquisition loop's capacity wait bookkeeping. + +import ( + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/quartz" +) + +func newCapacityWaitWorker(t *testing.T) (*chatWorker, *quartz.Mock) { + t.Helper() + clock := quartz.NewMock(t) + tracer, _ := newStageTestTracer(t) + return &chatWorker{ + server: &Server{stages: tracer}, + opts: chatWorkerOptions{Clock: clock}, + capacityWaitSince: make(map[uuid.UUID]time.Time), + }, clock +} + +func candidateRows(ids ...uuid.UUID) []database.GetChatWorkerAcquisitionCandidatesRow { + rows := make([]database.GetChatWorkerAcquisitionCandidatesRow, 0, len(ids)) + for _, id := range ids { + rows = append(rows, database.GetChatWorkerAcquisitionCandidatesRow{ID: id}) + } + return rows +} + +func TestCapacityWaitBookkeeping(t *testing.T) { + t.Parallel() + + t.Run("FirstRefusalStartsTheClock", func(t *testing.T) { + t.Parallel() + worker, clock := newCapacityWaitWorker(t) + chatID := uuid.New() + + worker.noteCapacityRefused(chatID) + first := worker.capacityWaitSince[chatID] + clock.Advance(time.Second) + worker.noteCapacityRefused(chatID) + require.Equal(t, first, worker.capacityWaitSince[chatID], "a later refusal keeps the first start") + }) + + t.Run("SkippedChatForgetsItsWait", func(t *testing.T) { + t.Parallel() + worker, _ := newCapacityWaitWorker(t) + chatID := uuid.New() + + worker.noteCapacityRefused(chatID) + worker.forgetCapacityWait(chatID) + require.NotContains(t, worker.capacityWaitSince, chatID) + }) + + t.Run("PruneKeepsCandidates", func(t *testing.T) { + t.Parallel() + worker, _ := newCapacityWaitWorker(t) + waiting, gone := uuid.New(), uuid.New() + + worker.noteCapacityRefused(waiting) + worker.noteCapacityRefused(gone) + worker.pruneCapacityWaits(candidateRows(waiting, uuid.New())) + require.Contains(t, worker.capacityWaitSince, waiting) + require.NotContains(t, worker.capacityWaitSince, gone) + }) + + t.Run("RecordClearsTheWait", func(t *testing.T) { + t.Parallel() + worker, clock := newCapacityWaitWorker(t) + chat := database.Chat{ID: uuid.New()} + + worker.noteCapacityRefused(chat.ID) + clock.Advance(time.Second) + worker.recordCapacityWait(t.Context(), chat) + require.NotContains(t, worker.capacityWaitSince, chat.ID) + }) +} diff --git a/coderd/x/chatd/worker.go b/coderd/x/chatd/worker.go index 83123f48ccb..980fc4243a1 100644 --- a/coderd/x/chatd/worker.go +++ b/coderd/x/chatd/worker.go @@ -28,6 +28,10 @@ type chatWorker struct { unsubscribe func() wakeCh chan struct{} wg sync.WaitGroup + + // capacityWaitSince tracks when each chat was first refused a + // capacity slot. Only the acquisition loop reads or writes it. + capacityWaitSince map[uuid.UUID]time.Time } // newChatWorker constructs a chat worker. The worker is idle until Start is @@ -40,7 +44,11 @@ func newChatWorker(server *Server, opts chatWorkerOptions) (*chatWorker, error) if err != nil { return nil, err } - return &chatWorker{server: server, opts: withDefaults}, nil + return &chatWorker{ + server: server, + opts: withDefaults, + capacityWaitSince: make(map[uuid.UUID]time.Time), + }, nil } // chatWorkerID returns this worker's configured worker ID. @@ -193,9 +201,10 @@ func (w *chatWorker) acquisitionLoop( func (w *chatWorker) acquireOnce(ctx context.Context, workerID uuid.UUID, manager *runnerManager) { // Fetch twice the budget so one full pool cannot hide candidates in the other. + limit := w.opts.AcquisitionBatchSize * 2 rows, err := w.opts.Store.GetChatWorkerAcquisitionCandidates(ctx, database.GetChatWorkerAcquisitionCandidatesParams{ StaleSeconds: w.opts.HeartbeatStaleSeconds, - LimitCount: w.opts.AcquisitionBatchSize * 2, + LimitCount: limit, }) if err != nil { if ctx.Err() == nil { @@ -207,6 +216,11 @@ func (w *chatWorker) acquireOnce(ctx context.Context, workerID uuid.UUID, manage acquired := int32(0) rootPoolRefused := false subagentPoolRefused := false + // A batch shorter than the limit holds every candidate, so a chat + // absent from it has left the candidate set. + if len(rows) < int(limit) { + w.pruneCapacityWaits(rows) + } for _, row := range rows { if acquired >= w.opts.AcquisitionBatchSize { return @@ -220,6 +234,7 @@ func (w *chatWorker) acquireOnce(ctx context.Context, workerID uuid.UUID, manage } candidateAcquired, err := w.acquireCandidateSafely(ctx, workerID, manager, row.ID) if errors.Is(err, errCapacityRefused) { + w.noteCapacityRefused(row.ID) if isSubagent { subagentPoolRefused = true } else { @@ -266,6 +281,7 @@ func (w *chatWorker) acquireCandidate( chatID uuid.UUID, ) (bool, error) { runnerID := uuid.New() + var acquiredChat database.Chat machine := chatstate.NewChatMachine(w.opts.Store, w.opts.Pubsub, chatID) err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { chat, err := store.GetChatByID(ctx, chatID) @@ -304,6 +320,7 @@ func (w *chatWorker) acquireCandidate( // worker into an immediate retry of this unowned chat. return errCapacityRefused } + acquiredChat = chat _, err = tx.Acquire(chatstate.AcquireInput{WorkerID: workerID, RunnerID: runnerID}) return err }) @@ -311,11 +328,13 @@ func (w *chatWorker) acquireCandidate( return false, errCapacityRefused } if errors.Is(err, errSkipAcquire) || errors.Is(err, chatstate.ErrChatNotFound) { + w.forgetCapacityWait(chatID) return false, nil } if err != nil { return false, err } + w.recordCapacityWait(ctx, acquiredChat) if err := manager.Spawn(ctx, spawnRunnerRequest{ChatID: chatID, WorkerID: workerID, RunnerID: runnerID}); err != nil { if errAbandon := w.abandonAcquiredChat(ctx, workerID, runnerID, chatID); errAbandon != nil { return false, errors.Join(err, errAbandon)