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:
Vincenzo Palazzo
2026-08-10 13:38:25 +02:00
committed by GitHub
parent d6ee97b4ce
commit 701e93ab41
+56 -1
View File
@@ -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();