From 7ae831991792744434027c431dfbc7e81dd5c12c Mon Sep 17 00:00:00 2001 From: willwang-openai Date: Thu, 15 Oct 2026 17:55:14 +0000 Subject: [PATCH] Retry SQLite operations on transient writer contention (#54369) ## Why Competing SQLite writers can cause database opening and writes to fail with `SQLITE_BUSY`. Deferred write transactions can also encounter contention after taking a read snapshot. ## What changed - Add up to two retries with 100 ms and 200 ms delays for SQLite busy errors during pool opening, writer admission, and standalone writes. - Use a shared `BEGIN IMMEDIATE` helper across state, thread history, and agent message board writes to acquire the writer slot before reading. - Preserve bound arguments when retrying standalone writes, and commit writes using `RETURNING` before exposing their results. Retry transaction admission without replaying an already-started transaction. ## Testing Add tests using real SQLite writer contention to verify recovery after a competing writer releases its lock. Queue tests cover enqueue, update, and delete, checking preserved payloads, single revision increments, and no duplicate messages or unintended deletions. GitOrigin-RevId: ac5bd31a87c67cc47fb385bf3fa5435a54ff70f0 --- .../src/local/lifecycle.rs | 12 +- .../ext/agent-message-board/src/local/pool.rs | 6 +- codex-rs/state/src/lib.rs | 1 + codex-rs/state/src/migrations.rs | 3 +- codex-rs/state/src/runtime.rs | 3 +- codex-rs/state/src/runtime/backfill.rs | 9 +- .../runtime/external_agent_config_imports.rs | 3 +- codex-rs/state/src/runtime/goals.rs | 27 +- codex-rs/state/src/runtime/memories.rs | 35 ++- codex-rs/state/src/runtime/projects.rs | 11 +- codex-rs/state/src/runtime/queued_items.rs | 12 +- codex-rs/state/src/runtime/remote_control.rs | 7 +- .../state/src/runtime/rollout_migration.rs | 7 +- .../state/src/runtime/thread_attachments.rs | 7 +- .../state/src/runtime/thread_section_order.rs | 3 +- codex-rs/state/src/runtime/thread_sections.rs | 8 +- codex-rs/state/src/runtime/threads.rs | 36 +-- codex-rs/state/src/sqlite.rs | 9 +- codex-rs/state/src/sqlite/busy_retry.rs | 191 +++++++++++++ codex-rs/state/src/sqlite/retry_tests.rs | 267 ++++++++++++++++++ .../thread-store/src/local/thread_history.rs | 8 +- 21 files changed, 572 insertions(+), 93 deletions(-) create mode 100644 codex-rs/state/src/sqlite/busy_retry.rs create mode 100644 codex-rs/state/src/sqlite/retry_tests.rs diff --git a/codex-rs/ext/agent-message-board/src/local/lifecycle.rs b/codex-rs/ext/agent-message-board/src/local/lifecycle.rs index 8b3daa3196..12e30dda61 100644 --- a/codex-rs/ext/agent-message-board/src/local/lifecycle.rs +++ b/codex-rs/ext/agent-message-board/src/local/lifecycle.rs @@ -13,6 +13,7 @@ use super::storage_error; use codex_protocol::SessionId; use codex_protocol::error::Result; use codex_state::SqliteConfig; +use codex_state::SqlitePoolWriteExt; use sqlx::Sqlite; use sqlx::Transaction; use std::sync::Weak; @@ -62,10 +63,7 @@ impl LocalAgentMessageBoard { .execute(&pool) .await .map_err(storage_error)?; - let mut tx = pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_err(storage_error)?; + let mut tx = pool.begin_immediate().await.map_err(storage_error)?; for root in roots { sqlx::query("INSERT OR IGNORE INTO deleted_boards (board) VALUES (?)") .bind(root.to_string()) @@ -92,11 +90,7 @@ impl LocalAgentMessageBoard { } pub(super) async fn begin_write(&self) -> Result> { - let mut tx = self - .pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_err(storage_error)?; + let mut tx = self.pool.begin_immediate().await.map_err(storage_error)?; let deleted: bool = sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM deleted_boards WHERE board=?)") .bind(self.identity.to_string()) diff --git a/codex-rs/ext/agent-message-board/src/local/pool.rs b/codex-rs/ext/agent-message-board/src/local/pool.rs index b27417ca74..c39b48d5ea 100644 --- a/codex-rs/ext/agent-message-board/src/local/pool.rs +++ b/codex-rs/ext/agent-message-board/src/local/pool.rs @@ -4,6 +4,7 @@ use super::storage_error; use codex_protocol::error::Result; use codex_state::SqliteConfig; +use codex_state::SqlitePoolWriteExt; use sqlx::SqlitePool; use std::collections::HashMap; use std::path::PathBuf; @@ -70,10 +71,7 @@ pub(super) async fn open(sqlite: &SqliteConfig) -> Result> { .open_read_write_pool(&path) .await .map_err(storage_error)?; - let mut tx = pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_err(storage_error)?; + let mut tx = pool.begin_immediate().await.map_err(storage_error)?; sqlx::raw_sql(SCHEMA) .execute(&mut *tx) .await diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index d1de678a86..78aea77025 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -41,6 +41,7 @@ pub use runtime::MAX_GUARDIAN_REVIEW_RECORDS_PER_THREAD; /// Preferred entrypoint: owns configuration and metrics. pub use runtime::StateRuntime; pub use sqlite::SqliteConfig; +pub use sqlite::SqlitePoolWriteExt; pub use audit::ThreadStateAuditRow; pub use audit::read_thread_state_audit_rows; diff --git a/codex-rs/state/src/migrations.rs b/codex-rs/state/src/migrations.rs index efbb76cc09..33eed1bb96 100644 --- a/codex-rs/state/src/migrations.rs +++ b/codex-rs/state/src/migrations.rs @@ -1,5 +1,6 @@ use std::borrow::Cow; +use crate::sqlite::SqliteWriteQueryExt; use sqlx::migrate::Migrator; use sqlx_sqlite::SqlitePool; @@ -112,7 +113,7 @@ WHERE version = ? .bind(38_i64) .bind(recency_migration.checksum.as_ref()) .bind(recency_migration.version) - .execute(pool) + .execute_with_busy_retry(pool) .await?; Ok(()) } diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index 244f219cd5..58dd668469 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -19,6 +19,7 @@ use crate::model::datetime_to_epoch_millis; use crate::model::datetime_to_epoch_seconds; use crate::model::epoch_millis_to_datetime; use crate::paths::file_modified_time_utc; +use crate::sqlite::SqliteWriteQueryExt; use crate::telemetry::DbKind; use crate::telemetry::DbTelemetry; use chrono::DateTime; @@ -391,7 +392,7 @@ ON CONFLICT(id) DO NOTHING .bind(1_i64) .bind(crate::BackfillStatus::Pending.as_str()) .bind(Utc::now().timestamp()) - .execute(pool) + .execute_with_busy_retry(pool) .await?; Ok(()) } diff --git a/codex-rs/state/src/runtime/backfill.rs b/codex-rs/state/src/runtime/backfill.rs index 1acbf01274..1ef89e54ee 100644 --- a/codex-rs/state/src/runtime/backfill.rs +++ b/codex-rs/state/src/runtime/backfill.rs @@ -1,4 +1,5 @@ use super::*; +use crate::sqlite::SqliteWriteQueryExt; impl StateRuntime { pub async fn get_backfill_state(&self) -> anyhow::Result { @@ -38,7 +39,7 @@ WHERE id = 1 .bind(crate::BackfillStatus::Complete.as_str()) .bind(crate::BackfillStatus::Running.as_str()) .bind(lease_cutoff) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected() == 1) } @@ -55,7 +56,7 @@ WHERE id = 1 ) .bind(crate::BackfillStatus::Running.as_str()) .bind(Utc::now().timestamp()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(()) } @@ -73,7 +74,7 @@ WHERE id = 1 .bind(crate::BackfillStatus::Running.as_str()) .bind(watermark) .bind(Utc::now().timestamp()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(()) } @@ -97,7 +98,7 @@ WHERE id = 1 .bind(last_watermark) .bind(now) .bind(now) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(()) } diff --git a/codex-rs/state/src/runtime/external_agent_config_imports.rs b/codex-rs/state/src/runtime/external_agent_config_imports.rs index 6a6c2edf6d..7b2ea28f39 100644 --- a/codex-rs/state/src/runtime/external_agent_config_imports.rs +++ b/codex-rs/state/src/runtime/external_agent_config_imports.rs @@ -1,5 +1,6 @@ use super::StateRuntime; use crate::model::datetime_to_epoch_millis; +use crate::sqlite::SqliteWriteQueryExt; use chrono::Utc; use serde::Deserialize; use serde::Serialize; @@ -72,7 +73,7 @@ ON CONFLICT(import_id) DO UPDATE SET .bind(datetime_to_epoch_millis(Utc::now())) .bind(serde_json::to_string(successes)?) .bind(serde_json::to_string(failures)?) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(()) diff --git a/codex-rs/state/src/runtime/goals.rs b/codex-rs/state/src/runtime/goals.rs index ed616e16b7..9f58554314 100644 --- a/codex-rs/state/src/runtime/goals.rs +++ b/codex-rs/state/src/runtime/goals.rs @@ -1,5 +1,7 @@ use super::*; use crate::model::ThreadGoalRow; +use crate::sqlite::SqlitePoolWriteExt; +use crate::sqlite::SqliteWriteQueryExt; use uuid::Uuid; #[derive(Clone)] @@ -69,7 +71,7 @@ WHERE thread_id = ? &self, goal: &crate::ThreadGoal, ) -> anyhow::Result<()> { - let mut transaction = self.pool.begin().await?; + let mut transaction = self.pool.begin_immediate().await?; sqlx::query( r#" INSERT INTO thread_goals ( @@ -147,7 +149,7 @@ SELECT EXISTS( ) -> anyhow::Result<()> { sqlx::query("DELETE FROM thread_goal_continuation_deferrals WHERE thread_id = ?") .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(()) @@ -204,7 +206,7 @@ RETURNING .bind(token_budget) .bind(now_ms) .bind(now_ms) - .fetch_one(self.pool.as_ref()) + .fetch_one_write(self.pool.as_ref()) .await?; thread_goal_from_row(&row) @@ -262,7 +264,7 @@ RETURNING .bind(token_budget) .bind(now_ms) .bind(now_ms) - .fetch_optional(self.pool.as_ref()) + .fetch_optional_write(self.pool.as_ref()) .await?; row.map(|row| thread_goal_from_row(&row)).transpose() @@ -315,7 +317,7 @@ WHERE thread_id = ? .bind(thread_id.to_string()) .bind(expected_goal_id) .bind(expected_goal_id) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await? } (Some(status), None) => { @@ -346,7 +348,7 @@ WHERE thread_id = ? .bind(thread_id.to_string()) .bind(expected_goal_id) .bind(expected_goal_id) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await? } (None, Some(token_budget)) => { @@ -374,7 +376,7 @@ WHERE thread_id = ? .bind(thread_id.to_string()) .bind(expected_goal_id) .bind(expected_goal_id) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await? } (None, None) => { @@ -394,7 +396,7 @@ WHERE thread_id = ? .bind(thread_id.to_string()) .bind(expected_goal_id) .bind(expected_goal_id) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await? } else { let goal = self.get_thread_goal(thread_id).await?; @@ -459,7 +461,7 @@ WHERE thread_id = ? .bind(now_ms) .bind(thread_id.to_string()) .bind(status.as_str()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; if result.rows_affected() == 0 { @@ -490,7 +492,7 @@ RETURNING "#, ) .bind(thread_id.to_string()) - .fetch_optional(self.pool.as_ref()) + .fetch_optional_write(self.pool.as_ref()) .await?; row.map(|row| thread_goal_from_row(&row)).transpose() @@ -598,7 +600,10 @@ RETURNING "#, ); - let row = builder.build().fetch_optional(self.pool.as_ref()).await?; + let row = builder + .build() + .fetch_optional_write(self.pool.as_ref()) + .await?; let Some(row) = row else { return Ok(GoalAccountingOutcome::Unchanged( diff --git a/codex-rs/state/src/runtime/memories.rs b/codex-rs/state/src/runtime/memories.rs index 2cb66d1f8c..f11a779fb6 100644 --- a/codex-rs/state/src/runtime/memories.rs +++ b/codex-rs/state/src/runtime/memories.rs @@ -11,6 +11,8 @@ use crate::model::Stage1JobClaimOutcome; use crate::model::Stage1Output; use crate::model::Stage1StartupClaimParams; use crate::model::ThreadRow; +use crate::sqlite::SqlitePoolWriteExt; +use crate::sqlite::SqliteWriteQueryExt; use chrono::DateTime; use chrono::Duration; use sqlx::Executor; @@ -64,7 +66,7 @@ impl MemoryStore { } let now = Utc::now().timestamp(); - let mut tx = self.pool.begin().await?; + let mut tx = self.pool.begin_immediate().await?; let mut updated_rows = 0; for thread_id in thread_ids { @@ -299,7 +301,7 @@ FROM threads pub(super) async fn delete_thread_memory(&self, thread_id: ThreadId) -> anyhow::Result<()> { let now = Utc::now().timestamp(); let thread_id = thread_id.to_string(); - let mut tx = self.pool.begin().await?; + let mut tx = self.pool.begin_immediate().await?; let existing_output = sqlx::query( r#" @@ -428,7 +430,7 @@ WHERE thread_id IN ( ) .bind(cutoff) .bind(limit as i64) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await? .rows_affected(); @@ -652,7 +654,7 @@ WHERE id = ? AND memory_mode != 'polluted' "#, ) .bind(thread_id.as_str()) - .execute(self.state_pool.as_ref()) + .execute_with_busy_retry(self.state_pool.as_ref()) .await? .rows_affected(); @@ -694,7 +696,7 @@ WHERE id = ? AND memory_mode != 'polluted' let thread_id = thread_id.to_string(); let worker_id = worker_id.to_string(); - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut tx = self.pool.begin_immediate().await?; let existing_output = sqlx::query( r#" @@ -878,7 +880,7 @@ WHERE kind = ? AND job_key = ? let now = Utc::now().timestamp(); let thread_id = thread_id.to_string(); - let mut tx = self.pool.begin().await?; + let mut tx = self.pool.begin_immediate().await?; let rows_affected = sqlx::query( r#" UPDATE jobs @@ -955,7 +957,7 @@ WHERE excluded.source_updated_at >= stage1_outputs.source_updated_at let now = Utc::now().timestamp(); let thread_id = thread_id.to_string(); - let mut tx = self.pool.begin().await?; + let mut tx = self.pool.begin_immediate().await?; let rows_affected = sqlx::query( r#" UPDATE jobs @@ -1053,7 +1055,7 @@ WHERE kind = ? AND job_key = ? .bind(JOB_KIND_MEMORY_STAGE1) .bind(thread_id.as_str()) .bind(ownership_token) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await? .rows_affected(); @@ -1068,7 +1070,10 @@ WHERE kind = ? AND job_key = ? /// Phase 2 does not use this watermark as a dirty check; git workspace diffing /// decides whether consolidation work exists after the lock is claimed. pub async fn enqueue_global_consolidation(&self, input_watermark: i64) -> anyhow::Result<()> { - enqueue_global_consolidation_with_executor(self.pool.as_ref(), input_watermark).await + crate::sqlite::retry_on_busy(|| { + enqueue_global_consolidation_with_executor(self.pool.as_ref(), input_watermark) + }) + .await } /// Attempts to claim the global phase-2 consolidation lock. @@ -1095,7 +1100,7 @@ WHERE kind = ? AND job_key = ? let ownership_token = Uuid::new_v4().to_string(); let worker_id = worker_id.to_string(); - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut tx = self.pool.begin_immediate().await?; let existing_job = sqlx::query( r#" @@ -1240,7 +1245,7 @@ WHERE kind = ? AND job_key = ? .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(ownership_token) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await? .rows_affected(); @@ -1263,7 +1268,7 @@ WHERE kind = ? AND job_key = ? completed_watermark: i64, selected_outputs: &[Stage1Output], ) -> anyhow::Result { - let mut tx = self.pool.begin().await?; + let mut tx = self.pool.begin_immediate().await?; let rows_affected = mark_global_phase2_job_succeeded_row(&mut *tx, ownership_token, completed_watermark) .await?; @@ -1347,7 +1352,7 @@ WHERE kind = ? AND job_key = ? .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(ownership_token) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await? .rows_affected(); @@ -1389,7 +1394,7 @@ WHERE kind = ? AND job_key = ? .bind(JOB_KIND_MEMORY_CONSOLIDATE_GLOBAL) .bind(MEMORY_CONSOLIDATION_JOB_KEY) .bind(ownership_token) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await? .rows_affected(); @@ -1432,7 +1437,7 @@ WHERE kind = ? AND job_key = ? } pub(super) async fn clear_memory_data_in_pool(pool: &SqlitePool) -> anyhow::Result<()> { - let mut tx = pool.begin().await?; + let mut tx = pool.begin_immediate().await?; sqlx::query("UPDATE consolidation_progress SET max_thread_count = 0") .execute(&mut *tx) .await?; diff --git a/codex-rs/state/src/runtime/projects.rs b/codex-rs/state/src/runtime/projects.rs index 8a94d2e9f0..b21ce3397d 100644 --- a/codex-rs/state/src/runtime/projects.rs +++ b/codex-rs/state/src/runtime/projects.rs @@ -15,6 +15,7 @@ use crate::ProjectRoot; use crate::ProjectSortKey; use crate::ProjectsPage; use crate::SortDirection; +use crate::sqlite::SqlitePoolWriteExt; const PROJECT_SELECT: &str = "SELECT projects.*, (SELECT MAX(recency_at_ms) FROM threads @@ -27,7 +28,7 @@ impl StateRuntime { thread_id: &str, project_id: Option<&str>, ) -> anyhow::Result>> { - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut tx = self.pool.begin_immediate().await?; if let Some(project_id) = project_id { let exists = sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM projects WHERE id = ?") .bind(project_id) @@ -143,7 +144,7 @@ impl StateRuntime { thread_ids: &[String], idempotency_key: &str, ) -> anyhow::Result { - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut tx = self.pool.begin_immediate().await?; let existing_project_id = sqlx::query_scalar::<_, String>( "SELECT project_id FROM project_idempotency_keys WHERE key = ?", ) @@ -234,7 +235,7 @@ impl StateRuntime { roots: Option>, metadata: Option>, ) -> anyhow::Result> { - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut tx = self.pool.begin_immediate().await?; let row = QueryBuilder::::new(PROJECT_SELECT) .push(" WHERE id = ") .push_bind(id) @@ -291,7 +292,7 @@ impl StateRuntime { project_id: &str, before_project_id: Option<&str>, ) -> anyhow::Result> { - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut tx = self.pool.begin_immediate().await?; let mut project_ids = sqlx::query_scalar::<_, String>( "SELECT id FROM projects ORDER BY position ASC, id ASC", ) @@ -340,7 +341,7 @@ impl StateRuntime { &self, id: &str, ) -> anyhow::Result, Vec)>> { - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut tx = self.pool.begin_immediate().await?; let exists = sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM projects WHERE id = ?") .bind(id) .fetch_one(&mut *tx) diff --git a/codex-rs/state/src/runtime/queued_items.rs b/codex-rs/state/src/runtime/queued_items.rs index 3cbf32b6ee..9a6b387522 100644 --- a/codex-rs/state/src/runtime/queued_items.rs +++ b/codex-rs/state/src/runtime/queued_items.rs @@ -1,6 +1,8 @@ use super::*; use crate::MAX_QUEUE_ITEMS; use crate::QueuedUserSubmissionRecord; +use crate::sqlite::SqlitePoolWriteExt; +use crate::sqlite::SqliteWriteQueryExt; use sqlx::Connection; use tokio::sync::Mutex; use uuid::Uuid; @@ -99,7 +101,7 @@ impl SqliteQueueStore { .bind(now_ms) .bind(thread_id.to_string()) .bind(i64::try_from(MAX_QUEUE_ITEMS)?) - .fetch_one(self.pool.as_ref()) + .fetch_one_write(self.pool.as_ref()) .await?; QueuedUserSubmissionRecord::try_from_row(&row) } @@ -142,7 +144,7 @@ impl SqliteQueueStore { .bind(datetime_to_epoch_millis(Utc::now())) .bind(thread_id.to_string()) .bind(item_id) - .fetch_optional(self.pool.as_ref()) + .fetch_optional_write(self.pool.as_ref()) .await?; row.as_ref() .map(QueuedUserSubmissionRecord::try_from_row) @@ -154,7 +156,7 @@ impl SqliteQueueStore { sqlx::query("DELETE FROM queued_items WHERE thread_id = ? AND id = ?") .bind(thread_id.to_string()) .bind(item_id) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await? .rows_affected() > 0, @@ -162,7 +164,7 @@ impl SqliteQueueStore { } pub async fn reorder(&self, thread_id: ThreadId, ordered_ids: &[String]) -> anyhow::Result<()> { - let mut transaction = self.pool.begin().await?; + let mut transaction = self.pool.begin_immediate().await?; let rows: Vec<(String, i64)> = sqlx::query_as( "SELECT id, queue_order FROM queued_items WHERE thread_id = ? ORDER BY queue_order", @@ -203,7 +205,7 @@ impl SqliteQueueStore { pub(crate) async fn delete_thread_queue(&self, thread_id: ThreadId) -> anyhow::Result { Ok(sqlx::query("DELETE FROM queued_items WHERE thread_id = ?") .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await? .rows_affected() > 0) diff --git a/codex-rs/state/src/runtime/remote_control.rs b/codex-rs/state/src/runtime/remote_control.rs index f09bb015d3..0ab312d895 100644 --- a/codex-rs/state/src/runtime/remote_control.rs +++ b/codex-rs/state/src/runtime/remote_control.rs @@ -1,4 +1,5 @@ use super::*; +use crate::sqlite::SqliteWriteQueryExt; const REMOTE_CONTROL_APP_SERVER_CLIENT_NAME_NONE: &str = ""; @@ -97,7 +98,7 @@ ON CONFLICT(websocket_url, account_id, app_server_client_name) DO UPDATE SET .bind(&enrollment.server_name) .bind(enrollment.remote_control_enabled) .bind(Utc::now().timestamp()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(()) } @@ -123,7 +124,7 @@ WHERE websocket_url = ? AND account_id = ? AND app_server_client_name = ? .bind(remote_control_app_server_client_name_key( app_server_client_name, )) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected()) } @@ -145,7 +146,7 @@ WHERE websocket_url = ? AND account_id = ? AND app_server_client_name = ? .bind(remote_control_app_server_client_name_key( app_server_client_name, )) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected()) } diff --git a/codex-rs/state/src/runtime/rollout_migration.rs b/codex-rs/state/src/runtime/rollout_migration.rs index bc1c55720e..8c26474c7e 100644 --- a/codex-rs/state/src/runtime/rollout_migration.rs +++ b/codex-rs/state/src/runtime/rollout_migration.rs @@ -5,6 +5,7 @@ //! old rollouts on every startup while still retrying skipped files that later change. use super::*; +use crate::sqlite::SqliteWriteQueryExt; impl StateRuntime { pub async fn get_rollout_migration_state( @@ -70,7 +71,7 @@ WHERE excluded.last_checked_thread_created_at IS NOT NULL .bind(thread_created_at) .bind(thread_id) .bind(now) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(()) } @@ -124,7 +125,7 @@ ON CONFLICT(migration_id, rollout_path) DO UPDATE SET .bind(skipped_rollout.rollout_modified_at_ns) .bind(skipped_rollout.skip_reason.as_str()) .bind(now) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(()) } @@ -142,7 +143,7 @@ WHERE migration_id = ? AND rollout_path = ? ) .bind(migration_id) .bind(rollout_path) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(()) } diff --git a/codex-rs/state/src/runtime/thread_attachments.rs b/codex-rs/state/src/runtime/thread_attachments.rs index e776015469..9dc8f67d59 100644 --- a/codex-rs/state/src/runtime/thread_attachments.rs +++ b/codex-rs/state/src/runtime/thread_attachments.rs @@ -13,6 +13,7 @@ use crate::ThreadAttachmentArchiveFilter; use crate::ThreadAttachmentOwner; use crate::ThreadAttachmentOwnerPage; use crate::ThreadAttachmentPage; +use crate::sqlite::SqlitePoolWriteExt; use anyhow::Context; use chrono::Utc; use codex_protocol::ThreadId; @@ -32,7 +33,7 @@ impl StateRuntime { source_thread_id: ThreadId, destination_thread_id: ThreadId, ) -> anyhow::Result<()> { - let mut transaction = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut transaction = self.pool.begin_immediate().await?; let destination = destination_thread_id.to_string(); let exists = sqlx::query_scalar::<_, i64>("SELECT 1 FROM threads WHERE id = ?") .bind(&destination) @@ -84,7 +85,7 @@ impl StateRuntime { ); } - let mut transaction = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut transaction = self.pool.begin_immediate().await?; let thread_id_string = thread_id.to_string(); let thread_exists = sqlx::query_scalar::<_, i64>("SELECT 1 FROM threads WHERE id = ?") .bind(&thread_id_string) @@ -152,7 +153,7 @@ impl StateRuntime { identity_key: &str, ) -> anyhow::Result { validate_attachment_identity(attachment_type, identity_key)?; - let mut transaction = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut transaction = self.pool.begin_immediate().await?; let thread_id_string = thread_id.to_string(); let thread_exists = sqlx::query_scalar::<_, i64>("SELECT 1 FROM threads WHERE id = ?") .bind(&thread_id_string) diff --git a/codex-rs/state/src/runtime/thread_section_order.rs b/codex-rs/state/src/runtime/thread_section_order.rs index f5b919847f..b5aa3ee625 100644 --- a/codex-rs/state/src/runtime/thread_section_order.rs +++ b/codex-rs/state/src/runtime/thread_section_order.rs @@ -1,4 +1,5 @@ use super::StateRuntime; +use crate::sqlite::SqlitePoolWriteExt; use chrono::DateTime; use chrono::Utc; use codex_protocol::ThreadId; @@ -113,7 +114,7 @@ LIMIT ? )); } - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut tx = self.pool.begin_immediate().await?; let thread_id = thread_id.to_string(); let current_section = sqlx::query_scalar::<_, Option>( "SELECT thread_section_id FROM threads WHERE id = ?", diff --git a/codex-rs/state/src/runtime/thread_sections.rs b/codex-rs/state/src/runtime/thread_sections.rs index f8c6f02378..71bb203107 100644 --- a/codex-rs/state/src/runtime/thread_sections.rs +++ b/codex-rs/state/src/runtime/thread_sections.rs @@ -2,6 +2,8 @@ use super::StateRuntime; use crate::PINNED_THREAD_SECTION_ID; use crate::ThreadSection; use crate::ThreadSectionAppearance; +use crate::sqlite::SqlitePoolWriteExt; +use crate::sqlite::SqliteWriteQueryExt; use uuid::Uuid; impl StateRuntime { @@ -27,7 +29,7 @@ impl StateRuntime { .map(serde_json::to_string) .transpose()?, ) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(section) @@ -56,7 +58,7 @@ impl StateRuntime { .bind(replace_appearance) .bind(appearance) .bind(id) - .fetch_optional(self.pool.as_ref()) + .fetch_optional_write(self.pool.as_ref()) .await?; section.map(ThreadSection::from_row).transpose() @@ -68,7 +70,7 @@ impl StateRuntime { anyhow::bail!("built-in pinned thread section cannot be deleted"); } - let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?; + let mut tx = self.pool.begin_immediate().await?; sqlx::query( "UPDATE threads SET section_position = NULL, section_entered_at_ms = NULL WHERE thread_section_id = ?", ) diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 1179126f62..96704dfa08 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -1,5 +1,7 @@ use super::*; use crate::SortDirection; +use crate::sqlite::SqlitePoolWriteExt; +use crate::sqlite::SqliteWriteQueryExt; use codex_protocol::SanitizedGitUrl; use codex_protocol::protocol::SessionSource; use std::sync::atomic::AtomicI64; @@ -94,7 +96,7 @@ WHERE id = ? .bind(crate::GUARDIAN_THREAD_TITLE) .bind(legacy_name) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected() > 0) } @@ -124,7 +126,7 @@ WHERE id = ? AND preview = '' ) .bind(preview) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected() > 0) } @@ -151,7 +153,7 @@ ON CONFLICT(child_thread_id) DO UPDATE SET .bind(parent_thread_id.to_string()) .bind(child_thread_id.to_string()) .bind(status.as_ref()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(()) } @@ -165,7 +167,7 @@ ON CONFLICT(child_thread_id) DO UPDATE SET sqlx::query("UPDATE thread_spawn_edges SET status = ? WHERE child_thread_id = ?") .bind(status.as_ref()) .bind(child_thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(()) } @@ -362,7 +364,7 @@ ON CONFLICT(child_thread_id) DO NOTHING .bind(parent_thread_id.to_string()) .bind(child_thread_id.to_string()) .bind(crate::DirectionalThreadSpawnEdgeStatus::Open.as_ref()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(()) } @@ -418,7 +420,7 @@ ON CONFLICT(child_thread_id) DO NOTHING .bind(replacement.display().to_string()) .bind(id.to_string()) .bind(expected.display().to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected() == 1) } @@ -716,7 +718,7 @@ ON CONFLICT(id) DO NOTHING .bind("enabled") .bind(metadata.project_id.as_deref()) .bind(metadata.daybreak_enabled) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; self.insert_thread_spawn_edge_from_source_if_absent(metadata.id, metadata.source.as_str()) .await?; @@ -732,7 +734,7 @@ ON CONFLICT(id) DO NOTHING let result = sqlx::query("UPDATE threads SET daybreak_enabled = ? WHERE id = ?") .bind(daybreak_enabled) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected() > 0) } @@ -745,7 +747,7 @@ ON CONFLICT(id) DO NOTHING let result = sqlx::query("UPDATE threads SET memory_mode = ? WHERE id = ?") .bind(memory_mode) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected() > 0) } @@ -758,7 +760,7 @@ ON CONFLICT(id) DO NOTHING let result = sqlx::query("UPDATE threads SET title = ? WHERE id = ?") .bind(title) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected() > 0) } @@ -771,7 +773,7 @@ ON CONFLICT(id) DO NOTHING let result = sqlx::query("UPDATE threads SET name = ? WHERE id = ?") .bind(name) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected() > 0) } @@ -787,7 +789,7 @@ ON CONFLICT(id) DO NOTHING .bind(datetime_to_epoch_seconds(updated_at)) .bind(datetime_to_epoch_millis(updated_at)) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected() > 0) } @@ -813,7 +815,7 @@ WHERE id = ? .bind(recency_at_millis) .bind(recency_at_millis) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected() > 0) } @@ -902,7 +904,7 @@ WHERE id = ? .bind(git_origin_url.is_some()) .bind(git_origin_url.flatten().map(SanitizedGitUrl::as_str)) .bind(thread_id.to_string()) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; Ok(result.rows_affected() > 0) } @@ -1052,7 +1054,7 @@ ON CONFLICT(id) DO UPDATE SET .bind(creation_memory_mode.unwrap_or("enabled")) .bind(metadata.project_id.as_deref()) .bind(metadata.daybreak_enabled) - .execute(self.pool.as_ref()) + .execute_with_busy_retry(self.pool.as_ref()) .await?; self.insert_thread_spawn_edge_from_source_if_absent(metadata.id, metadata.source.as_str()) .await?; @@ -1173,14 +1175,14 @@ ON CONFLICT(id) DO UPDATE SET for (thread_id, thread_id_string) in thread_ids.iter().zip(&thread_id_strings) { sqlx::query("DELETE FROM logs WHERE thread_id = ?") .bind(thread_id_string) - .execute(self.logs_pool.as_ref()) + .execute_with_busy_retry(self.logs_pool.as_ref()) .await?; self.thread_queue.delete_thread_queue(*thread_id).await?; self.delete_versioned_thread_memory(*thread_id).await?; self.thread_goals.delete_thread_goal(*thread_id).await?; } - let mut tx = self.pool.begin().await?; + let mut tx = self.pool.begin_immediate().await?; for thread_id_string in &thread_id_strings { sqlx::query("DELETE FROM thread_dynamic_tools WHERE thread_id = ?") .bind(thread_id_string) diff --git a/codex-rs/state/src/sqlite.rs b/codex-rs/state/src/sqlite.rs index e35a074d68..d8c72c6031 100644 --- a/codex-rs/state/src/sqlite.rs +++ b/codex-rs/state/src/sqlite.rs @@ -1,5 +1,6 @@ //! Shared SQLite connection configuration. //! Writable pools preserve existing vacuum modes and report initialization errors to callers. +//! External lock contention while opening gets a bounded retry before the original error surfaces. #![expect( clippy::disallowed_methods, @@ -25,8 +26,12 @@ use std::path::PathBuf; use std::time::Duration; use std::time::Instant; +mod busy_retry; mod validation; +pub use busy_retry::SqlitePoolWriteExt; +pub(crate) use busy_retry::SqliteWriteQueryExt; +pub(crate) use busy_retry::retry_on_busy; use validation::SqliteQuickCheckManager; const DEFAULT_BUSY_TIMEOUT: Duration = Duration::from_secs(5); @@ -409,7 +414,7 @@ impl SqliteConfig { result = pool_options.connect_with(options.clone()) => result, } }; - let pool = connect().await?; + let pool = retry_on_busy(|| async { Ok(connect().await?) }).await?; let validation_result = self .quick_check_manager .quick_check_once( @@ -439,7 +444,7 @@ impl SqliteConfig { ); } - return connect().await.map_err(anyhow::Error::from); + return retry_on_busy(|| async { Ok(connect().await?) }).await; } } diff --git a/codex-rs/state/src/sqlite/busy_retry.rs b/codex-rs/state/src/sqlite/busy_retry.rs new file mode 100644 index 0000000000..9ed22a7b4e --- /dev/null +++ b/codex-rs/state/src/sqlite/busy_retry.rs @@ -0,0 +1,191 @@ +//! Bounded retries for SQLite operations that can safely restart on external contention. +//! Callers must complete one atomic transaction (or only open a connection) per attempt. +//! Never retry cancellation: a commit whose result was not received may have succeeded. + +use std::error::Error as StdError; +use std::future::Future; +use std::time::Duration; + +use sqlx::Execute; +use sqlx::Executor; +use sqlx::FromRow; +use sqlx::Sqlite; +use sqlx::SqlitePool; +use sqlx::Transaction; +use sqlx::error::DatabaseError; +use sqlx::query::Query; +use sqlx::query::QueryAs; +use sqlx::sqlite::SqliteArguments; +use sqlx::sqlite::SqliteError; +use sqlx::sqlite::SqliteQueryResult; +use sqlx::sqlite::SqliteRow; + +/// SQLite writer admission for pools shared by the state and thread stores. +/// Implementations must acquire the writer slot before any read snapshot and +/// retry only admission, never an already-started transaction or external effects. +pub trait SqlitePoolWriteExt { + fn begin_immediate( + &self, + ) -> impl Future, sqlx::Error>> + Send; +} + +impl SqlitePoolWriteExt for SqlitePool { + async fn begin_immediate(&self) -> Result, sqlx::Error> { + retry_when( + || self.begin_with("BEGIN IMMEDIATE"), + |error| is_busy(error), + ) + .await + } +} + +/// Safe retry boundaries for single prepared SQLite writes, including RETURNING. +/// Implementations decode their row before committing, just as SQLx does in an +/// explicit transaction. Never use this for multiple statements or existing +/// transactions; those must keep using their transaction's connection. +pub(crate) trait SqliteWriteQueryExt<'q>: Execute<'q, Sqlite> + 'q { + type Output: Send; + + fn decode_returning_row(row: SqliteRow) -> Result; + + /// Retry a standalone autocommit write with exactly the same bound inputs. + fn execute_with_busy_retry<'e>( + mut self, + pool: &'e SqlitePool, + ) -> impl Future> + Send + 'e + where + 'q: 'e, + { + async move { + let persistent = Execute::persistent(&self); + let arguments = self + .take_arguments() + .map_err(sqlx::Error::Encode)? + .unwrap_or_default(); + let sql = self.sql(); + retry_when( + || { + sqlx::query_with(sql.clone(), arguments.clone()) + .persistent(persistent) + .execute(pool) + }, + |error| is_busy(error), + ) + .await + } + } + + /// Reserve the writer slot for a RETURNING write; expose its row only after commit. + fn fetch_one_write<'e>( + self, + pool: &'e SqlitePool, + ) -> impl Future> + Send + 'e + where + 'q: 'e, + { + async move { + let mut tx = pool.begin_immediate().await?; + let row = (&mut *tx).fetch_one(self).await?; + let output = Self::decode_returning_row(row)?; + tx.commit().await?; + Ok(output) + } + } + + /// As with fetch_one_write, but commit a write that returns no row too. + fn fetch_optional_write<'e>( + self, + pool: &'e SqlitePool, + ) -> impl Future, sqlx::Error>> + Send + 'e + where + 'q: 'e, + { + async move { + let mut tx = pool.begin_immediate().await?; + let row = (&mut *tx).fetch_optional(self).await?; + let output = row.map(Self::decode_returning_row).transpose()?; + tx.commit().await?; + Ok(output) + } + } +} + +impl<'q> SqliteWriteQueryExt<'q> for Query<'q, Sqlite, SqliteArguments> { + type Output = SqliteRow; + + fn decode_returning_row(row: SqliteRow) -> Result { + Ok(row) + } +} + +impl<'q, O> SqliteWriteQueryExt<'q> for QueryAs<'q, Sqlite, O, SqliteArguments> +where + O: Send + Unpin + for<'r> FromRow<'r, SqliteRow> + 'q, +{ + type Output = O; + + fn decode_returning_row(row: SqliteRow) -> Result { + O::from_row(&row) + } +} + +/// Retries only received SQLite BUSY errors. Each operation must drop or roll back +/// its transaction before returning an error so even BUSY_SNAPSHOT starts fresh. +/// SQLite already waits up to five seconds per lock; allow two more attempts. +pub(crate) async fn retry_on_busy(operation: F) -> anyhow::Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + retry_when(operation, |error| error.chain().any(is_busy)).await +} + +async fn retry_when( + mut operation: F, + is_retryable: impl Fn(&E) -> bool, +) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut delays = [Duration::from_millis(100), Duration::from_millis(200)].into_iter(); + loop { + match operation().await { + Ok(value) => return Ok(value), + Err(error) => { + if is_retryable(&error) + && let Some(delay) = delays.next() + { + tokio::time::sleep(delay).await; + continue; + } + return Err(error); + } + } + } +} + +fn is_busy(source: &(dyn StdError + 'static)) -> bool { + let Some(sqlite_error) = source + .downcast_ref::() + .and_then(sqlx::Error::as_database_error) + .and_then(|database_error| database_error.try_downcast_ref::()) + else { + return false; + }; + matches!( + sqlite_error + .code() + .and_then(|code| code.parse::().ok()), + Some( + libsqlite3_sys::SQLITE_BUSY + | libsqlite3_sys::SQLITE_BUSY_RECOVERY + | libsqlite3_sys::SQLITE_BUSY_SNAPSHOT + | libsqlite3_sys::SQLITE_BUSY_TIMEOUT + ) + ) +} + +#[cfg(test)] +#[path = "retry_tests.rs"] +mod tests; diff --git a/codex-rs/state/src/sqlite/retry_tests.rs b/codex-rs/state/src/sqlite/retry_tests.rs new file mode 100644 index 0000000000..e597f83e21 --- /dev/null +++ b/codex-rs/state/src/sqlite/retry_tests.rs @@ -0,0 +1,267 @@ +//! Tests safe retries against real SQLite writer contention. + +use super::retry_on_busy; +use crate::SqliteConfig; +use crate::SqliteQueueStore; +use crate::migrations::runtime_logs_migrator; +use crate::migrations::runtime_queue_migrator; +use crate::runtime::test_support::unique_temp_dir; +use anyhow::Context; +use codex_protocol::ThreadId; +use codex_utils_absolute_path::test_support::PathExt; +use pretty_assertions::assert_eq; +use sqlx::sqlite::SqlitePoolOptions; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Mutex; +use tokio::sync::mpsc; + +#[tokio::test] +#[expect( + clippy::await_holding_invalid_type, + reason = "Only this sequential retry closure owns these connections. The guards keep the real competing SQLite lock alive across attempts; no other task takes either guard." +)] +async fn retries_after_a_competing_writer_releases_its_lock() -> anyhow::Result<()> { + let sqlite_home = unique_temp_dir(); + tokio::fs::create_dir_all(&sqlite_home).await?; + let _cleanup = scopeguard::guard(sqlite_home.clone(), |sqlite_home| { + let _ = std::fs::remove_dir_all(sqlite_home); + }); + let sqlite = SqliteConfig::new_for_testing(sqlite_home.as_path().abs()); + let pool = sqlite + .open_logs_db(&runtime_logs_migrator(), /*telemetry_override*/ None) + .await?; + let mut competing_writer = pool.acquire().await?; + let mut writer = pool.acquire().await?; + sqlx::query("PRAGMA busy_timeout=0") + .execute(&mut *writer) + .await?; + sqlx::query("BEGIN IMMEDIATE") + .execute(&mut *competing_writer) + .await?; + let competing_writer = Mutex::new(competing_writer); + let writer = Mutex::new(writer); + + retry_on_busy(|| async { + let mut writer = writer.lock().await; + let result = sqlx::query( + "INSERT INTO logs (ts, ts_nanos, level, target, feedback_log_body) \ + VALUES (1, 0, 'INFO', 'retry-test', 'saved once')", + ) + .execute(&mut **writer) + .await; + // Release the competing transaction only after SQLite actually rejected + // the write. This cannot pass by racing the first attempt's busy wait. + if result.is_err() { + let mut competing_writer = competing_writer.lock().await; + sqlx::query("ROLLBACK") + .execute(&mut **competing_writer) + .await?; + } + result.map_err(anyhow::Error::from) + }) + .await?; + + let rows = sqlx::query_scalar::<_, String>("SELECT feedback_log_body FROM logs ORDER BY id") + .fetch_all(&pool) + .await?; + assert_eq!(rows, vec!["saved once".to_string()]); + drop(writer); + drop(competing_writer); + pool.close().await; + Ok(()) +} + +#[tokio::test] +async fn queue_enqueue_survives_received_busy_without_duplicating_the_message() -> anyhow::Result<()> +{ + let sqlite_home = unique_temp_dir(); + tokio::fs::create_dir_all(&sqlite_home).await?; + let _cleanup = scopeguard::guard(sqlite_home.clone(), |sqlite_home| { + let _ = std::fs::remove_dir_all(sqlite_home); + }); + let sqlite = SqliteConfig::new_for_testing(sqlite_home.as_path().abs()); + let pool = sqlite + .open_queue_db(&runtime_queue_migrator(), /*telemetry_override*/ None) + .await?; + let (attempt_finished, mut finished) = mpsc::unbounded_channel(); + // No connection is opened or released until enqueue's first attempt. + // SQLite must have returned BUSY before after_release can run: the other + // writer is still holding the lock, and this connection never waits for it. + let writer_pool = SqlitePoolOptions::new() + .max_connections(/*max*/ 1) + .after_release(move |_, _| { + let _ = attempt_finished.send(()); + Box::pin(async { Ok(true) }) + }) + .connect_lazy_with( + (*pool.connect_options()) + .clone() + .busy_timeout(Duration::ZERO), + ); + let queue = SqliteQueueStore::new(Arc::new(writer_pool)); + let observer = SqliteQueueStore::new(Arc::new(pool.clone())); + let thread_id = ThreadId::new(); + let mut competing_writer = pool.acquire().await?; + sqlx::query("BEGIN IMMEDIATE") + .execute(&mut *competing_writer) + .await?; + let (item, ()) = tokio::try_join!( + queue.enqueue(thread_id, r#"{"text":"save this bound value"}"#), + async { + finished + .recv() + .await + .context("no completed writer attempt")?; + sqlx::query("ROLLBACK") + .execute(&mut *competing_writer) + .await?; + Ok::<(), anyhow::Error>(()) + } + )?; + assert_eq!( + (thread_id, r#"{"text":"save this bound value"}"#), + (item.thread_id, item.payload.as_str()) + ); + assert_eq!( + vec![item], + observer + .list_page(thread_id, /*offset*/ 0, /*limit*/ 10) + .await? + ); + drop(competing_writer); + queue.close().await; + pool.close().await; + Ok(()) +} + +#[tokio::test] +async fn queue_update_survives_received_busy_and_commits_once() -> anyhow::Result<()> { + let sqlite_home = unique_temp_dir(); + tokio::fs::create_dir_all(&sqlite_home).await?; + let _cleanup = scopeguard::guard(sqlite_home.clone(), |sqlite_home| { + let _ = std::fs::remove_dir_all(sqlite_home); + }); + let sqlite = SqliteConfig::new_for_testing(sqlite_home.as_path().abs()); + let pool = sqlite + .open_queue_db(&runtime_queue_migrator(), /*telemetry_override*/ None) + .await?; + let observer = SqliteQueueStore::new(Arc::new(pool.clone())); + let thread_id = ThreadId::new(); + let mut expected = observer + .enqueue(thread_id, r#"{"text":"original"}"#) + .await?; + let revision_before = observer.changes_since(/*revision*/ 0, &[thread_id]).await?[0].1; + expected.payload = r#"{"text":"edited bound value"}"#.to_string(); + let (attempt_finished, mut finished) = mpsc::unbounded_channel(); + let writer_pool = SqlitePoolOptions::new() + .max_connections(/*max*/ 1) + .after_release(move |_, _| { + let _ = attempt_finished.send(()); + Box::pin(async { Ok(true) }) + }) + .connect_lazy_with( + (*pool.connect_options()) + .clone() + .busy_timeout(Duration::ZERO), + ); + let queue = SqliteQueueStore::new(Arc::new(writer_pool)); + let mut competing_writer = pool.acquire().await?; + sqlx::query("BEGIN IMMEDIATE") + .execute(&mut *competing_writer) + .await?; + let (updated, ()) = tokio::try_join!( + queue.update(thread_id, &expected.id, &expected.payload), + async { + // Release only after update actually failed admission, not while + // SQLite's busy handler could still turn the first attempt into success. + finished + .recv() + .await + .context("no completed writer attempt")?; + sqlx::query("ROLLBACK") + .execute(&mut *competing_writer) + .await?; + Ok::<(), anyhow::Error>(()) + } + )?; + assert_eq!(Some(expected.clone()), updated); + assert_eq!( + (vec![expected], vec![(thread_id, revision_before + 1)]), + ( + observer + .list_page(thread_id, /*offset*/ 0, /*limit*/ 10) + .await?, + observer + .changes_since(revision_before, &[thread_id]) + .await? + ) + ); + drop(competing_writer); + queue.close().await; + pool.close().await; + Ok(()) +} + +#[tokio::test] +async fn queue_delete_survives_received_busy_and_preserves_other_messages() -> anyhow::Result<()> { + let sqlite_home = unique_temp_dir(); + tokio::fs::create_dir_all(&sqlite_home).await?; + let _cleanup = scopeguard::guard(sqlite_home.clone(), |sqlite_home| { + let _ = std::fs::remove_dir_all(sqlite_home); + }); + let sqlite = SqliteConfig::new_for_testing(sqlite_home.as_path().abs()); + let pool = sqlite + .open_queue_db(&runtime_queue_migrator(), /*telemetry_override*/ None) + .await?; + let observer = SqliteQueueStore::new(Arc::new(pool.clone())); + let thread_id = ThreadId::new(); + let target = observer.enqueue(thread_id, r#"{"text":"remove"}"#).await?; + let survivor = observer.enqueue(thread_id, r#"{"text":"keep"}"#).await?; + let revision_before = observer.changes_since(/*revision*/ 0, &[thread_id]).await?[0].1; + let (attempt_finished, mut finished) = mpsc::unbounded_channel(); + let writer_pool = SqlitePoolOptions::new() + .max_connections(/*max*/ 1) + .after_release(move |_, _| { + let _ = attempt_finished.send(()); + Box::pin(async { Ok(true) }) + }) + .connect_lazy_with( + (*pool.connect_options()) + .clone() + .busy_timeout(Duration::ZERO), + ); + let queue = SqliteQueueStore::new(Arc::new(writer_pool)); + let mut competing_writer = pool.acquire().await?; + sqlx::query("BEGIN IMMEDIATE") + .execute(&mut *competing_writer) + .await?; + let (deleted, ()) = tokio::try_join!(queue.delete(thread_id, &target.id), async { + // The first autocommit DELETE has returned BUSY and released its + // connection. Its bound thread and item IDs must survive the retry. + finished + .recv() + .await + .context("no completed writer attempt")?; + sqlx::query("ROLLBACK") + .execute(&mut *competing_writer) + .await?; + Ok::<(), anyhow::Error>(()) + })?; + assert!(deleted); + assert_eq!( + (vec![survivor], vec![(thread_id, revision_before + 1)]), + ( + observer + .list_page(thread_id, /*offset*/ 0, /*limit*/ 10) + .await?, + observer + .changes_since(revision_before, &[thread_id]) + .await? + ) + ); + drop(competing_writer); + queue.close().await; + pool.close().await; + Ok(()) +} diff --git a/codex-rs/thread-store/src/local/thread_history.rs b/codex-rs/thread-store/src/local/thread_history.rs index ee59619a13..7be33d6720 100644 --- a/codex-rs/thread-store/src/local/thread_history.rs +++ b/codex-rs/thread-store/src/local/thread_history.rs @@ -4,6 +4,7 @@ use codex_app_server_protocol::TurnStatus; use codex_protocol::ThreadId; use codex_protocol::models::MessagePhase; use codex_protocol::realtime::RealtimeItem; +use codex_state::SqlitePoolWriteExt; use super::LocalThreadStore; use crate::ThreadStoreError; @@ -112,10 +113,7 @@ pub(super) async fn apply_projection( // Write the projected rows and advance the JSONL offset and ordinal in one transaction. If // SQLite fails, it stays behind the durable rollout instead of claiming data it did not // materialize. - let mut transaction = pool - .begin_with("BEGIN IMMEDIATE") - .await - .map_err(thread_history_error)?; + let mut transaction = pool.begin_immediate().await.map_err(thread_history_error)?; let thread_id = thread_id.to_string(); let projection_state = sqlx::query_as::<_, (i64, i64)>( r#" @@ -255,7 +253,7 @@ pub(super) async fn delete_thread( let pool = store.thread_history_db().await?; let mut transaction = pool - .begin_with("BEGIN IMMEDIATE") + .begin_immediate() .await .map_err(thread_history_delete_error)?; let thread_id = thread_id.to_string();