From b7b318366d882887c088b8a9b66de5632a13b7a2 Mon Sep 17 00:00:00 2001 From: Eric Curtin Date: Sun, 11 Oct 2026 21:33:30 +0100 Subject: [PATCH 1/2] fix(server): return SQLite write versions atomically Use RETURNING instead of a follow-up read. Signed-off-by: Eric Curtin --- .../src/persistence/sqlite.rs | 67 +++++++-------- .../openshell-server/src/persistence/tests.rs | 86 +++++++++++++++++++ 2 files changed, 115 insertions(+), 38 deletions(-) diff --git a/crates/openshell-server/src/persistence/sqlite.rs b/crates/openshell-server/src/persistence/sqlite.rs index ce269cd241..2d6012e1d6 100644 --- a/crates/openshell-server/src/persistence/sqlite.rs +++ b/crates/openshell-server/src/persistence/sqlite.rs @@ -423,12 +423,12 @@ ON CONFLICT ("object_type", "workspace", "name") WHERE "name" IS NOT NULL DO UPD .await } WriteCondition::MatchResourceVersion(expected_version) => { - // Update with version check - let result = sqlx::query( + let row = sqlx::query( r#" UPDATE "objects" SET "payload" = ?4, "labels" = ?5, "updated_at_ms" = ?6, "resource_version" = "resource_version" + 1 WHERE "object_type" = ?1 AND "id" = ?2 AND "resource_version" = ?3 +RETURNING "resource_version", "created_at_ms", "updated_at_ms" "#, ) .bind(object_type) @@ -437,37 +437,28 @@ WHERE "object_type" = ?1 AND "id" = ?2 AND "resource_version" = ?3 .bind(payload) .bind(labels.unwrap_or("{}")) .bind(now_ms) - .execute(&self.pool) + .fetch_optional(&self.pool) .await .map_err(|e| map_db_error(&e))?; - if result.rows_affected() == 0 { - // The version-matched UPDATE matched no row. Distinguish a - // version mismatch (row present, different version) from an - // absent row (deleted / never existed). Both are CAS - // precondition failures, so report them as typed `Conflict` - // rather than a backend-dependent error string: absent rows - // carry `current_resource_version: None`. - let existing = self.get(object_type, id).await?; - return Err(PersistenceError::Conflict { - current_resource_version: existing.map(|record| record.resource_version), - }); + if let Some(row) = row { + return Ok(write_result_from_row(&row)); } - // Fetch the updated record to get the new resource_version - let updated = self.get(object_type, id).await?.ok_or_else(|| { - PersistenceError::Database("object disappeared after update".to_string()) - })?; - - Ok(WriteResult { - resource_version: updated.resource_version, - created_at_ms: updated.created_at_ms, - updated_at_ms: updated.updated_at_ms, + // The version-matched UPDATE matched no row. Distinguish a + // version mismatch (row present, different version) from an + // absent row (deleted / never existed). Both are CAS + // precondition failures, so report them as typed `Conflict` + // rather than a backend-dependent error string: absent rows + // carry `current_resource_version: None`. + let existing = self.get(object_type, id).await?; + Err(PersistenceError::Conflict { + current_resource_version: existing.map(|record| record.resource_version), }) } WriteCondition::Unconditional => { // Unconditional upsert by name - sqlx::query( + let row = sqlx::query( r#" INSERT INTO "objects" ("object_type", "id", "name", "workspace", "payload", "created_at_ms", "updated_at_ms", "labels", "resource_version") VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?6, ?7, 1) @@ -476,6 +467,7 @@ ON CONFLICT ("object_type", "workspace", "name") WHERE "name" IS NOT NULL DO UPD "updated_at_ms" = excluded."updated_at_ms", "labels" = excluded."labels", "resource_version" = "objects"."resource_version" + 1 +RETURNING "resource_version", "created_at_ms", "updated_at_ms" "#, ) .bind(object_type) @@ -485,23 +477,11 @@ ON CONFLICT ("object_type", "workspace", "name") WHERE "name" IS NOT NULL DO UPD .bind(payload) .bind(now_ms) .bind(labels.unwrap_or("{}")) - .execute(&self.pool) + .fetch_one(&self.pool) .await .map_err(|e| map_db_error(&e))?; - // Fetch the result to get the resource_version - let record = self - .get_by_name(object_type, workspace, name) - .await? - .ok_or_else(|| { - PersistenceError::Database("object disappeared after upsert".to_string()) - })?; - - Ok(WriteResult { - resource_version: record.resource_version, - created_at_ms: record.created_at_ms, - updated_at_ms: record.updated_at_ms, - }) + Ok(write_result_from_row(&row)) } } } @@ -1913,6 +1893,17 @@ pub(super) fn sqlite_sidecar_paths(path: &Path) -> [PathBuf; 2] { [with_suffix("-wal"), with_suffix("-shm")] } +/// Read the `RETURNING` columns of a write, so the result describes that write +/// and not whichever one landed last. +fn write_result_from_row(row: &sqlx::sqlite::SqliteRow) -> WriteResult { + let resource_version_i64: i64 = row.try_get("resource_version").unwrap_or(1); + WriteResult { + resource_version: resource_version_i64.max(1).cast_unsigned(), + created_at_ms: row.get("created_at_ms"), + updated_at_ms: row.get("updated_at_ms"), + } +} + fn row_to_object_record(row: sqlx::sqlite::SqliteRow) -> ObjectRecord { let resource_version_i64: i64 = row.try_get("resource_version").unwrap_or(1); ObjectRecord { diff --git a/crates/openshell-server/src/persistence/tests.rs b/crates/openshell-server/src/persistence/tests.rs index a62cbd2b42..67c23f18d0 100644 --- a/crates/openshell-server/src/persistence/tests.rs +++ b/crates/openshell-server/src/persistence/tests.rs @@ -2200,6 +2200,92 @@ async fn cas_concurrent_updates_one_succeeds() { assert_eq!(record.resource_version, 2); } +/// Each write must report the version it produced, even while other writers +/// race on the same row. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn sqlite_concurrent_writes_return_their_own_resource_version() { + use super::WriteCondition; + use std::collections::BTreeSet; + use std::sync::Arc; + + const WRITERS: u64 = 8; + const WRITES: u64 = 25; + + let tmp = tempfile::tempdir().expect("tempdir"); + let url = on_disk_store_url(&tmp.path().join("openshell.db")); + let store = Arc::new(Store::connect(&url).await.expect("connect to sqlite")); + store + .put_if( + "sandbox", + "id-1", + "sandbox-1", + "default", + b"initial", + None, + WriteCondition::MustCreate, + ) + .await + .unwrap(); + + let writers = (0..WRITERS).map(|writer| { + let store = Arc::clone(&store); + tokio::spawn(async move { + let mut versions = Vec::new(); + for write in 0..WRITES { + let payload = format!("{writer}-{write}"); + let unconditional = store + .put_if( + "sandbox", + "id-1", + "sandbox-1", + "default", + payload.as_bytes(), + None, + WriteCondition::Unconditional, + ) + .await + .unwrap(); + versions.push(unconditional.resource_version); + + // Retry the version-checked update until it wins a race. + loop { + let current = store.get("sandbox", "id-1").await.unwrap().unwrap(); + match store + .put_if( + "sandbox", + "id-1", + "sandbox-1", + "default", + payload.as_bytes(), + None, + WriteCondition::MatchResourceVersion(current.resource_version), + ) + .await + { + Ok(result) => { + versions.push(result.resource_version); + break; + } + Err(PersistenceError::Conflict { .. }) => {} + Err(err) => panic!("unexpected error: {err}"), + } + } + } + versions + }) + }); + + let mut seen = BTreeSet::new(); + for handle in futures::future::join_all(writers).await { + for version in handle.unwrap() { + assert!(seen.insert(version), "version {version} returned twice"); + } + } + + let total = WRITERS * WRITES * 2; + assert_eq!(seen, (2..=total + 1).collect::>()); +} + #[tokio::test] async fn cas_update_message_cas_succeeds() { use openshell_core::proto::Sandbox; From 73eddf6ca34af086b7326616e351420699dda2b4 Mon Sep 17 00:00:00 2001 From: Eric Curtin Date: Mon, 12 Oct 2026 13:05:32 +0100 Subject: [PATCH 2/2] refactor(server): dedupe sqlite write helpers Signed-off-by: Eric Curtin --- .../src/persistence/sqlite.rs | 11 ++-- .../openshell-server/src/persistence/tests.rs | 55 ++++++++----------- 2 files changed, 30 insertions(+), 36 deletions(-) diff --git a/crates/openshell-server/src/persistence/sqlite.rs b/crates/openshell-server/src/persistence/sqlite.rs index 2d6012e1d6..19cb420094 100644 --- a/crates/openshell-server/src/persistence/sqlite.rs +++ b/crates/openshell-server/src/persistence/sqlite.rs @@ -1896,16 +1896,19 @@ pub(super) fn sqlite_sidecar_paths(path: &Path) -> [PathBuf; 2] { /// Read the `RETURNING` columns of a write, so the result describes that write /// and not whichever one landed last. fn write_result_from_row(row: &sqlx::sqlite::SqliteRow) -> WriteResult { - let resource_version_i64: i64 = row.try_get("resource_version").unwrap_or(1); WriteResult { - resource_version: resource_version_i64.max(1).cast_unsigned(), + resource_version: row_resource_version(row), created_at_ms: row.get("created_at_ms"), updated_at_ms: row.get("updated_at_ms"), } } +fn row_resource_version(row: &sqlx::sqlite::SqliteRow) -> u64 { + let resource_version: i64 = row.try_get("resource_version").unwrap_or(1); + resource_version.max(1).cast_unsigned() +} + fn row_to_object_record(row: sqlx::sqlite::SqliteRow) -> ObjectRecord { - let resource_version_i64: i64 = row.try_get("resource_version").unwrap_or(1); ObjectRecord { object_type: row.get("object_type"), id: row.get("id"), @@ -1915,7 +1918,7 @@ fn row_to_object_record(row: sqlx::sqlite::SqliteRow) -> ObjectRecord { created_at_ms: row.get("created_at_ms"), updated_at_ms: row.get("updated_at_ms"), labels: row.get("labels"), - resource_version: resource_version_i64.max(1).cast_unsigned(), + resource_version: row_resource_version(&row), } } diff --git a/crates/openshell-server/src/persistence/tests.rs b/crates/openshell-server/src/persistence/tests.rs index 67c23f18d0..7a3b435ce9 100644 --- a/crates/openshell-server/src/persistence/tests.rs +++ b/crates/openshell-server/src/persistence/tests.rs @@ -2200,6 +2200,25 @@ async fn cas_concurrent_updates_one_succeeds() { assert_eq!(record.resource_version, 2); } +/// Write `payload` to the shared `sandbox-1` row under `condition`. +async fn put_sandbox_1( + store: &Store, + payload: &str, + condition: super::WriteCondition, +) -> Result { + store + .put_if( + "sandbox", + "id-1", + "sandbox-1", + "default", + payload.as_bytes(), + None, + condition, + ) + .await +} + /// Each write must report the version it produced, even while other writers /// race on the same row. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] @@ -2214,16 +2233,7 @@ async fn sqlite_concurrent_writes_return_their_own_resource_version() { let tmp = tempfile::tempdir().expect("tempdir"); let url = on_disk_store_url(&tmp.path().join("openshell.db")); let store = Arc::new(Store::connect(&url).await.expect("connect to sqlite")); - store - .put_if( - "sandbox", - "id-1", - "sandbox-1", - "default", - b"initial", - None, - WriteCondition::MustCreate, - ) + put_sandbox_1(&store, "initial", WriteCondition::MustCreate) .await .unwrap(); @@ -2233,16 +2243,7 @@ async fn sqlite_concurrent_writes_return_their_own_resource_version() { let mut versions = Vec::new(); for write in 0..WRITES { let payload = format!("{writer}-{write}"); - let unconditional = store - .put_if( - "sandbox", - "id-1", - "sandbox-1", - "default", - payload.as_bytes(), - None, - WriteCondition::Unconditional, - ) + let unconditional = put_sandbox_1(&store, &payload, WriteCondition::Unconditional) .await .unwrap(); versions.push(unconditional.resource_version); @@ -2250,18 +2251,8 @@ async fn sqlite_concurrent_writes_return_their_own_resource_version() { // Retry the version-checked update until it wins a race. loop { let current = store.get("sandbox", "id-1").await.unwrap().unwrap(); - match store - .put_if( - "sandbox", - "id-1", - "sandbox-1", - "default", - payload.as_bytes(), - None, - WriteCondition::MatchResourceVersion(current.resource_version), - ) - .await - { + let condition = WriteCondition::MatchResourceVersion(current.resource_version); + match put_sandbox_1(&store, &payload, condition).await { Ok(result) => { versions.push(result.resource_version); break;