diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 5d1ab88ad..1230b1ada 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -297,12 +297,14 @@ jobs: # The control plane for vector indexes is not reachable over the wire while # this backend declares no vector search capability, so its tests drive the - # storage layer directly against this job's PostgreSQL. They build their own - # throwaway databases; the connection string is the server, not a database. + # storage layer directly against this job's PostgreSQL. The write-transaction + # conflict tests force a deadlock, which also needs direct access. They build + # their own throwaway databases; the connection string is the server, not a + # database. - name: Run PostgreSQL storage-level tests env: EXTENDDB_TEST_PG_CONNECTION_STRING: postgresql://postgres:devpass@127.0.0.1:5432 - run: cargo test --release -p extenddb-storage-postgres --test vector_control_plane + run: cargo test --release -p extenddb-storage-postgres --test vector_control_plane --test twi_conflict # The daemonized server logs to syslog; dump it so server-side failures # are diagnosable from the job log. diff --git a/crates/core/src/types/transaction.rs b/crates/core/src/types/transaction.rs index d15e608ef..5010cbb47 100755 --- a/crates/core/src/types/transaction.rs +++ b/crates/core/src/types/transaction.rs @@ -296,6 +296,16 @@ impl CancellationReason { } } + /// Create a reason for an item that another transaction holds. + #[must_use] + pub fn transaction_conflict() -> Self { + Self { + code: "TransactionConflict".to_owned(), + message: Some("Transaction is ongoing for the item".to_owned()), + item: None, + } + } + /// Create a reason for a validation error. #[must_use] pub fn validation_error(msg: impl Into) -> Self { diff --git a/crates/storage-postgres/src/data/mod.rs b/crates/storage-postgres/src/data/mod.rs index d64ceb539..e16ee3b65 100755 --- a/crates/storage-postgres/src/data/mod.rs +++ b/crates/storage-postgres/src/data/mod.rs @@ -92,7 +92,7 @@ macro_rules! bind_sk_fetch_optional { .await } } - .map_err(|e| extenddb_storage::error::StorageError::Internal(e.to_string())) + .map_err($crate::data::index::db_error) }; } @@ -124,7 +124,7 @@ macro_rules! bind_sk_execute { .await } } - .map_err(|e| extenddb_storage::error::StorageError::Internal(e.to_string())) + .map_err($crate::data::index::db_error) }; } diff --git a/crates/storage-postgres/src/data/transactions.rs b/crates/storage-postgres/src/data/transactions.rs index d69f90043..b576a40a4 100644 --- a/crates/storage-postgres/src/data/transactions.rs +++ b/crates/storage-postgres/src/data/transactions.rs @@ -3,26 +3,32 @@ //! Transactional read/write implementations for the `PostgreSQL` backend. +use std::borrow::Cow; use std::collections::HashMap; use extenddb_core::expression::{self, ExpressionMaps}; use extenddb_core::types::{ - AttributeValue, CancellationReason, Item, ReturnValuesOnConditionCheckFailure, + AttributeValue, CancellationReason, Item, ReturnValuesOnConditionCheckFailure, TableKeyInfo, }; use extenddb_core::validation; use extenddb_storage::error::StorageError; +use extenddb_storage::util::pk_to_text; use extenddb_storage::{IdempotencyKey, TransactGetOp, TransactWriteOp}; -use super::index::{IndexMeta, enqueue_async_indexes, fetch_write_path_indexes, sync_indexes}; +use super::index::{ + IndexMeta, db_error, enqueue_async_indexes, fetch_write_path_indexes, sync_indexes, +}; use super::tx_helpers::{ check_idempotency_token_in_tx, delete_item_in_tx, fetch_item_for_update, fetch_item_in_tx, insert_item_if_absent_in_tx, upsert_item_in_tx, write_stream_record_in_tx, }; use crate::PostgresEngine; +use crate::pg_util::is_conflict_abort; -/// Bound on insert retries when a transactional write to a nonexistent item -/// keeps losing the create race to writers that then roll back. Mirrors the -/// same bound on the non-transactional `UpdateItem` path (`update_item.rs`). +/// Bound on insert retries when a transactional write to a nonexistent item, +/// or the reservation of a missing key, keeps losing the create race to a +/// winner that is deleted again before the re-read. Mirrors the same bound on +/// the non-transactional `UpdateItem` path (`update_item.rs`). const MAX_CREATE_RACE_ATTEMPTS: u32 = 5; impl PostgresEngine { @@ -109,19 +115,32 @@ impl PostgresEngine { .await .map_err(|e| StorageError::Internal(e.to_string()))?; + // A conflict abort outside the per-op loop cannot be tied to one item. + // Unreachable at READ COMMITTED, where the ops already hold every row + // lock; a stricter operator isolation can raise 40001 at commit. + let cancel_all = |e: StorageError| conflict_cancels_all(e, ops.len()); + // Check the idempotency token within the transaction so token storage // and data writes commit together. The token is scoped to its account. if let Some(key) = idempotency { check_idempotency_token_in_tx(&mut tx, key.account_id, key.token, key.fingerprint) - .await?; + .await + .map_err(cancel_all)?; } - let mut reasons: Vec = Vec::with_capacity(ops.len()); + let mut reasons: Vec = vec![CancellationReason::none(); ops.len()]; // M-3: Collect old/new items from each op for async GSI enqueue after commit. - let mut op_items: Vec<(Option, Option)> = Vec::with_capacity(ops.len()); + let mut op_items: Vec<(Option, Option)> = vec![(None, None); ops.len()]; let mut any_failed = false; + let mut first_invalid: Option<(usize, String)> = None; + let mut failed_op: Option<(usize, StorageError)> = None; + let mut ran = vec![false; ops.len()]; - for op in ops { + // Run the ops in key order, not request order, so every transaction + // locks its items in one global order and two cannot deadlock on each + // other. Results keep their request positions. + for i in execution_order(ops) { + let op = &ops[i]; let indexes = &table_indexes[transact_op_table_name(op)]; let reason = execute_transact_write_op( &mut tx, @@ -132,27 +151,54 @@ impl PostgresEngine { ) .await; match reason { - Ok(items) => { - op_items.push(items); - reasons.push(CancellationReason::none()); - } + Ok(items) => op_items[i] = items, Err(TxnOpError::Cancel(r)) => { - op_items.push((None, None)); any_failed = true; - reasons.push(r); + reasons[i] = r; } Err(TxnOpError::Validation(msg)) => { - // Up-front input validation (e.g. empty secondary-index key): - // abort the whole transaction with a top-level - // ValidationException, not a per-item cancellation reason. - return Err(StorageError::Validation(msg)); + // Up-front input validation (e.g. empty secondary-index key) + // fails the request with a top-level ValidationException. + // Report the earliest invalid op in request order. + if first_invalid.as_ref().is_none_or(|(j, _)| i < *j) { + first_invalid = Some((i, msg)); + } } Err(TxnOpError::Storage(e)) => { - // Infrastructure error — abort the entire transaction - // without leaking internal details into cancellation reasons. - return Err(StorageError::Internal(e.to_string())); + // Infrastructure error, or PostgreSQL aborted the + // transaction: run no further ops. A request-earlier op + // that sorts later is then not validated. + failed_op = Some((i, e)); + break; } } + ran[i] = true; + // Once every op before the earliest invalid one has run, the + // answer is fixed: stop instead of locking the rest. + if let Some((j, _)) = &first_invalid + && ran[..*j].iter().all(|r| *r) + { + break; + } + } + + // An invalid request fails validation even if a storage error or a + // conflict abort ended the loop. After an early break, a request-earlier + // op that sorts later was never checked, so the op named is the earliest + // invalid one that ran. + if let Some((_, msg)) = first_invalid { + return Err(StorageError::Validation(msg)); + } + if let Some((i, e)) = failed_op { + if is_conflict_abort(&e) { + // PostgreSQL broke a lock conflict by aborting this transaction. + // Cancel it with the contended item named, as the service does. + reasons[i] = CancellationReason::transaction_conflict(); + return Err(StorageError::TransactionCanceled(reasons)); + } + // Infrastructure error: abort without leaking internal details + // into cancellation reasons. + return Err(StorageError::Internal(e.to_string())); } if any_failed { @@ -180,7 +226,8 @@ impl PostgresEngine { old_item.as_ref(), new_item.as_ref(), ) - .await?; + .await + .map_err(cancel_all)?; } } @@ -207,7 +254,8 @@ impl PostgresEngine { new_item.as_ref(), sys_delay, ) - .await?; + .await + .map_err(cancel_all)?; // Vector maintenance for all three write kinds in one place, rather // than in each branch above: this loop already visits exactly the ops @@ -224,15 +272,14 @@ impl PostgresEngine { new_item.as_ref(), sys_delay, ) - .await?; + .await + .map_err(cancel_all)?; if n > 0 || vector_n > 0 { needs_notify = true; } } - tx.commit() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + tx.commit().await.map_err(db_error).map_err(cancel_all)?; if needs_notify && let Some(ref q) = self.gsi_queue { q.notify_workers(); @@ -260,6 +307,43 @@ impl PostgresEngine { } } +/// Map a conflict abort to a cancellation that names every item. Other +/// errors pass through unchanged. +fn conflict_cancels_all(e: StorageError, n_ops: usize) -> StorageError { + if is_conflict_abort(&e) { + StorageError::TransactionCanceled(vec![CancellationReason::transaction_conflict(); n_ops]) + } else { + e + } +} + +/// Request positions of `ops`, sorted by table and primary key. +fn execution_order(ops: &[TransactWriteOp<'_>]) -> Vec { + let mut order: Vec = (0..ops.len()).collect(); + order.sort_by_cached_key(|&i| lock_key(&ops[i])); + order +} + +/// The table and the primary key values of the item an op touches. +fn lock_key<'a>(op: &'a TransactWriteOp<'_>) -> (&'a str, Vec>) { + let (key_info, key) = match op { + TransactWriteOp::Put { key_info, item, .. } => (key_info, *item), + TransactWriteOp::Delete { key_info, key, .. } + | TransactWriteOp::Update { key_info, key, .. } + | TransactWriteOp::ConditionCheck { key_info, key, .. } => (key_info, *key), + }; + let values = key_info + .key_schema + .iter() + .map(|k| { + key.get(&k.attribute_name) + .and_then(|v| pk_to_text(v).ok()) + .unwrap_or_default() + }) + .collect(); + (&key_info.table_id, values) +} + /// Extract the table name from a transactional write operation. fn transact_op_table_name<'a>(op: &'a TransactWriteOp<'_>) -> &'a str { match op { @@ -423,11 +507,9 @@ async fn execute_transact_write_op( // canceled transaction with a TransactionConflict // reason, the DDB-canonical contention shape, never a // 500 (matches the MongoDB backend's exhaustion path). - return Err(TxnOpError::Cancel(CancellationReason { - code: "TransactionConflict".to_owned(), - message: Some("Transaction is ongoing for the item".to_owned()), - item: None, - })); + return Err(TxnOpError::Cancel( + CancellationReason::transaction_conflict(), + )); } } } @@ -460,9 +542,12 @@ async fn execute_transact_write_op( &key_info.attribute_definitions, ) .map_err(|e| TxnOpError::Cancel(CancellationReason::validation_error(e.to_string())))?; - let existing = fetch_item_for_update(tx, key_info, key) + let mut existing = fetch_item_for_update(tx, key_info, key) .await .map_err(TxnOpError::Storage)?; + if existing.is_none() { + existing = reserve_missing_item(tx, key_info, key).await?; + } let empty = Item::new(); eval_condition( *condition, @@ -471,14 +556,8 @@ async fn execute_transact_write_op( *return_values_on_ccf, existing.as_ref(), )?; - // Only delete a row the locking read actually saw (and locked). - // When the read found nothing there is nothing to lock, so a - // concurrent transaction can create and commit the item before our - // DELETE runs; its fresh READ COMMITTED snapshot would then see - // and kill the winner's row with no stream record and orphaned - // index rows. Deleting a nonexistent item is a no-op in the real - // service, so skipping the write is the faithful serialization - // (this delete simply ordered before the concurrent create). + // Deleting a missing item is a no-op. Its key is reserved, so no + // concurrent create can commit before this transaction ends. if existing.is_some() { delete_item_in_tx(tx, key_info, key) .await @@ -612,11 +691,9 @@ async fn execute_transact_write_op( if attempt >= MAX_CREATE_RACE_ATTEMPTS { // Sustained create-then-delete churn. Same // TransactionConflict cancellation as the Put arm. - return Err(TxnOpError::Cancel(CancellationReason { - code: "TransactionConflict".to_owned(), - message: Some("Transaction is ongoing for the item".to_owned()), - item: None, - })); + return Err(TxnOpError::Cancel( + CancellationReason::transaction_conflict(), + )); } } } @@ -648,9 +725,12 @@ async fn execute_transact_write_op( &key_info.attribute_definitions, ) .map_err(|e| TxnOpError::Cancel(CancellationReason::validation_error(e.to_string())))?; - let existing = fetch_item_for_update(tx, key_info, key) + let mut existing = fetch_item_for_update(tx, key_info, key) .await .map_err(TxnOpError::Storage)?; + if existing.is_none() { + existing = reserve_missing_item(tx, key_info, key).await?; + } let empty = Item::new(); let check_against = existing.as_ref().unwrap_or(&empty); eval_condition( @@ -665,6 +745,52 @@ async fn execute_transact_write_op( } } +/// Hold the key of a missing item until commit, for an op that reads the +/// item but does not create it (ConditionCheck, Delete). +/// +/// A missing item has no row for `FOR UPDATE` to lock, so a concurrent +/// transaction could create it while this one still relies on its absence. +/// Inserting a key-only row and deleting it again leaves the key's unique +/// index entry owned by this open transaction: every insert of the key, +/// transactional or not, waits for this transaction to end, and no other +/// transaction ever sees the row. PostgreSQL documents that wait in +/// . +/// Returns the item, locked, when another transaction created it first. +async fn reserve_missing_item( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + key_info: &TableKeyInfo, + key: &Item, +) -> Result, TxnOpError> { + for _ in 0..MAX_CREATE_RACE_ATTEMPTS { + if insert_item_if_absent_in_tx(tx, key_info, key) + .await + .map_err(TxnOpError::Storage)? + { + // A placeholder left behind would commit as a key-only item. + let deleted = delete_item_in_tx(tx, key_info, key) + .await + .map_err(TxnOpError::Storage)?; + if deleted != 1 { + return Err(TxnOpError::Storage(StorageError::Internal(format!( + "deleting the reserved key removed {deleted} rows, expected 1" + )))); + } + return Ok(None); + } + // Lost to a concurrent create, which has committed by now: lock it. + if let Some(item) = fetch_item_for_update(tx, key_info, key) + .await + .map_err(TxnOpError::Storage)? + { + return Ok(Some(item)); + } + } + // Sustained create-then-delete churn, as on the create paths. + Err(TxnOpError::Cancel( + CancellationReason::transaction_conflict(), + )) +} + /// Evaluate a condition expression, returning a `CancellationReason` on failure. /// /// When `return_values_on_ccf` is `AllOld`, the existing item is included in the @@ -693,3 +819,105 @@ fn eval_condition( } Ok(()) } + +#[cfg(test)] +mod tests { + use extenddb_core::types::{KeySchemaElement, KeyType, TableKeyInfo}; + + use super::*; + + fn table(id: &str) -> TableKeyInfo { + TableKeyInfo { + table_id: id.to_owned(), + key_schema: vec![ + KeySchemaElement { + attribute_name: "pk".to_owned(), + key_type: KeyType::Hash, + }, + KeySchemaElement { + attribute_name: "sk".to_owned(), + key_type: KeyType::Range, + }, + ], + ..TableKeyInfo::default() + } + } + + fn key(pk: &str, sk: &str) -> Item { + Item::from([ + ("pk".to_owned(), AttributeValue::S(pk.to_owned())), + ("sk".to_owned(), AttributeValue::S(sk.to_owned())), + ]) + } + + fn delete<'a>( + key_info: &'a TableKeyInfo, + key: &'a Item, + maps: &'a ExpressionMaps, + ) -> TransactWriteOp<'a> { + TransactWriteOp::Delete { + key_info, + key, + condition: None, + maps, + return_values_on_ccf: ReturnValuesOnConditionCheckFailure::None, + stream: None, + } + } + + /// The items each request locks, in the order it locks them. + fn locked(ops: &[TransactWriteOp<'_>]) -> Vec<(String, Vec)> { + execution_order(ops) + .into_iter() + .map(|i| { + let (t, k) = lock_key(&ops[i]); + (t.to_owned(), k.into_iter().map(Cow::into_owned).collect()) + }) + .collect() + } + + #[test] + fn requests_lock_shared_items_in_the_same_order() { + let (t1, t2) = (table("t1"), table("t2")); + let maps = ExpressionMaps::default(); + let (a, b, c) = (key("a", "1"), key("a", "2"), key("b", "1")); + let forward = [ + delete(&t2, &a, &maps), + delete(&t1, &c, &maps), + delete(&t1, &b, &maps), + delete(&t1, &a, &maps), + ]; + let backward = [ + delete(&t1, &a, &maps), + delete(&t1, &b, &maps), + delete(&t1, &c, &maps), + delete(&t2, &a, &maps), + ]; + assert_eq!(locked(&forward), locked(&backward)); + assert_eq!(execution_order(&forward), vec![3, 2, 1, 0]); + assert_eq!(execution_order(&backward), vec![0, 1, 2, 3]); + } + + #[test] + fn a_conflict_abort_outside_the_ops_cancels_every_item() { + let deadlock = StorageError::Internal("SQLSTATE 40P01: deadlock detected".to_owned()); + match conflict_cancels_all(deadlock, 3) { + StorageError::TransactionCanceled(reasons) => { + assert_eq!(reasons.len(), 3); + for r in reasons { + assert_eq!(r.code, "TransactionConflict"); + assert_eq!( + r.message.as_deref(), + Some("Transaction is ongoing for the item") + ); + } + } + other => panic!("expected a cancellation, got {other:?}"), + } + let other = StorageError::Internal("SQLSTATE 23505: duplicate key".to_owned()); + assert!(matches!( + conflict_cancels_all(other, 3), + StorageError::Internal(_) + )); + } +} diff --git a/crates/storage-postgres/src/data/tx_helpers.rs b/crates/storage-postgres/src/data/tx_helpers.rs index ed30a1bf3..6468f159b 100755 --- a/crates/storage-postgres/src/data/tx_helpers.rs +++ b/crates/storage-postgres/src/data/tx_helpers.rs @@ -14,6 +14,7 @@ use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; use extenddb_storage::util::{SortKeyValue, parse_sk, pk_to_text, sk_column, sk_info}; +use super::index::db_error; use super::{data_table_name, json_to_item}; /// Fetch a single item within an existing transaction. @@ -86,7 +87,7 @@ pub(super) async fn fetch_item_for_update( .bind(pk_text.as_ref()) .fetch_optional(&mut **tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(db_error)?; row.map(|(v,)| v) }; @@ -130,7 +131,7 @@ pub(super) async fn upsert_item_in_tx( .bind(&item_json) .execute(&mut **tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(db_error)?; } Ok(()) } @@ -183,18 +184,18 @@ pub(super) async fn insert_item_if_absent_in_tx( .bind(&item_json) .execute(&mut **tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))? + .map_err(db_error)? .rows_affected() }; Ok(rows_affected == 1) } -/// Delete an item by key within a transaction. +/// Delete an item by key within a transaction. Returns the rows deleted. pub(super) async fn delete_item_in_tx( tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, key_info: &TableKeyInfo, key: &Item, -) -> Result<(), StorageError> { +) -> Result { let ddb_table = data_table_name(&key_info.table_id); let pk_name = &key_info.key_schema[0].attribute_name; let pk_value = key @@ -210,7 +211,7 @@ pub(super) async fn delete_item_in_tx( let sk = parse_sk(sk_value, sk_type)?; let sk_col = sk_column(sk_type); let sql = format!("DELETE FROM {ddb_table} WHERE pk = $1 AND {sk_col} = $2"); - match &sk { + let deleted = match &sk { SortKeyValue::S(s) => { sqlx::query(&sql) .bind(pk_text.as_ref()) @@ -233,16 +234,17 @@ pub(super) async fn delete_item_in_tx( .await } } - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(db_error)?; + Ok(deleted.rows_affected()) } else { let sql = format!("DELETE FROM {ddb_table} WHERE pk = $1"); - sqlx::query(&sql) + let deleted = sqlx::query(&sql) .bind(pk_text.as_ref()) .execute(&mut **tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(db_error)?; + Ok(deleted.rows_affected()) } - Ok(()) } /// Write a stream record within an existing transaction. @@ -323,7 +325,7 @@ pub(super) async fn write_stream_record_in_tx( .bind(&key_info.table_id) .fetch_all(&mut **tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(db_error)?; if shards.is_empty() { // No shards — streams may not be fully set up yet. Skip silently. @@ -339,7 +341,7 @@ pub(super) async fn write_stream_record_in_tx( let (seq_val,): (i64,) = sqlx::query_as("SELECT nextval('stream_seq')") .fetch_one(&mut **tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(db_error)?; let seq = format!("{seq_val:021}"); let record = StreamRecord { @@ -380,7 +382,7 @@ pub(super) async fn write_stream_record_in_tx( .bind(&record_json) .execute(&mut **tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(db_error)?; Ok(()) } @@ -421,7 +423,7 @@ pub(super) async fn check_idempotency_token_in_tx( .bind(fingerprint) .fetch_optional(&mut **tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(db_error)?; match row { Some((_, true)) | None => Ok(()), diff --git a/crates/storage-postgres/src/data/vector_index.rs b/crates/storage-postgres/src/data/vector_index.rs index 0bc0a18b5..6259160cb 100644 --- a/crates/storage-postgres/src/data/vector_index.rs +++ b/crates/storage-postgres/src/data/vector_index.rs @@ -188,7 +188,7 @@ pub(crate) async fn maintain_vector_indexes( .bind(table_id) .fetch_optional(&mut **tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(crate::data::index::db_error)?; pending.is_none() } else { false diff --git a/crates/storage-postgres/src/gsi_queue.rs b/crates/storage-postgres/src/gsi_queue.rs index 39e60d443..d9b551fc3 100644 --- a/crates/storage-postgres/src/gsi_queue.rs +++ b/crates/storage-postgres/src/gsi_queue.rs @@ -262,7 +262,7 @@ pub(crate) async fn enqueue_gsi_pending( .bind(delay_interval) .execute(&mut **tx) .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + .map_err(crate::data::index::db_error)?; Ok(()) } diff --git a/crates/storage-postgres/src/pg_util.rs b/crates/storage-postgres/src/pg_util.rs index f87a22590..9c5b8e2db 100755 --- a/crates/storage-postgres/src/pg_util.rs +++ b/crates/storage-postgres/src/pg_util.rs @@ -18,3 +18,90 @@ pub(crate) fn is_fk_violation(e: &sqlx::Error) -> bool { } false } + +/// Check if a mapped error is `PostgreSQL` aborting the transaction to break a +/// lock conflict: `deadlock_detected` (40P01) or `serialization_failure` (40001). +/// Matches the `SQLSTATE` prefix that `data::index::db_error` writes. +pub(crate) fn is_conflict_abort(e: &extenddb_storage::error::StorageError) -> bool { + match e { + extenddb_storage::error::StorageError::Internal(msg) => { + msg.starts_with("SQLSTATE 40P01:") || msg.starts_with("SQLSTATE 40001:") + } + _ => false, + } +} + +#[cfg(test)] +mod tests { + use std::borrow::Cow; + use std::fmt; + + use extenddb_storage::error::StorageError; + + use super::is_conflict_abort; + use crate::data::index::db_error; + + /// A database error with a chosen SQLSTATE, standing in for the server. + #[derive(Debug)] + struct FakeDbError(&'static str, &'static str); + + impl fmt::Display for FakeDbError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.1) + } + } + + impl std::error::Error for FakeDbError {} + + impl sqlx::error::DatabaseError for FakeDbError { + fn message(&self) -> &str { + self.1 + } + fn code(&self) -> Option> { + Some(Cow::Borrowed(self.0)) + } + fn as_error(&self) -> &(dyn std::error::Error + Send + Sync + 'static) { + self + } + fn as_error_mut(&mut self) -> &mut (dyn std::error::Error + Send + Sync + 'static) { + self + } + fn into_error(self: Box) -> Box { + self + } + fn kind(&self) -> sqlx::error::ErrorKind { + sqlx::error::ErrorKind::Other + } + } + + fn mapped(code: &'static str, msg: &'static str) -> StorageError { + db_error(sqlx::Error::Database(Box::new(FakeDbError(code, msg)))) + } + + #[test] + fn deadlock_and_serialization_failure_are_conflict_aborts() { + assert!(is_conflict_abort(&mapped("40P01", "deadlock detected"))); + assert!(is_conflict_abort(&mapped( + "40001", + "could not serialize access due to concurrent update" + ))); + } + + #[test] + fn other_errors_are_not_conflict_aborts() { + // A unique violation, a lock timeout, and the same text without the + // code all stay internal errors. + assert!(!is_conflict_abort(&mapped("23505", "duplicate key value"))); + assert!(!is_conflict_abort(&mapped( + "55P03", + "could not obtain lock" + ))); + assert!(!is_conflict_abort(&StorageError::Internal( + "deadlock detected".to_owned() + ))); + assert!(!is_conflict_abort(&db_error(sqlx::Error::PoolTimedOut))); + assert!(!is_conflict_abort(&StorageError::TransactionConflict( + "SQLSTATE 40P01: deadlock detected".to_owned() + ))); + } +} diff --git a/crates/storage-postgres/tests/twi_conflict.rs b/crates/storage-postgres/tests/twi_conflict.rs new file mode 100644 index 000000000..6752c2e62 --- /dev/null +++ b/crates/storage-postgres/tests/twi_conflict.rs @@ -0,0 +1,859 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 +//! Storage-level tests for lock conflicts inside `TransactWriteItems`. +//! +//! Amazon DynamoDB cancels a write transaction that loses a conflict with a +//! `TransactionConflict` reason on the contended item. These tests force the +//! PostgreSQL side of that: a real deadlock with an outside lock holder, and +//! many transactions that take the same items in opposite request orders. They +//! also pin that a check or delete of a missing item holds its key until +//! commit, so a concurrent create orders around it instead of slipping by. +//! +//! Each test builds its own throwaway database, applies the shipped migrations +//! to it, and drops it when it passes. A failing test leaves its database behind +//! on purpose, named `eddb_twic_*`, so the state that failed can be inspected. +//! +//! Requires `EXTENDDB_TEST_PG_CONNECTION_STRING`, a base URL with no database +//! component (for example `postgresql://postgres@127.0.0.1:5432`), pointing at a +//! server whose role may create and drop databases. Without it every test here +//! reports a skip and passes, the same convention the wire suites use. + +use std::collections::BTreeMap; +use std::time::Duration; + +use extenddb_core::expression::{self, Expr, ExpressionMaps}; +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, BillingMode, CreateTableInput, GsiInput, Item, + KeySchemaElement, KeyType, LsiInput, Projection, ProjectionType, + ReturnValuesOnConditionCheckFailure, ScalarAttributeType, StreamRecord, StreamSpecification, + StreamViewType, TableKeyInfo, +}; +use extenddb_storage::error::StorageError; +use extenddb_storage::{DataEngine, StreamCapture, TableEngine, TransactWriteOp}; +use extenddb_storage_postgres::{PostgresConfig, PostgresEngine}; +use sqlx::postgres::PgPoolOptions; +use sqlx::{PgPool, Postgres, Transaction}; + +const ACCOUNT: &str = "123456789012"; +const REGION: &str = "us-east-1"; +const TABLE: &str = "t_twi_conflict"; + +struct Scratch { + engine: PostgresEngine, + db: PgPool, + admin: PgPool, + db_name: String, +} + +impl Scratch { + async fn cleanup(self) { + let Scratch { + engine, + db, + admin, + db_name, + } = self; + drop(engine); + db.close().await; + sqlx::query(&format!( + "DROP DATABASE IF EXISTS \"{db_name}\" WITH (FORCE)" + )) + .execute(&admin) + .await + .expect("drop the scratch database"); + admin.close().await; + } +} + +fn base_conn() -> Option { + let conn = std::env::var("EXTENDDB_TEST_PG_CONNECTION_STRING").ok()?; + (!conn.trim().is_empty()).then(|| conn.trim_end_matches('/').to_owned()) +} + +fn skip(test: &str) { + eprintln!( + "SKIP {test}: EXTENDDB_TEST_PG_CONNECTION_STRING is not set, so there is no PostgreSQL \ + to build a scratch catalog in." + ); +} + +async fn scratch() -> Scratch { + let base = base_conn().expect("caller checks base_conn() first"); + let db_name = format!("eddb_twic_{}", uuid::Uuid::new_v4().simple())[..24].to_owned(); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/postgres")) + .await + .expect("connect to the postgres maintenance database"); + sqlx::query(&format!("CREATE DATABASE \"{db_name}\"")) + .execute(&admin) + .await + .expect("create the scratch database"); + let url = format!("{base}/{db_name}"); + let db = PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await + .expect("connect to the scratch database"); + for sql in [ + include_str!("../migrations/001_schema.sql"), + include_str!("../migrations/002_vector_indexes.sql"), + include_str!("../data_migrations/001_data_schema.sql"), + include_str!("../data_migrations/002_gsi_pending.sql"), + include_str!("../data_migrations/003_idempotency_account_scope.sql"), + include_str!("../data_migrations/004_vector_index_state.sql"), + ] { + sqlx::raw_sql(sql) + .execute(&db) + .await + .expect("apply a shipped migration"); + } + sqlx::query("UPDATE settings SET value = '0' WHERE key = 'control_plane_delay_seconds'") + .execute(&db) + .await + .expect("pin the control-plane delay to zero"); + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES ($1, $2)") + .bind(ACCOUNT) + .bind(format!("acct-{db_name}")) + .execute(&db) + .await + .expect("seed the account row"); + let engine = PostgresEngine::new( + &PostgresConfig { + connection_string: url, + pool_size: 10, + max_item_size_bytes: 400_000, + }, + REGION, + ) + .await + .expect("open a PostgresEngine on the scratch database"); + Scratch { + engine, + db, + admin, + db_name, + } +} + +/// Create a hash-key table holding items `a` and `b`, and return its key info. +async fn seeded_table(s: &Scratch, maps: &ExpressionMaps) -> TableKeyInfo { + s.engine + .create_table( + ACCOUNT, + CreateTableInput { + table_name: TABLE.to_owned(), + key_schema: vec![KeySchemaElement { + attribute_name: "pk".to_owned(), + key_type: KeyType::Hash, + }], + attribute_definitions: vec![AttributeDefinition { + attribute_name: "pk".to_owned(), + attribute_type: ScalarAttributeType::S, + }], + billing_mode: Some(BillingMode::PayPerRequest), + ..Default::default() + }, + ) + .await + .expect("create the table"); + let key_info = s + .engine + .table_key_info(ACCOUNT, TABLE) + .await + .expect("read the key info"); + for pk in ["a", "b"] { + s.engine + .put_item(&key_info, item(pk, "seed"), false, None, maps, None) + .await + .expect("seed an item"); + } + key_info +} + +fn item(pk: &str, v: &str) -> Item { + BTreeMap::from([ + ("pk".to_owned(), AttributeValue::S(pk.to_owned())), + ("v".to_owned(), AttributeValue::S(v.to_owned())), + ]) +} + +fn put<'a>( + key_info: &'a TableKeyInfo, + item: &'a Item, + maps: &'a ExpressionMaps, +) -> TransactWriteOp<'a> { + TransactWriteOp::Put { + key_info, + item, + condition: None, + maps, + return_values_on_ccf: ReturnValuesOnConditionCheckFailure::None, + stream: None, + } +} + +async fn data_table(db: &PgPool) -> String { + let id: String = + sqlx::query_scalar("SELECT table_id FROM tables WHERE account_id = $1 AND table_name = $2") + .bind(ACCOUNT) + .bind(TABLE) + .fetch_one(db) + .await + .expect("look up the table id"); + format!("\"_ddb_{id}\"") +} + +#[tokio::test] +async fn a_deadlock_cancels_with_transaction_conflict_on_the_waiting_item() { + let test = "a_deadlock_cancels_with_transaction_conflict_on_the_waiting_item"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let maps = ExpressionMaps::default(); + let key_info = seeded_table(&s, &maps).await; + let table = data_table(&s.db).await; + + // An outside transaction holds `b`. The write transaction locks `a`, then + // waits on `b`. The outside transaction then asks for `a`, which closes the + // cycle. The write transaction waited first, so its deadlock check runs + // first and PostgreSQL aborts it. + let mut holder = s.db.begin().await.expect("begin the outside transaction"); + let holder_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *holder) + .await + .expect("read the backend pid"); + sqlx::query(&format!("SELECT 1 FROM {table} WHERE pk = 'b' FOR UPDATE")) + .execute(&mut *holder) + .await + .expect("lock b"); + + let (new_a, new_b) = (item("a", "twi"), item("b", "twi")); + // Request order is b, a: the reason must land on b's request position. + let ops = [put(&key_info, &new_b, &maps), put(&key_info, &new_a, &maps)]; + let twi = s.engine.transact_write_items(&ops, None); + let close_cycle = async { + wait_for_lock_waiters(&s.db, holder_pid, 1).await; + sqlx::query(&format!("SELECT 1 FROM {table} WHERE pk = 'a' FOR UPDATE")) + .execute(&mut *holder) + .await + .expect("the outside transaction gets a once the write transaction aborts"); + }; + // Bounded, so a transaction that never closes the cycle fails the test + // instead of hanging it. + let (result, ()) = tokio::time::timeout(Duration::from_secs(30), async { + tokio::join!(twi, close_cycle) + }) + .await + .expect("the write transaction finished"); + holder.rollback().await.expect("release the outside locks"); + + match result { + Err(StorageError::TransactionCanceled(reasons)) => { + let codes: Vec<&str> = reasons.iter().map(|r| r.code.as_str()).collect(); + assert_eq!(codes, ["TransactionConflict", "None"], "{reasons:?}"); + assert_eq!( + reasons[0].message.as_deref(), + Some("Transaction is ongoing for the item") + ); + assert_eq!(reasons[1].message, None); + } + other => panic!("expected a TransactionConflict cancellation, got {other:?}"), + } + // The canceled transaction applied nothing. + for pk in ["a", "b"] { + let v: String = sqlx::query_scalar(&format!( + "SELECT item_data->'v'->>'S' FROM {table} WHERE pk = $1" + )) + .bind(pk) + .fetch_one(&s.db) + .await + .expect("read the item"); + assert_eq!(v, "seed", "item {pk}"); + } + + s.cleanup().await; +} + +#[tokio::test] +async fn opposite_request_orders_do_not_deadlock() { + let test = "opposite_request_orders_do_not_deadlock"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let maps = ExpressionMaps::default(); + let key_info = seeded_table(&s, &maps).await; + + // Eight writers, half naming the items as (a, b) and half as (b, a). Taken + // in request order these deadlock within a few rounds. + let writers = (0..8).map(|w| { + let (engine, key_info, maps) = (&s.engine, &key_info, &maps); + async move { + for round in 0..50 { + let tag = format!("w{w}-r{round}"); + let (new_a, new_b) = (item("a", &tag), item("b", &tag)); + let ops = if w % 2 == 0 { + [put(key_info, &new_a, maps), put(key_info, &new_b, maps)] + } else { + [put(key_info, &new_b, maps), put(key_info, &new_a, maps)] + }; + engine + .transact_write_items(&ops, None) + .await + .unwrap_or_else(|e| panic!("writer {w} round {round}: {e:?}")); + } + } + }); + futures::future::join_all(writers).await; + + // The last commit is some writer's last round, and it wrote both items. + let table = data_table(&s.db).await; + let mut tags = Vec::new(); + for pk in ["a", "b"] { + let v: String = sqlx::query_scalar(&format!( + "SELECT item_data->'v'->>'S' FROM {table} WHERE pk = $1" + )) + .bind(pk) + .fetch_one(&s.db) + .await + .expect("read the item"); + tags.push(v); + } + assert_eq!(tags[0], tags[1], "a and b come from different transactions"); + assert!( + tags[0].ends_with("-r49"), + "last commit was not a final round: {}", + tags[0] + ); + + s.cleanup().await; +} + +/// Create a hash-key table with GSIs `gi1` on `a1` and `gi2` on `a2`. +async fn gsi_table(s: &Scratch) -> TableKeyInfo { + let s_attr = |name: &str| AttributeDefinition { + attribute_name: name.to_owned(), + attribute_type: ScalarAttributeType::S, + }; + let hash = |name: &str| KeySchemaElement { + attribute_name: name.to_owned(), + key_type: KeyType::Hash, + }; + let gsi = |index: &str, attr: &str| GsiInput { + index_name: index.to_owned(), + key_schema: vec![hash(attr)], + projection: Projection { + projection_type: ProjectionType::All, + non_key_attributes: None, + }, + provisioned_throughput: None, + }; + s.engine + .create_table( + ACCOUNT, + CreateTableInput { + table_name: TABLE.to_owned(), + key_schema: vec![hash("pk")], + attribute_definitions: vec![s_attr("pk"), s_attr("a1"), s_attr("a2")], + billing_mode: Some(BillingMode::PayPerRequest), + global_secondary_indexes: Some(vec![gsi("gi1", "a1"), gsi("gi2", "a2")]), + ..Default::default() + }, + ) + .await + .expect("create the table"); + s.engine + .table_key_info(ACCOUNT, TABLE) + .await + .expect("read the key info") +} + +/// An item whose `attr` index key is the empty string. +fn empty_key(pk: &str, attr: &str) -> Item { + BTreeMap::from([ + ("pk".to_owned(), AttributeValue::S(pk.to_owned())), + (attr.to_owned(), AttributeValue::S(String::new())), + ]) +} + +#[tokio::test] +async fn the_earliest_invalid_op_in_request_order_is_reported() { + // The ops run in key order, but Amazon DynamoDB names the first invalid + // item of the request: [z, a] names z's index, [a, z] names a's. + let test = "the_earliest_invalid_op_in_request_order_is_reported"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let maps = ExpressionMaps::default(); + let key_info = gsi_table(&s).await; + let (z, a) = (empty_key("z", "a2"), empty_key("a", "a1")); + + for (ops, expected) in [ + ( + [put(&key_info, &z, &maps), put(&key_info, &a, &maps)], + "IndexName: gi2, IndexKey: a2", + ), + ( + [put(&key_info, &a, &maps), put(&key_info, &z, &maps)], + "IndexName: gi1, IndexKey: a1", + ), + ] { + match s.engine.transact_write_items(&ops, None).await { + Err(StorageError::Validation(msg)) => { + assert!(msg.ends_with(expected), "{msg}"); + } + other => panic!("expected a ValidationException naming {expected}, got {other:?}"), + } + } + + s.cleanup().await; +} + +#[tokio::test] +async fn an_invalid_first_op_answers_without_waiting_on_later_locks() { + // Once the earliest invalid op is known, the rest must not run: here the + // next op's row is held by an outside transaction. + let test = "an_invalid_first_op_answers_without_waiting_on_later_locks"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let maps = ExpressionMaps::default(); + let key_info = gsi_table(&s).await; + let b = item("b", "seed"); + s.engine + .put_item(&key_info, b.clone(), false, None, &maps, None) + .await + .expect("seed b"); + let table = data_table(&s.db).await; + let mut holder = s.db.begin().await.expect("begin the outside transaction"); + sqlx::query(&format!("SELECT 1 FROM {table} WHERE pk = 'b' FOR UPDATE")) + .execute(&mut *holder) + .await + .expect("lock b"); + + let a = empty_key("a", "a1"); + let ops = [put(&key_info, &a, &maps), put(&key_info, &b, &maps)]; + let result = tokio::time::timeout( + Duration::from_secs(5), + s.engine.transact_write_items(&ops, None), + ) + .await + .expect("the transaction answered without waiting on b"); + holder.rollback().await.expect("release b"); + match result { + Err(StorageError::Validation(msg)) => { + assert!(msg.ends_with("IndexName: gi1, IndexKey: a1"), "{msg}"); + } + other => panic!("expected a ValidationException, got {other:?}"), + } + + s.cleanup().await; +} + +/// The condition `attribute_not_exists(pk)`. +fn not_exists() -> Expr { + let tokens = expression::tokenize("attribute_not_exists(pk)").expect("tokenize"); + expression::parse_condition(&tokens).expect("parse") +} + +fn key(pk: &str) -> Item { + BTreeMap::from([("pk".to_owned(), AttributeValue::S(pk.to_owned()))]) +} + +/// Begin an outside transaction that has inserted `pk` but not committed: a +/// create in flight. Returns it and its backend pid. +async fn create_in_flight( + db: &PgPool, + table: &str, + pk: &str, +) -> (Transaction<'static, Postgres>, i32) { + let mut creator = db.begin().await.expect("begin the outside transaction"); + let pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *creator) + .await + .expect("read the backend pid"); + sqlx::query(&format!( + "INSERT INTO {table} (pk, item_data) VALUES ($1, jsonb_build_object('pk', jsonb_build_object('S', $1::text)))" + )) + .bind(pk) + .execute(&mut *creator) + .await + .expect("insert the item"); + (creator, pid) +} + +/// Wait until `n` backends other than `own_pid` wait on a lock. +async fn wait_for_lock_waiters(db: &PgPool, own_pid: i32, n: i64) { + for _ in 0..500 { + let waiting: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_stat_activity \ + WHERE datname = current_database() AND wait_event_type = 'Lock' AND pid <> $1", + ) + .bind(own_pid) + .fetch_one(db) + .await + .expect("read pg_stat_activity"); + if waiting >= n { + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + panic!("{n} backends never waited on a lock"); +} + +#[tokio::test] +async fn a_check_of_a_missing_item_waits_for_its_concurrent_create() { + // The item does not exist yet, but an outside transaction is creating it. + // The check must wait for that create to commit and then fail, instead of + // passing on an absence that is about to end. + let test = "a_check_of_a_missing_item_waits_for_its_concurrent_create"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let maps = ExpressionMaps::default(); + let key_info = seeded_table(&s, &maps).await; + let table = data_table(&s.db).await; + let (creator, creator_pid) = create_in_flight(&s.db, &table, "c").await; + + let (cond, c) = (not_exists(), key("c")); + let ops = [TransactWriteOp::ConditionCheck { + key_info: &key_info, + key: &c, + condition: &cond, + maps: &maps, + return_values_on_ccf: ReturnValuesOnConditionCheckFailure::None, + }]; + let twi = s.engine.transact_write_items(&ops, None); + let commit = async { + wait_for_lock_waiters(&s.db, creator_pid, 1).await; + creator.commit().await.expect("commit the create"); + }; + let (result, ()) = + tokio::time::timeout(Duration::from_secs(30), async { tokio::join!(twi, commit) }) + .await + .expect("the write transaction finished"); + + match result { + Err(StorageError::TransactionCanceled(reasons)) => { + let codes: Vec<&str> = reasons.iter().map(|r| r.code.as_str()).collect(); + assert_eq!(codes, ["ConditionalCheckFailed"], "{reasons:?}"); + } + other => panic!("expected a ConditionalCheckFailed cancellation, got {other:?}"), + } + + s.cleanup().await; +} + +#[tokio::test] +async fn a_delete_of_a_missing_item_waits_for_its_concurrent_create() { + // Same race for a Delete: it must order after the create and remove the + // new item, not skip it as missing while the create commits around it. + let test = "a_delete_of_a_missing_item_waits_for_its_concurrent_create"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let maps = ExpressionMaps::default(); + let key_info = seeded_table(&s, &maps).await; + let table = data_table(&s.db).await; + let (creator, creator_pid) = create_in_flight(&s.db, &table, "c").await; + + let c = key("c"); + let ops = [TransactWriteOp::Delete { + key_info: &key_info, + key: &c, + condition: None, + maps: &maps, + return_values_on_ccf: ReturnValuesOnConditionCheckFailure::None, + stream: None, + }]; + let twi = s.engine.transact_write_items(&ops, None); + let commit = async { + wait_for_lock_waiters(&s.db, creator_pid, 1).await; + creator.commit().await.expect("commit the create"); + }; + let (result, ()) = + tokio::time::timeout(Duration::from_secs(30), async { tokio::join!(twi, commit) }) + .await + .expect("the write transaction finished"); + result.expect("the delete commits"); + + let left: i64 = sqlx::query_scalar(&format!("SELECT count(*) FROM {table} WHERE pk = 'c'")) + .fetch_one(&s.db) + .await + .expect("count c"); + assert_eq!(left, 0, "the delete ran after the create, so c is gone"); + + s.cleanup().await; +} + +#[tokio::test] +async fn a_plain_put_waits_for_a_check_of_the_missing_item() { + // A transaction checks that `0c` is missing, then waits on `b`, which an + // outside transaction holds. A non-transactional PutItem of `0c` meanwhile + // must wait for the transaction, which still relies on `0c` being absent. + let test = "a_plain_put_waits_for_a_check_of_the_missing_item"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let maps = ExpressionMaps::default(); + let key_info = seeded_table(&s, &maps).await; + let table = data_table(&s.db).await; + + let mut holder = s.db.begin().await.expect("begin the outside transaction"); + let holder_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *holder) + .await + .expect("read the backend pid"); + sqlx::query(&format!("SELECT 1 FROM {table} WHERE pk = 'b' FOR UPDATE")) + .execute(&mut *holder) + .await + .expect("lock b"); + + // `0c` sorts before `b`, so the check runs first and then the Put waits. + let (cond, c, new_b) = (not_exists(), key("0c"), item("b", "twi")); + let ops = [ + put(&key_info, &new_b, &maps), + TransactWriteOp::ConditionCheck { + key_info: &key_info, + key: &c, + condition: &cond, + maps: &maps, + return_values_on_ccf: ReturnValuesOnConditionCheckFailure::None, + }, + ]; + let twi = s.engine.transact_write_items(&ops, None); + let plain_put = async { + wait_for_lock_waiters(&s.db, holder_pid, 1).await; + // Readers neither see the reserved key nor wait for it. + let read = tokio::time::timeout(Duration::from_secs(5), async { + let got = s.engine.get_item(&key_info, &key("0c")).await; + let rows: i64 = + sqlx::query_scalar(&format!("SELECT count(*) FROM {table} WHERE pk = '0c'")) + .fetch_one(&s.db) + .await + .expect("count 0c"); + (got, rows) + }) + .await + .expect("a read of the reserved key does not wait"); + assert_eq!(read.0.expect("GetItem of 0c"), None, "GetItem sees no item"); + assert_eq!(read.1, 0, "no reader sees the placeholder row"); + s.engine + .put_item(&key_info, item("0c", "plain"), false, None, &maps, None) + .await + }; + let release = async { + // Both the transaction and the plain put wait. + wait_for_lock_waiters(&s.db, holder_pid, 2).await; + holder.rollback().await.expect("release b"); + }; + let (twi_result, put_result, ()) = tokio::time::timeout(Duration::from_secs(30), async { + tokio::join!(twi, plain_put, release) + }) + .await + .expect("both writes finished"); + twi_result.expect("the check passes: 0c was missing when it ran"); + put_result.expect("the plain put commits after the transaction"); + + let v: String = sqlx::query_scalar(&format!( + "SELECT item_data->'v'->>'S' FROM {table} WHERE pk = '0c'" + )) + .fetch_one(&s.db) + .await + .expect("read 0c"); + assert_eq!(v, "plain"); + + s.cleanup().await; +} + +const LSI_TABLE: &str = "t_twi_lsi_stream"; + +/// Create a (pk, sk) table with LSI `lsi1` on `lsk` and a stream. +async fn lsi_stream_table(s: &Scratch) -> TableKeyInfo { + // The scratch database holds the catalog and data schemas together, so it + // has the catalog's copy of `stream_shards`. Drop its foreign key to + // `tables`, which the data database does not have. + sqlx::query("ALTER TABLE stream_shards DROP CONSTRAINT stream_shards_table_id_fkey") + .execute(&s.db) + .await + .expect("match the data schema"); + let s_attr = |name: &str| AttributeDefinition { + attribute_name: name.to_owned(), + attribute_type: ScalarAttributeType::S, + }; + let key = |name: &str, key_type: KeyType| KeySchemaElement { + attribute_name: name.to_owned(), + key_type, + }; + s.engine + .create_table( + ACCOUNT, + CreateTableInput { + table_name: LSI_TABLE.to_owned(), + key_schema: vec![key("pk", KeyType::Hash), key("sk", KeyType::Range)], + attribute_definitions: vec![s_attr("pk"), s_attr("sk"), s_attr("lsk")], + billing_mode: Some(BillingMode::PayPerRequest), + local_secondary_indexes: Some(vec![LsiInput { + index_name: "lsi1".to_owned(), + key_schema: vec![key("pk", KeyType::Hash), key("lsk", KeyType::Range)], + projection: Projection { + projection_type: ProjectionType::All, + non_key_attributes: None, + }, + }]), + stream_specification: Some(StreamSpecification { + stream_enabled: true, + stream_view_type: Some(StreamViewType::NewAndOldImages), + }), + ..Default::default() + }, + ) + .await + .expect("create the table"); + s.engine + .table_key_info(ACCOUNT, LSI_TABLE) + .await + .expect("read the key info") +} + +fn range_item(pk: &str, extra: &[(&str, &str)]) -> Item { + let mut item = BTreeMap::from([ + ("pk".to_owned(), AttributeValue::S(pk.to_owned())), + ("sk".to_owned(), AttributeValue::S("1".to_owned())), + ]); + for (name, value) in extra { + item.insert((*name).to_owned(), AttributeValue::S((*value).to_owned())); + } + item +} + +#[tokio::test] +async fn a_delete_that_loses_its_reservation_removes_the_winner_everywhere() { + // A write transaction creates `c`, with its LSI row and stream record, and + // then waits on `d`, which an outside transaction holds. A Delete of the + // missing `c` meanwhile waits for that create. Once it commits, the Delete + // must remove the winner: one REMOVE record whose old image is the winner, + // and no LSI row left behind. + let test = "a_delete_that_loses_its_reservation_removes_the_winner_everywhere"; + if base_conn().is_none() { + return skip(test); + } + let s = scratch().await; + let maps = ExpressionMaps::default(); + let key_info = lsi_stream_table(&s).await; + s.engine + .put_item(&key_info, range_item("d", &[]), false, None, &maps, None) + .await + .expect("seed d"); + let table_id: String = + sqlx::query_scalar("SELECT table_id FROM tables WHERE account_id = $1 AND table_name = $2") + .bind(ACCOUNT) + .bind(LSI_TABLE) + .fetch_one(&s.db) + .await + .expect("look up the table id"); + let index_id: String = sqlx::query_scalar( + "SELECT index_id FROM indexes WHERE table_id = $1 AND index_name = 'lsi1'", + ) + .bind(&table_id) + .fetch_one(&s.db) + .await + .expect("look up the index id"); + let (table, lsi) = ( + format!("\"_ddb_{table_id}\""), + format!("\"_ddb_{index_id}\""), + ); + + let mut holder = s.db.begin().await.expect("begin the outside transaction"); + let holder_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *holder) + .await + .expect("read the backend pid"); + sqlx::query(&format!("SELECT 1 FROM {table} WHERE pk = 'd' FOR UPDATE")) + .execute(&mut *holder) + .await + .expect("lock d"); + + let capture = StreamCapture { + view_type: StreamViewType::NewAndOldImages, + user_identity: None, + region: REGION.into(), + }; + let winner = range_item("c", &[("lsk", "l"), ("v", "winner")]); + let new_d = range_item("d", &[("v", "creator")]); + let stream_put = |item| TransactWriteOp::Put { + key_info: &key_info, + item, + condition: None, + maps: &maps, + return_values_on_ccf: ReturnValuesOnConditionCheckFailure::None, + stream: Some(capture.clone()), + }; + let creator_ops = [stream_put(&winner), stream_put(&new_d)]; + let c = range_item("c", &[]); + let delete_ops = [TransactWriteOp::Delete { + key_info: &key_info, + key: &c, + condition: None, + maps: &maps, + return_values_on_ccf: ReturnValuesOnConditionCheckFailure::None, + stream: Some(capture.clone()), + }]; + + let creator = s.engine.transact_write_items(&creator_ops, None); + let deleter = async { + // The creator has inserted c and waits on d. + wait_for_lock_waiters(&s.db, holder_pid, 1).await; + s.engine.transact_write_items(&delete_ops, None).await + }; + let release = async { + // The Delete now waits on the creator's insert of c. + wait_for_lock_waiters(&s.db, holder_pid, 2).await; + holder.rollback().await.expect("release d"); + }; + let (created, deleted, ()) = tokio::time::timeout(Duration::from_secs(30), async { + tokio::join!(creator, deleter, release) + }) + .await + .expect("both transactions finished"); + created.expect("the create commits"); + deleted.expect("the delete commits after it"); + + for t in [&table, &lsi] { + let left: i64 = sqlx::query_scalar(&format!("SELECT count(*) FROM {t} WHERE pk = 'c'")) + .fetch_one(&s.db) + .await + .expect("count c"); + assert_eq!(left, 0, "c is gone from {t}"); + } + let records: Vec<(serde_json::Value,)> = sqlx::query_as( + "SELECT record_data FROM stream_records WHERE table_id = $1 ORDER BY sequence_number", + ) + .bind(&table_id) + .fetch_all(&s.db) + .await + .expect("read the stream records"); + let c_events: Vec<(String, Option)> = records + .into_iter() + .map(|(data,)| serde_json::from_value::(data).expect("a stream record")) + .filter(|r| r.dynamodb.keys.get("pk") == Some(&AttributeValue::S("c".to_owned()))) + .map(|r| (format!("{:?}", r.event_name), r.dynamodb.old_image)) + .collect(); + assert_eq!( + c_events, + [ + ("Insert".to_owned(), None), + ("Remove".to_owned(), Some(winner.clone())) + ], + "the create, then the delete of the winner" + ); + + s.cleanup().await; +} diff --git a/crates/storage-sqlite/docs/design-decisions.md b/crates/storage-sqlite/docs/design-decisions.md index b34f1b016..2c121b786 100644 --- a/crates/storage-sqlite/docs/design-decisions.md +++ b/crates/storage-sqlite/docs/design-decisions.md @@ -32,8 +32,9 @@ DynamoDB requires, and the Postgres backend provides: - Atomic stream-record capture in the same transaction as the data write. - Atomic idempotency-token check+store with the writes. -Postgres achieves this with `BEGIN ISOLATION LEVEL SERIALIZABLE`. SQLite has no -equivalent knob; a naive multi-connection pool with deferred `BEGIN` allows +Postgres does not use `SERIALIZABLE` for this. Its write transactions run at `READ COMMITTED` and hold locks until they end: `SELECT ... FOR UPDATE` on every item that exists, and a reserved unique key on every missing item that a ConditionCheck or Delete reads. + +SQLite has no row locks to build the same thing from; a naive multi-connection pool with deferred `BEGIN` allows write-skew and `SQLITE_BUSY` on the read-then-write path. ### Options considered diff --git a/crates/storage/src/lib.rs b/crates/storage/src/lib.rs index 607c1fafc..82593548d 100755 --- a/crates/storage/src/lib.rs +++ b/crates/storage/src/lib.rs @@ -427,8 +427,11 @@ pub trait DataEngine: Send + Sync { /// /// # Errors /// - /// Returns [`StorageError::TransactionCanceled`] if any condition fails. - /// Returns [`StorageError::Internal`] on transaction or query failure. + /// Returns [`StorageError::TransactionCanceled`] if any condition fails, or + /// with a `TransactionConflict` reason when the database aborts the + /// transaction to break a lock conflict (a deadlock or a serialization + /// failure). + /// Returns [`StorageError::Internal`] on any other transaction or query failure. /// Returns [`StorageError::IdempotentReplay`] if the token matches a previous request. /// Returns [`StorageError::IdempotentMismatch`] if the token exists with different ops. fn transact_write_items( diff --git a/docs/design/04-component-storage.md b/docs/design/04-component-storage.md index b85c60232..5587b8ccc 100755 --- a/docs/design/04-component-storage.md +++ b/docs/design/04-component-storage.md @@ -549,9 +549,9 @@ async fn transact_write_items(&self, input: TransactWriteInput) -> Result<...> { for item in &input.items { match item { TransactWriteItem::Put { .. } => { /* INSERT/UPSERT within tx */ } - TransactWriteItem::Delete { .. } => { /* DELETE within tx */ } + TransactWriteItem::Delete { .. } => { /* SELECT FOR UPDATE + DELETE within tx */ } TransactWriteItem::Update { .. } => { /* UPDATE within tx */ } - TransactWriteItem::ConditionCheck { .. } => { /* SELECT + evaluate */ } + TransactWriteItem::ConditionCheck { .. } => { /* SELECT FOR UPDATE + evaluate */ } } } tx.commit().await?; @@ -559,6 +559,8 @@ async fn transact_write_items(&self, input: TransactWriteInput) -> Result<...> { } ``` +The ops run in table and key order. A ConditionCheck or Delete that finds no row inserts the key and deletes it again, so the key's unique index entry stays with the transaction and any concurrent create of the item waits until the transaction commits or rolls back. + ### 5.6 Migrations Migrations are embedded in the binary at compile time via `include_str!` and applied in order by the diff --git a/docs/differences-from-dynamodb.md b/docs/differences-from-dynamodb.md index da6a00f7f..3177a4b4c 100755 --- a/docs/differences-from-dynamodb.md +++ b/docs/differences-from-dynamodb.md @@ -15,6 +15,7 @@ adaptation when switching between ExtendDB and the real service. | Numeric precision on partition/sort keys (MongoDB backend only) | 38 significant digits | 34 significant digits (BSON Decimal128). Values that exceed this precision are rejected at write and query time with a ValidationException rather than silently downcast. PostgreSQL backend supports the full 38 digits. | | Inverted numeric `BETWEEN` on a sort key (MongoDB backend only) | ValidationException ("The BETWEEN operator requires upper bound to be greater than or equal to lower bound") | Same error in all practical cases. The inversion guard compares bounds via `f64`, so a `KeyConditionExpression` `BETWEEN` whose bounds are inverted only beyond f64's ~15–17 significant digits (e.g. `BETWEEN 10000000000000002 AND 10000000000000001`) is not rejected and returns an empty result set instead. Valid ranges are never wrongly rejected. | | Transaction read concern (MongoDB backend only) | No user-configurable equivalent | `snapshot` is the default, fidelity-preserving mode. With `majority` or `local`, `TransactGetItems` is not guaranteed a single point-in-time snapshot and transaction condition reads have weaker isolation. With `local`, a condition can be evaluated against data that is later rolled back after failover. | +| Contended `TransactWriteItems` | A transaction that conflicts with another in-flight transaction on the same item is canceled with a `TransactionConflict` reason ("Transaction is ongoing for the item") | **PostgreSQL** queues the second transaction on the first one's row locks and **SQLite** runs one writer at a time, so both normally commit and no `TransactionConflict` is returned. PostgreSQL locks each transaction's items in table and key order, so two `TransactWriteItems` cannot deadlock on each other. A PutItem, UpdateItem, DeleteItem or BatchWriteItem of an item that a `TransactWriteItems` has locked also waits for it, where the Amazon DynamoDB developer guide lists that request as a conflict. On PostgreSQL this includes an item that the transaction checks or deletes while it is missing: the transaction holds the key until it commits or rolls back, so a request that creates the item waits. Many transactions that check one hot missing item queue less fairly than on an existing item, because every waiter wakes at each commit and only one proceeds. Each such check or delete of a missing item also inserts and deletes its key, so it writes WAL and leaves a dead row and index entry for autovacuum, even though it changes no item. When PostgreSQL still aborts a transaction (a deadlock with another lock holder, or a serialization failure under a stricter `default_transaction_isolation`), the answer is a `TransactionCanceledException` with a `TransactionConflict` reason, not a 500. **MongoDB** retries write conflicts internally and cancels with `TransactionConflict` on every item only after sustained contention. A client that relies on `TransactionConflict` to back off does not see it under ordinary contention. | ## Authentication and Authorization (AWS IAM/STS auth surface used by DynamoDB) diff --git a/docs/manuals/01-architecture-guide.md b/docs/manuals/01-architecture-guide.md index bd40427c9..ddcfc6fb8 100755 --- a/docs/manuals/01-architecture-guide.md +++ b/docs/manuals/01-architecture-guide.md @@ -96,7 +96,7 @@ PostgreSQL implementation of all storage traits using `sqlx`. Features: - Schema migrations managed by version-stamped SQL files - Items stored as JSONB with indexed key columns - GSI/LSI implemented as separate PostgreSQL tables -- Transactions use `SELECT FOR UPDATE` + single-transaction commits +- Transactions use `SELECT FOR UPDATE` on existing items, a reserved unique key on missing items, and single-transaction commits - Stream records stored in a dedicated table with background cleanup - All queries parameterized (no dynamic SQL construction) diff --git a/docs/manuals/12-extending-extenddb-storage.md b/docs/manuals/12-extending-extenddb-storage.md index d8f7a9b09..75f701d54 100755 --- a/docs/manuals/12-extending-extenddb-storage.md +++ b/docs/manuals/12-extending-extenddb-storage.md @@ -93,6 +93,8 @@ Key design decisions: - **Condition expressions** are evaluated inside the storage transaction. The engine parses and compiles expressions; the storage layer receives an AST (`Expr`) and evaluates it against the existing item within the same transaction that performs the write. This is critical for correctness — condition checks and writes must be atomic. - **Stream capture** is passed as `Option<&StreamCapture>`. When present, the stream record must be written in the same transaction as the data write. - **Idempotency tokens** for `TransactWriteItems` must be checked and stored atomically with the writes, and must be scoped per account: a `ClientRequestToken` is unique per account in DynamoDB, not globally, so the store must key on `(account_id, token)`. Keying on the token alone lets the same token value from two accounts collide (one account's transaction wrongly replayed or rejected against another's). +- **Lock order in `transact_write_items`**: a backend that takes blocking row locks must lock a transaction's items in one global order, so two write transactions cannot deadlock on each other. The PostgreSQL backend runs the ops sorted by table ID and key and keeps the results in request order. When the database still aborts a transaction to break a conflict (a deadlock or a serialization failure), return `StorageError::TransactionCanceled`, never an internal error. Put the `TransactionConflict` reason on the item whose statement hit the abort and `None` on the others. When the abort cannot be tied to one item (for example, it comes at commit), mark every item `TransactionConflict`. +- **Missing items in `transact_write_items`**: a ConditionCheck or Delete that finds no item still depends on that item staying absent until commit. A lock on rows alone does not cover this, and two transactions that each check an item the other creates can then both commit (write skew). The PostgreSQL backend inserts the missing key and deletes it again in the same transaction, so any concurrent insert of the key waits for the transaction to end. - **Items** are `BTreeMap`. A new backend must handle the full `AttributeValue` type (S, N, B, SS, NS, BS, L, M, BOOL, NULL). - **Query** must support forward/reverse sort order, exclusive start key pagination, and routing to secondary index storage. - **Parallel scan** uses `segment` and `total_segments` to partition the keyspace. @@ -395,7 +397,7 @@ The PostgreSQL implementation makes backend-specific choices. These are implemen 1. **JSONB item storage** — items are stored as JSONB, enabling PostgreSQL-specific query optimizations. The traits pass `Item` = `BTreeMap` — your backend can use any serialization format. -2. **Transaction isolation** — the PostgreSQL backend uses `BEGIN ISOLATION LEVEL SERIALIZABLE` for transactions. Your backend needs equivalent isolation guarantees. +2. **Transaction isolation**: the PostgreSQL backend runs `transact_write_items` at `READ COMMITTED` and makes it serializable with locks held until commit: `SELECT ... FOR UPDATE` on every item that exists, and a reserved unique key on every item that a ConditionCheck or Delete finds missing. Your backend needs equivalent isolation guarantees. 3. **Sequence generation** — stream sequence numbers use PostgreSQL sequences (`nextval`). Your backend needs a monotonic counter mechanism. diff --git a/tests/test_transact_absent_item_isolation.py b/tests/test_transact_absent_item_isolation.py new file mode 100644 index 000000000..74104c0cb --- /dev/null +++ b/tests/test_transact_absent_item_isolation.py @@ -0,0 +1,298 @@ +# Copyright 2026 ExtendDB contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Concurrent TransactWriteItems that read items that do not exist. + +A ConditionCheck or a Delete on a missing item reads that the item is absent. +Amazon DynamoDB keeps that read valid until the transaction commits: a +concurrent transaction that creates the item conflicts with it. So when each +of two transactions reads an item the other one creates, they never both +commit as though neither saw the other (write skew). + +Each test runs two such transactions at once, many times, on fresh keys. It +checks every outcome against the two serial orders. A failure must be a +TransactionCanceledException with ConditionalCheckFailed, TransactionConflict, +or None reasons. Amazon DynamoDB answers a few contended transactions with +InternalServerError, so a small fixed number of 5xx is tolerated. +""" + +from __future__ import annotations + +import os +import threading +import uuid + +import boto3 +import pytest +from botocore.config import Config +from botocore.exceptions import ClientError + +from conftest import scoped_table + +ROUNDS = 40 +# Same bound as the contended-write tests: Amazon DynamoDB returns a few +# InternalServerError per run of contended transactions. +MAX_SERVER_ERRORS = 5 +NOT_EXISTS = "attribute_not_exists(pk)" +REASON_CODES = {"None", "ConditionalCheckFailed", "TransactionConflict"} + +# TEMPORARY: MongoDB has the same write skew, and the MongoDB write-race fix +# repairs it. That fix lands separately. The marker is strict, so these tests +# fail on MongoDB once the fix is on main, until the marker is removed. Only an +# assertion counts as the expected failure. The MongoDB test runner sets +# EXTENDDB_TEST_MONGODB_CONTAINER. +XFAIL_UNTIL_MONGODB_FIX = pytest.mark.xfail( + bool(os.environ.get("EXTENDDB_TEST_MONGODB_CONTAINER", "").strip()), + reason="MongoDB write skew on missing items, fixed by the MongoDB write-race fix", + raises=AssertionError, + strict=True, +) + + +@pytest.fixture(scope="module") +def raw_client(endpoint_url): + """A client that never retries, so every 5xx and every cancellation is seen.""" + kwargs: dict = { + "service_name": "dynamodb", + "region_name": os.environ.get("AWS_DEFAULT_REGION", "us-east-1"), + "config": Config( + retries={"total_max_attempts": 1, "mode": "standard"}, + max_pool_connections=4, + ), + } + if endpoint_url: + kwargs["endpoint_url"] = endpoint_url + if endpoint_url.startswith("https://"): + kwargs["verify"] = False + return boto3.client(**kwargs) + + +@pytest.fixture(scope="module") +def table(dynamodb_client): + with scoped_table(dynamodb_client) as name: + yield name + + +@pytest.fixture(scope="module") +def range_table(dynamodb_client): + with scoped_table( + dynamodb_client, + attribute_definitions=[ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "N"}, + ], + key_schema=[ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"}, + ], + ) as name: + yield name + + +# The sort key of every item in the range table: a number with a fraction. +SK = {"N": "1.50"} + + +def _key(pk: str, sk: dict | None = None) -> dict: + key = {"pk": {"S": pk}} + if sk is not None: + key["sk"] = sk + return key + + +def _put(table: str, pk: str, sk: dict | None = None) -> dict: + return {"Put": {"TableName": table, "Item": _key(pk, sk)}} + + +def _check_absent(table: str, pk: str, sk: dict | None = None) -> dict: + return { + "ConditionCheck": { + "TableName": table, + "Key": _key(pk, sk), + "ConditionExpression": NOT_EXISTS, + } + } + + +def _delete(table: str, pk: str, condition: str | None = None) -> dict: + op: dict = {"TableName": table, "Key": _key(pk)} + if condition: + op["ConditionExpression"] = condition + return {"Delete": op} + + +def _update_absent(table: str, pk: str) -> dict: + return { + "Update": { + "TableName": table, + "Key": _key(pk), + "UpdateExpression": "SET o = :o", + "ConditionExpression": NOT_EXISTS, + "ExpressionAttributeValues": {":o": {"S": pk}}, + } + } + + +def _attempt(client, items: list[dict]) -> dict | None: + """Run one transaction. Returns None when it commits, else the failure.""" + try: + client.transact_write_items(TransactItems=items) + return None + except ClientError as e: + r = e.response + return { + "status": r["ResponseMetadata"]["HTTPStatusCode"], + "code": r["Error"]["Code"], + "reasons": r.get("CancellationReasons"), + } + + +def _race(client, t1: list[dict], t2: list[dict]) -> list[dict | None]: + """Start both transactions together and return both outcomes.""" + start = threading.Barrier(2, timeout=30) + out: list[dict | None] = [None, None] + errors: list[BaseException] = [] + + def run(i: int, items: list[dict]): + try: + start.wait() + out[i] = _attempt(client, items) + except BaseException as e: # noqa: BLE001 - surfaced below + errors.append(e) + + threads = [threading.Thread(target=run, args=(i, t)) for i, t in enumerate((t1, t2))] + for t in threads: + t.start() + for t in threads: + t.join() + if errors: + raise errors[0] + return out + + +def _exists(client, table: str, pk: str, sk: dict | None = None) -> bool: + return "Item" in client.get_item(TableName=table, Key=_key(pk, sk), ConsistentRead=True) + + +def _run_rounds(dynamodb_client, raw_client, table, build, sk=None) -> list[tuple]: + """Race `build(x, y)` on fresh keys each round. + + Returns one record per round: both outcomes and whether x and y exist after. + Every failure must be a cancellation, apart from a few 5xx. + """ + rounds = [] + server_errors = [] + for _ in range(ROUNDS): + x, y = f"x-{uuid.uuid4().hex}", f"y-{uuid.uuid4().hex}" + t1, t2 = build(x, y) + r1, r2 = _race(raw_client, t1, t2) + for f in (r1, r2): + if f is None: + continue + if f["status"] >= 500: + assert (f["status"], f["code"]) == (500, "InternalServerError"), f + server_errors.append(f) + continue + assert (f["status"], f["code"]) == (400, "TransactionCanceledException"), f + codes = [r["Code"] for r in f["reasons"]] + assert len(codes) == 2 and set(codes) <= REASON_CODES, f + rounds.append( + (r1, r2, _exists(dynamodb_client, table, x, sk), _exists(dynamodb_client, table, y, sk)) + ) + assert len(server_errors) <= MAX_SERVER_ERRORS, server_errors + return rounds + + +def _both_committed(rounds) -> list[tuple]: + return [r for r in rounds if r[0] is None and r[1] is None] + + +@XFAIL_UNTIL_MONGODB_FIX +def test_condition_checks_on_missing_items_do_not_write_skew( + dynamodb_client, raw_client, table +): + """[check y absent, put x] vs [check x absent, put y]: one at most commits.""" + rounds = _run_rounds( + dynamodb_client, + raw_client, + table, + lambda x, y: ( + [_check_absent(table, y), _put(table, x)], + [_check_absent(table, x), _put(table, y)], + ), + ) + skewed = _both_committed(rounds) + assert not skewed, f"{len(skewed)} of {ROUNDS} rounds committed both: {skewed[0]}" + + +@XFAIL_UNTIL_MONGODB_FIX +def test_condition_checks_on_missing_items_of_a_range_table_do_not_write_skew( + dynamodb_client, raw_client, range_table +): + """The same race on a hash and range table, keyed by a number sort key.""" + rounds = _run_rounds( + dynamodb_client, + raw_client, + range_table, + lambda x, y: ( + [_check_absent(range_table, y, SK), _put(range_table, x, SK)], + [_check_absent(range_table, x, SK), _put(range_table, y, SK)], + ), + sk=SK, + ) + skewed = _both_committed(rounds) + assert not skewed, f"{len(skewed)} of {ROUNDS} rounds committed both: {skewed[0]}" + + +@XFAIL_UNTIL_MONGODB_FIX +def test_conditional_deletes_of_missing_items_do_not_write_skew( + dynamodb_client, raw_client, table +): + """[delete y if absent, put x] vs [delete x if absent, put y]: one at most commits.""" + rounds = _run_rounds( + dynamodb_client, + raw_client, + table, + lambda x, y: ( + [_delete(table, y, NOT_EXISTS), _put(table, x)], + [_delete(table, x, NOT_EXISTS), _put(table, y)], + ), + ) + skewed = _both_committed(rounds) + assert not skewed, f"{len(skewed)} of {ROUNDS} rounds committed both: {skewed[0]}" + + +@XFAIL_UNTIL_MONGODB_FIX +def test_deletes_of_missing_items_serialize(dynamodb_client, raw_client, table): + """[delete y, put x] vs [delete x, put y]: when both commit, the later one + deletes the earlier one's item, so exactly one of x and y remains.""" + rounds = _run_rounds( + dynamodb_client, + raw_client, + table, + lambda x, y: ( + [_delete(table, y), _put(table, x)], + [_delete(table, x), _put(table, y)], + ), + ) + # Both serial orders leave exactly one item: both or neither is skew. + skewed = [r for r in _both_committed(rounds) if r[2] == r[3]] + assert not skewed, f"{len(skewed)} of {ROUNDS} rounds kept both or neither: {skewed[0]}" + + +def test_conditional_updates_of_missing_items_serialize( + dynamodb_client, raw_client, table +): + """[update x if absent, put y] vs [update y if absent, put x]: an update + creates its item, so one at most commits.""" + rounds = _run_rounds( + dynamodb_client, + raw_client, + table, + lambda x, y: ( + [_update_absent(table, x), _put(table, y)], + [_update_absent(table, y), _put(table, x)], + ), + ) + skewed = _both_committed(rounds) + assert not skewed, f"{len(skewed)} of {ROUNDS} rounds committed both: {skewed[0]}" diff --git a/tests/test_transact_write_conflict.py b/tests/test_transact_write_conflict.py new file mode 100644 index 000000000..af8ca7f8c --- /dev/null +++ b/tests/test_transact_write_conflict.py @@ -0,0 +1,219 @@ +# Copyright 2026 ExtendDB contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Concurrent TransactWriteItems that contend for the same items. + +Amazon DynamoDB answers a write transaction that loses a conflict with +TransactionCanceledException (HTTP 400). The contended items carry a +TransactionConflict reason and the other items carry None. Under this load +Amazon DynamoDB also answers a few transactions (about 1 in 1000) with +InternalServerError, so a small fixed number of 5xx is tolerated. Every +committed transaction must apply in full. + +ExtendDB on PostgreSQL queues contending transactions instead of canceling +them, so there the cancellation shape checks have nothing to check and these +tests prove only that contention never surfaces as a 5xx. Only the +opposite-order test can deadlock; the shared-plus-private test checks the +reason positions on Amazon DynamoDB. The mapping of a +database-detected deadlock to TransactionConflict is pinned by the storage +tests in crates/storage-postgres/tests/twi_conflict.rs. +""" + +from __future__ import annotations + +import os +import threading +import time +import uuid +from concurrent.futures import ThreadPoolExecutor + +import boto3 +import pytest +from botocore.config import Config +from botocore.exceptions import ClientError + +from conftest import scoped_table + +WORKERS = 4 +TXNS_PER_WORKER = 25 +# Bounds the run on a server that stalls on each conflict. +TIME_BUDGET_S = 20.0 +# Amazon DynamoDB returned at most 3 InternalServerError in one run of 100 +# contended transactions; they come in bursts. +MAX_SERVER_ERRORS = 5 +CONFLICT_MESSAGE = "Transaction is ongoing for the item" +CANCEL_PREFIX = ( + "Transaction cancelled, please refer cancellation reasons for specific reasons [" +) + + +@pytest.fixture(scope="module") +def raw_client(endpoint_url): + """A client that never retries, so every 5xx and every cancellation is seen.""" + kwargs: dict = { + "service_name": "dynamodb", + "region_name": os.environ.get("AWS_DEFAULT_REGION", "us-east-1"), + "config": Config( + retries={"total_max_attempts": 1, "mode": "standard"}, + max_pool_connections=WORKERS * 2, + ), + } + if endpoint_url: + kwargs["endpoint_url"] = endpoint_url + if endpoint_url.startswith("https://"): + kwargs["verify"] = False + return boto3.client(**kwargs) + + +@pytest.fixture(scope="module") +def table(dynamodb_client): + with scoped_table(dynamodb_client) as name: + yield name + + +def _add_one(table: str, pk: str) -> dict: + return { + "Update": { + "TableName": table, + "Key": {"pk": {"S": pk}}, + "UpdateExpression": "ADD n :one", + "ExpressionAttributeValues": {":one": {"N": "1"}}, + } + } + + +def _run(client, builders) -> tuple[int, int, list[dict]]: + """Run each builder on WORKERS / len(builders) threads at once. + + Returns the attempts, the committed transactions, and one record per failure. + """ + lock = threading.Lock() + attempts = 0 + committed = 0 + failures: list[dict] = [] + start = threading.Barrier(WORKERS) + deadline = time.monotonic() + TIME_BUDGET_S + + def worker(which: int): + nonlocal attempts, committed + start.wait() + for _ in range(TXNS_PER_WORKER): + if time.monotonic() > deadline: + return + with lock: + attempts += 1 + try: + client.transact_write_items(TransactItems=builders[which]()) + with lock: + committed += 1 + except ClientError as e: + r = e.response + with lock: + failures.append( + { + "builder": which, + "status": r["ResponseMetadata"]["HTTPStatusCode"], + "code": r["Error"]["Code"], + "message": r["Error"]["Message"], + "reasons": r.get("CancellationReasons"), + } + ) + + with ThreadPoolExecutor(max_workers=WORKERS) as pool: + futures = [ + pool.submit(worker, i % len(builders)) for i in range(WORKERS) + ] + for f in futures: + f.result() + return attempts, committed, failures + + +def _assert_conflict_shape(failures: list[dict], n_items: int) -> tuple[list[list[str]], int]: + """Assert the failures are conflict cancellations. + + Returns the reason codes of each cancellation and the number of 5xx. + """ + server_errors = [f for f in failures if f["status"] >= 500] + cancellations = [f for f in failures if f["status"] < 500] + for f in server_errors: + assert (f["status"], f["code"]) == (500, "InternalServerError"), f + assert len(server_errors) <= MAX_SERVER_ERRORS, ( + f"{len(server_errors)} 5xx in {len(failures)} failures, first: {server_errors[0]}" + ) + shapes = [] + for f in cancellations: + assert f["status"] == 400, f + assert f["code"] == "TransactionCanceledException", f + reasons = f["reasons"] + assert reasons is not None and len(reasons) == n_items, f + codes = [r["Code"] for r in reasons] + assert set(codes) <= {"None", "TransactionConflict"}, f + assert "TransactionConflict" in codes, f + for r in reasons: + if r["Code"] == "TransactionConflict": + assert r.get("Message") == CONFLICT_MESSAGE, f + else: + assert "Message" not in r, f + assert f["message"] == CANCEL_PREFIX + ", ".join(codes) + "]", f + shapes.append(codes) + return shapes, len(server_errors) + + +def _counter(client, table: str, pk: str) -> int: + item = client.get_item(TableName=table, Key={"pk": {"S": pk}}, ConsistentRead=True) + return int(item["Item"]["n"]["N"]) + + +def test_opposite_order_transactions_cancel_instead_of_failing( + dynamodb_client, raw_client, table +): + """Two items updated in opposite orders by many clients at once.""" + a, b = f"hot-a-{uuid.uuid4().hex[:8]}", f"hot-b-{uuid.uuid4().hex[:8]}" + for pk in (a, b): + dynamodb_client.put_item(TableName=table, Item={"pk": {"S": pk}, "n": {"N": "0"}}) + + attempts, committed, failures = _run( + raw_client, + [ + lambda: [_add_one(table, a), _add_one(table, b)], + lambda: [_add_one(table, b), _add_one(table, a)], + ], + ) + + _, n_5xx = _assert_conflict_shape(failures, 2) + assert committed + len(failures) == attempts + assert attempts == WORKERS * TXNS_PER_WORKER, f"time budget cut the run at {attempts}" + # Every committed transaction applied both updates and no canceled one + # applied any. A 5xx leaves the outcome unknown, so it may count either way. + n_a, n_b = _counter(dynamodb_client, table, a), _counter(dynamodb_client, table, b) + assert n_a == n_b, (n_a, n_b) + assert committed <= n_a <= committed + n_5xx, (n_a, committed, n_5xx) + + +def test_conflict_reason_names_only_the_contended_item( + dynamodb_client, raw_client, table +): + """One shared item plus one private item per transaction, in either position.""" + hot = f"hot-{uuid.uuid4().hex[:8]}" + dynamodb_client.put_item(TableName=table, Item={"pk": {"S": hot}, "n": {"N": "0"}}) + + def private() -> str: + return f"own-{uuid.uuid4().hex}" + + attempts, committed, failures = _run( + raw_client, + [ + lambda: [_add_one(table, hot), _add_one(table, private())], + lambda: [_add_one(table, private()), _add_one(table, hot)], + ], + ) + + shapes, n_5xx = _assert_conflict_shape(failures, 2) + cancellations = [f for f in failures if f["status"] < 500] + for codes, f in zip(shapes, cancellations): + # The private item never conflicts, so only the shared item is named. + expected = ["TransactionConflict", "None"] + assert codes == (expected if f["builder"] == 0 else expected[::-1]), f + assert committed + len(failures) == attempts + assert attempts == WORKERS * TXNS_PER_WORKER, f"time budget cut the run at {attempts}" + assert committed <= _counter(dynamodb_client, table, hot) <= committed + n_5xx