Thanks to visit codestin.com
Credit goes to github.com

Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions coderd/x/chatd/ARCHITECTURE.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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:

Expand Down
58 changes: 58 additions & 0 deletions coderd/x/chatd/capacity.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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,
Expand Down
80 changes: 80 additions & 0 deletions coderd/x/chatd/capacity_internal_test.go
Original file line number Diff line number Diff line change
@@ -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)
})
}
23 changes: 21 additions & 2 deletions coderd/x/chatd/worker.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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 {
Expand All @@ -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
Expand All @@ -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 {
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -304,18 +320,21 @@ 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
})
if errors.Is(err, errCapacityRefused) {
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)
Expand Down
Loading