opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-21 20:51:01 +08:00
Code Issues Packages Projects Releases Wiki Activity

perf(coderd/chatd): reduce lock contention in instruction cache and persistStep (#23144)

Browse Source 
## Summary

Two targeted performance improvements to the chatd server, identified
through benchmarking.

### 1. RWMutex for instruction cache

The instruction cache is read on every chat turn to fetch the home
instruction file for a workspace agent. Writes only occur on cache
misses (once per agent per 5-minute TTL window), making the access
pattern ~90%+ reads.

Switching from `sync.Mutex` to `sync.RWMutex` and using
`RLock`/`RUnlock` on the read path allows concurrent readers instead of
serializing them.

**Benchmark (200 concurrent chats):**
| | ns/op |
|---|---|
| Mutex | 108 |
| RWMutex | 32 |
| **Speedup** | **3.4x** |

### 2. Hoist JSON marshaling out of persistStep transaction

`MarshalParts`, `PartFromContent`, `CalculateTotalCostMicros`, and the
`usageForCost` struct population are pure CPU work that ran inside the
`FOR UPDATE` transaction in `persistStep`. They have zero dependency on
the database transaction.

Moving all marshal and cost-calculation calls above `p.db.InTx()` means
the row lock is held only for `GetChatByIDForUpdate` +
`InsertChatMessage` calls.

**Benchmark (16 goroutines contending on same lock):**
| Tool calls | Inside lock | Outside lock | Speedup |
|---|---|---|---|
| 1 | 13,977 ns/op | 1,055 ns/op | 13x |
| 5 | 38,203 ns/op | 3,769 ns/op | 10x |
| 10 | 67,353 ns/op | 7,284 ns/op | 9x |
| 20 | 145,864 ns/op | 14,045 ns/op | 10x |

No behavioral changes in either commit.
This commit is contained in:
Ethan
2026-03-18 16:12:14 +11:00
committed by GitHub
parent f3bf5baba0
commit 11481d7bed
1 changed files with 54 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/chatd/chatd.go
+54 -56
View File
@@ -93,7 +93,7 @@ type Server struct {
// instructionCache caches home instruction file contents by
// workspace agent ID so we don't re-dial on every chat turn.
instructionCacheMu sync.Mutex
instructionCacheMu sync.RWMutex
instructionCache map[uuid.UUID]cachedInstruction
// Configuration
@@ -2563,13 +2563,56 @@ func (p *Server) runChat(
assistantBlocks = append(assistantBlocks, block)
}
// Pre-marshal all content outside the transaction so the
// FOR UPDATE lock is held only for the INSERT statements.
// Marshaling is pure CPU work with no database dependency.
var assistantContent pqtype.NullRawMessage
if len(assistantBlocks) > 0 {
sdkParts := make([]codersdk.ChatMessagePart, 0, len(assistantBlocks))
for _, block := range assistantBlocks {
sdkParts = append(sdkParts, chatprompt.PartFromContent(block))
}
finalAssistantText = strings.TrimSpace(contentBlocksToText(sdkParts))
var marshalErr error
assistantContent, marshalErr = chatprompt.MarshalParts(sdkParts)
if marshalErr != nil {
return xerrors.Errorf("marshal assistant content: %w", marshalErr)
}
}
toolResultContents := make([]pqtype.NullRawMessage, len(toolResults))
for i, tr := range toolResults {
trPart := chatprompt.PartFromContent(tr)
var marshalErr error
toolResultContents[i], marshalErr = chatprompt.MarshalParts([]codersdk.ChatMessagePart{trPart})
if marshalErr != nil {
return xerrors.Errorf("marshal tool result %d: %w", i, marshalErr)
}
}
hasUsage := step.Usage != (fantasy.Usage{})
var usageForCost codersdk.ChatMessageUsage
if hasUsage {
if step.Usage.InputTokens != 0 {
usageForCost.InputTokens = int64Ptr(step.Usage.InputTokens)
}
if step.Usage.OutputTokens != 0 {
usageForCost.OutputTokens = int64Ptr(step.Usage.OutputTokens)
}
if step.Usage.ReasoningTokens != 0 {
usageForCost.ReasoningTokens = int64Ptr(step.Usage.ReasoningTokens)
}
if step.Usage.CacheCreationTokens != 0 {
usageForCost.CacheCreationTokens = int64Ptr(step.Usage.CacheCreationTokens)
}
if step.Usage.CacheReadTokens != 0 {
usageForCost.CacheReadTokens = int64Ptr(step.Usage.CacheReadTokens)
}
}
totalCostMicros := chatcost.CalculateTotalCostMicros(usageForCost, callConfig.Cost)
var insertedMessages []database.ChatMessage
err := p.db.InTx(func(tx database.Store) error {
// Verify this worker still owns the chat before
// inserting messages. This closes the race where
// EditMessage truncates history and clears worker_id
// while persistInterruptedStep (which uses an
// uncancelable context) is still running.
lockedChat, lockErr := tx.GetChatByIDForUpdate(persistCtx, chat.ID)
if lockErr != nil {
return xerrors.Errorf("lock chat for persist: %w", lockErr)
@@ -2578,42 +2621,7 @@ func (p *Server) runChat(
return chatloop.ErrInterrupted
}
if len(assistantBlocks) > 0 {
sdkParts := make([]codersdk.ChatMessagePart, 0, len(assistantBlocks))
for _, block := range assistantBlocks {
sdkParts = append(sdkParts, chatprompt.PartFromContent(block))
}
finalAssistantText = strings.TrimSpace(contentBlocksToText(sdkParts))
assistantContent, marshalErr := chatprompt.MarshalParts(sdkParts)
if marshalErr != nil {
return marshalErr
}
hasUsage := step.Usage != (fantasy.Usage{})
var usageForCost codersdk.ChatMessageUsage
if hasUsage {
// Only populate fields that the provider explicitly
// reported. Nil fields tell the calculator "no data"
// vs zero meaning "reported as zero tokens."
if step.Usage.InputTokens != 0 {
usageForCost.InputTokens = int64Ptr(step.Usage.InputTokens)
}
if step.Usage.OutputTokens != 0 {
usageForCost.OutputTokens = int64Ptr(step.Usage.OutputTokens)
}
if step.Usage.ReasoningTokens != 0 {
usageForCost.ReasoningTokens = int64Ptr(step.Usage.ReasoningTokens)
}
if step.Usage.CacheCreationTokens != 0 {
usageForCost.CacheCreationTokens = int64Ptr(step.Usage.CacheCreationTokens)
}
if step.Usage.CacheReadTokens != 0 {
usageForCost.CacheReadTokens = int64Ptr(step.Usage.CacheReadTokens)
}
}
totalCostMicros := chatcost.CalculateTotalCostMicros(usageForCost, callConfig.Cost)
if assistantContent.Valid {
assistantMessage, insertErr := tx.InsertChatMessage(persistCtx, database.InsertChatMessageParams{
ChatID: chat.ID,
CreatedBy: uuid.NullUUID{},
@@ -2636,10 +2644,6 @@ func (p *Server) runChat(
CacheReadTokens: usageNullInt64(step.Usage.CacheReadTokens, hasUsage),
ContextLimit: step.ContextLimit,
Compressed: sql.NullBool{},
// TotalCostMicros is nullable: NULL means "unpriced"
// (pricing config was missing or no priced token
// breakdown available), while 0 means "priced at
// zero cost" (e.g., a free model).
TotalCostMicros: usageNullInt64Ptr(totalCostMicros),
})
if insertErr != nil {
@@ -2648,13 +2652,7 @@ func (p *Server) runChat(
insertedMessages = append(insertedMessages, assistantMessage)
}
for _, tr := range toolResults {
trPart := chatprompt.PartFromContent(tr)
resultContent, marshalErr := chatprompt.MarshalParts([]codersdk.ChatMessagePart{trPart})
if marshalErr != nil {
return marshalErr
}
for i, resultContent := range toolResultContents {
toolMessage, insertErr := tx.InsertChatMessage(persistCtx, database.InsertChatMessageParams{
ChatID: chat.ID,
CreatedBy: uuid.NullUUID{},
@@ -2674,7 +2672,7 @@ func (p *Server) runChat(
Compressed: sql.NullBool{},
})
if insertErr != nil {
return xerrors.Errorf("insert tool result: %w", insertErr)
return xerrors.Errorf("insert tool result %d: %w", i, insertErr)
}
insertedMessages = append(insertedMessages, toolMessage)
}
@@ -3277,9 +3275,9 @@ func (p *Server) resolveInstructions(
}
agentID := agent.ID
p.instructionCacheMu.Lock()
p.instructionCacheMu.RLock()
cached, ok := p.instructionCache[agentID]
p.instructionCacheMu.Unlock()
p.instructionCacheMu.RUnlock()
if ok && time.Since(cached.fetchedAt) < instructionCacheTTL {
return cached.instruction