mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
perf(coderd/database): skip redundant chat row update in InsertChatMessage (#23111)
## Summary
- add an `IS DISTINCT FROM` guard to `InsertChatMessage`'s
`updated_chat` CTE so `chats.last_model_config_id` is only rewritten
when the incoming `model_config_id` actually changes
- regenerate the query layer
- add focused regression coverage for the two meaningful behaviors:
same-model inserts and real model switches
- trim redundant message-field assertions so the new test stays focused
on the guard behavior
## Proof this is an improvement
This PR reduces work in the hottest chat write query without changing
the insert behavior.
### Why the old query did unnecessary work
Before this change, `InsertChatMessage` always ran this update whenever
`model_config_id` was non-null:
```sql
UPDATE chats
SET last_model_config_id = sqlc.narg('model_config_id')::uuid
WHERE id = @chat_id::uuid
AND sqlc.narg('model_config_id')::uuid IS NOT NULL
```
That means the query rewrote the `chats` row even when
`chats.last_model_config_id` was already equal to the incoming value.
### What changes in this PR
This PR adds:
```sql
AND chats.last_model_config_id IS DISTINCT FROM sqlc.narg('model_config_id')::uuid
```
So same-model inserts still insert the message, but they no longer
perform a redundant `UPDATE chats`.
### Why this matters on the hot path
From the chat scaletest investigation that motivated this change:
- `InsertChatMessage` (+ `updated_chat` CTE) was the hottest write query
- about **104k calls**
- about **0.69 ms average latency**
- about **71.8 s total DB execution time**
We also verified common callsites where the update is provably
redundant:
- `CreateChat` inserts the chat with `LastModelConfigID =
opts.ModelConfigID`, then immediately inserts initial system/user
messages with that same model config
- follow-up user messages commonly pass `lockedChat.LastModelConfigID`
straight into `InsertChatMessage`
- assistant/tool/summary persistence keeps the current model in the
common case; only real switches or fallback cases need the chat row
update
That means a meaningful fraction of executions of the hottest DB write
query move from:
- **before:** insert message **+** rewrite chat row
- **after:** insert message only
This should reduce row churn and write contention on `chats`, especially
against other chat-row writers like `UpdateChatStatus` and
`GetChatByIDForUpdate`.
This commit is contained in:
@@ -8997,6 +8997,131 @@ func TestInsertWorkspaceAgentDevcontainers(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestInsertChatMessage(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
insertModelConfig := func(
|
||||
t *testing.T,
|
||||
store database.Store,
|
||||
ctx context.Context,
|
||||
userID uuid.UUID,
|
||||
provider string,
|
||||
model string,
|
||||
displayName string,
|
||||
isDefault bool,
|
||||
) database.ChatModelConfig {
|
||||
t.Helper()
|
||||
|
||||
modelConfig, err := store.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
|
||||
Provider: provider,
|
||||
Model: model,
|
||||
DisplayName: displayName,
|
||||
CreatedBy: uuid.NullUUID{UUID: userID, Valid: true},
|
||||
UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true},
|
||||
Enabled: true,
|
||||
IsDefault: isDefault,
|
||||
ContextLimit: 128000,
|
||||
CompressionThreshold: 80,
|
||||
Options: json.RawMessage(`{}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
return modelConfig
|
||||
}
|
||||
|
||||
setupChat := func(t *testing.T) (database.Store, context.Context, database.User, database.Chat, string, database.ChatModelConfig) {
|
||||
t.Helper()
|
||||
|
||||
store, _ := dbtestutil.NewDB(t)
|
||||
ctx := context.Background()
|
||||
|
||||
dbgen.Organization(t, store, database.Organization{})
|
||||
user := dbgen.User(t, store, database.User{})
|
||||
provider := "openai"
|
||||
|
||||
_, err := store.InsertChatProvider(ctx, database.InsertChatProviderParams{
|
||||
Provider: provider,
|
||||
DisplayName: "OpenAI",
|
||||
APIKey: "test-key",
|
||||
Enabled: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
modelConfigA := insertModelConfig(
|
||||
t,
|
||||
store,
|
||||
ctx,
|
||||
user.ID,
|
||||
provider,
|
||||
"test-model-a-"+uuid.NewString(),
|
||||
"Test Model A",
|
||||
true,
|
||||
)
|
||||
|
||||
chat, err := store.InsertChat(ctx, database.InsertChatParams{
|
||||
OwnerID: user.ID,
|
||||
LastModelConfigID: modelConfigA.ID,
|
||||
Title: "test-chat-" + uuid.NewString(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
return store, ctx, user, chat, provider, modelConfigA
|
||||
}
|
||||
|
||||
insertMessage := func(t *testing.T, store database.Store, ctx context.Context, chatID, userID, modelConfigID uuid.UUID, content string) {
|
||||
t.Helper()
|
||||
|
||||
_, err := store.InsertChatMessage(ctx, database.InsertChatMessageParams{
|
||||
ChatID: chatID,
|
||||
CreatedBy: uuid.NullUUID{UUID: userID, Valid: true},
|
||||
ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true},
|
||||
Role: database.ChatMessageRoleUser,
|
||||
ContentVersion: chatprompt.CurrentContentVersion,
|
||||
Visibility: database.ChatMessageVisibilityBoth,
|
||||
Content: pqtype.NullRawMessage{
|
||||
RawMessage: json.RawMessage(fmt.Sprintf("%q", content)),
|
||||
Valid: true,
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
t.Run("ModelSwitchUpdatesLastModelConfigID", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
store, ctx, user, chat, provider, modelConfigA := setupChat(t)
|
||||
modelConfigB := insertModelConfig(
|
||||
t,
|
||||
store,
|
||||
ctx,
|
||||
user.ID,
|
||||
provider,
|
||||
"test-model-b-"+uuid.NewString(),
|
||||
"Test Model B",
|
||||
false,
|
||||
)
|
||||
|
||||
insertMessage(t, store, ctx, chat.ID, user.ID, modelConfigB.ID, "switch models")
|
||||
|
||||
gotChat, err := store.GetChatByID(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, modelConfigA.ID, chat.LastModelConfigID)
|
||||
require.Equal(t, modelConfigB.ID, gotChat.LastModelConfigID)
|
||||
})
|
||||
|
||||
t.Run("SameModelDoesNotBreakAnything", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
store, ctx, user, chat, _, modelConfigA := setupChat(t)
|
||||
|
||||
insertMessage(t, store, ctx, chat.ID, user.ID, modelConfigA.ID, "same model")
|
||||
|
||||
gotChat, err := store.GetChatByID(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, modelConfigA.ID, gotChat.LastModelConfigID)
|
||||
})
|
||||
}
|
||||
|
||||
func TestGetChatMessagesForPromptByChatID(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
@@ -4211,6 +4211,7 @@ WITH updated_chat AS (
|
||||
WHERE
|
||||
id = $1::uuid
|
||||
AND $3::uuid IS NOT NULL
|
||||
AND chats.last_model_config_id IS DISTINCT FROM $3::uuid
|
||||
)
|
||||
INSERT INTO chat_messages (
|
||||
chat_id,
|
||||
|
||||
@@ -165,6 +165,7 @@ WITH updated_chat AS (
|
||||
WHERE
|
||||
id = @chat_id::uuid
|
||||
AND sqlc.narg('model_config_id')::uuid IS NOT NULL
|
||||
AND chats.last_model_config_id IS DISTINCT FROM sqlc.narg('model_config_id')::uuid
|
||||
)
|
||||
INSERT INTO chat_messages (
|
||||
chat_id,
|
||||
|
||||
Reference in New Issue
Block a user