github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 16:17:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Retry SQLite operations on transient writer contention (#54369)

Browse Source 
## 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
This commit is contained in:
willwang-openai
2026-10-15 18:06:57 +00:00
committed by copyberry
parent c313eb75fd
commit 7ae8319917
21 changed files with 572 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/ext/agent-message-board/src/local/lifecycle.rs
+3 -9
View File
@@ -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<Transaction<'_, Sqlite>> {
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())
codex-rs/ext/agent-message-board/src/local/pool.rs
+2 -4
View File
@@ -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<Arc<SqlitePool>> {
.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
codex-rs/state/src/lib.rs
+1
View File
@@ -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;
codex-rs/state/src/migrations.rs
+2 -1
View File
@@ -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(())
}
codex-rs/state/src/runtime.rs
+2 -1
View File
@@ -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(())
}
codex-rs/state/src/runtime/backfill.rs
+5 -4
View File
@@ -1,4 +1,5 @@
use super::*;
use crate::sqlite::SqliteWriteQueryExt;
impl StateRuntime {
pub async fn get_backfill_state(&self) -> anyhow::Result<crate::BackfillState> {
@@ -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(())
}
codex-rs/state/src/runtime/external_agent_config_imports.rs
+2 -1
View File
@@ -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(())
codex-rs/state/src/runtime/goals.rs
+16 -11
View File
@@ -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(
codex-rs/state/src/runtime/memories.rs
+20 -15
View File
@@ -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<bool> {
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?;
codex-rs/state/src/runtime/projects.rs
+6 -5
View File
@@ -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<Option<Option<String>>> {
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<CreatedProject> {
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<Vec<ProjectRoot>>,
metadata: Option<BTreeMap<String, String>>,
) -> anyhow::Result<Option<(Project, bool)>> {
let mut tx = self.pool.begin_with("BEGIN IMMEDIATE").await?;
let mut tx = self.pool.begin_immediate().await?;
let row = QueryBuilder::<Sqlite>::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<Option<bool>> {
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<Option<(Vec<String>, Vec<String>)>> {
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)
codex-rs/state/src/runtime/queued_items.rs
+7 -5
View File
@@ -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<bool> {
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)
codex-rs/state/src/runtime/remote_control.rs
+4 -3
View File
@@ -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())
}
codex-rs/state/src/runtime/rollout_migration.rs
+4 -3
View File
@@ -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(())
}
codex-rs/state/src/runtime/thread_attachments.rs
+4 -3
View File
@@ -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<RemoveThreadAttachmentOutcome> {
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)
codex-rs/state/src/runtime/thread_section_order.rs
+2 -1
View File
@@ -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<String>>(
"SELECT thread_section_id FROM threads WHERE id = ?",
codex-rs/state/src/runtime/thread_sections.rs
+5 -3
View File
@@ -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 = ?",
)
codex-rs/state/src/runtime/threads.rs
+19 -17
View File
@@ -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)
codex-rs/state/src/sqlite.rs
+7 -2
View File
@@ -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;
}
}
codex-rs/state/src/sqlite/busy_retry.rs
+191
View File
@@ -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<Output = Result<Transaction<'static, Sqlite>, sqlx::Error>> + Send;
}
impl SqlitePoolWriteExt for SqlitePool {
async fn begin_immediate(&self) -> Result<Transaction<'static, Sqlite>, 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<Self::Output, sqlx::Error>;
/// Retry a standalone autocommit write with exactly the same bound inputs.
fn execute_with_busy_retry<'e>(
mut self,
pool: &'e SqlitePool,
) -> impl Future<Output = Result<SqliteQueryResult, sqlx::Error>> + 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<Output = Result<Self::Output, sqlx::Error>> + 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<Output = Result<Option<Self::Output>, 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<Self::Output, sqlx::Error> {
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<Self::Output, sqlx::Error> {
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<T, F, Fut>(operation: F) -> anyhow::Result<T>
where
F: FnMut() -> Fut,
Fut: Future<Output = anyhow::Result<T>>,
{
retry_when(operation, |error| error.chain().any(is_busy)).await
}
async fn retry_when<T, E, F, Fut>(
mut operation: F,
is_retryable: impl Fn(&E) -> bool,
) -> Result<T, E>
where
F: FnMut() -> Fut,
Fut: Future<Output = Result<T, E>>,
{
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::<sqlx::Error>()
.and_then(sqlx::Error::as_database_error)
.and_then(|database_error| database_error.try_downcast_ref::<SqliteError>())
else {
return false;
};
matches!(
sqlite_error
.code()
.and_then(|code| code.parse::<i32>().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;
codex-rs/state/src/sqlite/retry_tests.rs
+267
View File
@@ -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(())
}
codex-rs/thread-store/src/local/thread_history.rs
+3 -5
View File
@@ -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();