Index messages by (session_id, created_timestamp, id) to stop on-disk sort storms (#10874)
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Lifei Zhou <lifei@squareup.com>
This commit is contained in:
@@ -24,7 +24,7 @@ use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, LazyLock};
|
||||
use tracing::{info, warn};
|
||||
|
||||
pub const CURRENT_SCHEMA_VERSION: i32 = 15;
|
||||
pub const CURRENT_SCHEMA_VERSION: i32 = 16;
|
||||
pub const SESSIONS_FOLDER: &str = "sessions";
|
||||
pub const DB_NAME: &str = "sessions.db";
|
||||
const MILLISECOND_TIMESTAMP_THRESHOLD: i64 = 10_000_000_000;
|
||||
@@ -1079,6 +1079,11 @@ impl SessionStorage {
|
||||
sqlx::query("CREATE INDEX IF NOT EXISTS idx_messages_message_id ON messages(message_id)")
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS idx_messages_session_created ON messages(session_id, created_timestamp, id)",
|
||||
)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
sqlx::query("CREATE INDEX IF NOT EXISTS idx_sessions_updated ON sessions(updated_at DESC)")
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
@@ -1558,6 +1563,13 @@ impl SessionStorage {
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
}
|
||||
16 => {
|
||||
sqlx::query(
|
||||
"CREATE INDEX IF NOT EXISTS idx_messages_session_created ON messages(session_id, created_timestamp, id)",
|
||||
)
|
||||
.execute(&mut **tx)
|
||||
.await?;
|
||||
}
|
||||
_ => {
|
||||
anyhow::bail!("Unknown migration version: {}", version);
|
||||
}
|
||||
@@ -2963,6 +2975,49 @@ mod tests {
|
||||
assert_eq!(texts, ["appended first", "built first"]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_messages_session_created_index_avoids_disk_sort() {
|
||||
use sqlx::Row;
|
||||
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
let sm = SessionManager::new(temp_dir.path().to_path_buf());
|
||||
let pool = sm.storage.pool().await.unwrap();
|
||||
|
||||
let index_exists: bool = sqlx::query_scalar(
|
||||
"SELECT EXISTS (SELECT 1 FROM sqlite_master WHERE type='index' AND name='idx_messages_session_created')",
|
||||
)
|
||||
.fetch_one(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
index_exists,
|
||||
"idx_messages_session_created should exist after schema init"
|
||||
);
|
||||
|
||||
let plan_rows = sqlx::query(
|
||||
"EXPLAIN QUERY PLAN \
|
||||
SELECT content_json FROM messages WHERE session_id = ? ORDER BY created_timestamp, id",
|
||||
)
|
||||
.bind("nonexistent_session")
|
||||
.fetch_all(pool)
|
||||
.await
|
||||
.unwrap();
|
||||
let plan_text: String = plan_rows
|
||||
.iter()
|
||||
.map(|r| r.try_get::<String, _>("detail").unwrap_or_default())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
|
||||
assert!(
|
||||
plan_text.contains("idx_messages_session_created"),
|
||||
"loading a session's messages should use idx_messages_session_created, got: {plan_text}"
|
||||
);
|
||||
assert!(
|
||||
!plan_text.contains("TEMP B-TREE"),
|
||||
"loading a session's messages must not require an on-disk sort, got: {plan_text}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_last_message_at_is_derived_from_messages() {
|
||||
let temp_dir = TempDir::new().unwrap();
|
||||
|
||||
Reference in New Issue
Block a user