diff --git a/crates/storage-postgres/src/data/update_item.rs b/crates/storage-postgres/src/data/update_item.rs index fceb7463..a81439c3 100755 --- a/crates/storage-postgres/src/data/update_item.rs +++ b/crates/storage-postgres/src/data/update_item.rs @@ -4,7 +4,7 @@ //! `update_item` implementation for the `PostgreSQL` backend. use extenddb_core::expression::{self, Expr, ExpressionMaps, UpdateAction}; -use extenddb_core::types::{Item, KeyType, TableKeyInfo}; +use extenddb_core::types::{Item, TableKeyInfo}; use extenddb_core::validation; use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; @@ -54,153 +54,178 @@ impl PostgresEngine { .load(std::sync::atomic::Ordering::Relaxed) }; - // Fetch existing item - let old_json = if let Some((sk_name, sk_type)) = + // Sort-key binding, computed once: the key is immutable across attempts. + let sk_parts = if let Some((sk_name, sk_type)) = sk_info(&key_info.key_schema, &key_info.attribute_definitions) { let sk_value = key .get(sk_name) .ok_or_else(|| StorageError::Internal("missing sort key".to_owned()))?; - let sk = parse_sk(sk_value, sk_type)?; - let sk_col = sk_column(sk_type); - let select_sql = format!( - "SELECT item_data FROM {ddb_table} WHERE pk = $1 AND {sk_col} = $2 FOR UPDATE" - ); - let row: Option<(serde_json::Value,)> = - bind_sk_fetch_optional!(&select_sql, pk_text.as_ref(), &sk, &mut *tx)?; - row.map(|(v,)| v) + Some((parse_sk(sk_value, sk_type)?, sk_column(sk_type))) } else { - let select_sql = format!("SELECT item_data FROM {ddb_table} WHERE pk = $1 FOR UPDATE"); - let row: Option<(serde_json::Value,)> = sqlx::query_as(&select_sql) - .bind(pk_text.as_ref()) - .fetch_optional(&mut *tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - row.map(|(v,)| v) + None }; - // Build the working item: existing or new with key attributes only (upsert) - let mut item = if let Some(json) = old_json.clone() { - json_to_item(json)? - } else { - key.clone() - }; + // Read-modify-write, retried when a concurrent writer creates the row first. + // + // The locking read below serialises writers to an item that EXISTS. It cannot + // lock a row that does not exist yet, so two writers can both decide to + // insert. The loser's `ON CONFLICT DO NOTHING` affects no rows, and at that + // moment the winner may still be uncommitted, so a plain re-read can see + // nothing at all. Re-reading `FOR UPDATE` is what makes this terminate: it + // blocks until the winner commits, returning its row, or aborts, returning + // none so the insert can be retried. + // + // Returning ConditionalCheckFailedException here instead was wrong. DynamoDB + // serialises writes to one item, and an UpdateItem carrying no condition is + // an upsert, so every writer must succeed and each update expression applies + // on top of whatever the previous writer committed. Measured against the + // service: four writers setting four different attributes on one new key + // yield an item holding all four, so the loser's expression must be re-applied + // to the winner's item rather than overwriting it. + // + // A supplied condition is re-evaluated against the winner for the same + // reason: after losing the race `attribute_exists` is now true, and failing + // it unconditionally was wrong in the opposite direction. + const MAX_CREATE_RACE_ATTEMPTS: u32 = 5; + let mut attempt: u32 = 0; + let (old_item, new_item, item, pre_mutation_item) = loop { + attempt += 1; - // Save pre-mutation item for index sync and stream capture. - let pre_mutation_item = if (!indexes.is_empty() || stream.is_some()) && old_json.is_some() { - Some(item.clone()) - } else { - None - }; + // Current row, locked when present. First pass: the initial read. Later + // passes: the post-conflict re-read that waits on the race winner. + let old_json: Option = match &sk_parts { + Some((sk, sk_col)) => { + let select_sql = format!( + "SELECT item_data FROM {ddb_table} WHERE pk = $1 AND {sk_col} = $2 FOR UPDATE" + ); + let row: Option<(serde_json::Value,)> = + bind_sk_fetch_optional!(&select_sql, pk_text.as_ref(), sk, &mut *tx)?; + row.map(|(v,)| v) + } + None => { + let select_sql = + format!("SELECT item_data FROM {ddb_table} WHERE pk = $1 FOR UPDATE"); + let row: Option<(serde_json::Value,)> = sqlx::query_as(&select_sql) + .bind(pk_text.as_ref()) + .fetch_optional(&mut *tx) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + row.map(|(v,)| v) + } + }; - let old_item = if return_old && old_json.is_some() { - Some(item.clone()) - } else { - None - }; + // Build the working item: existing or new with key attributes only (upsert) + let mut item = if let Some(json) = old_json.clone() { + json_to_item(json)? + } else { + key.clone() + }; - // Evaluate condition against the existing item (empty if non-existent). - // DynamoDB treats a non-existent item as having no attributes at all. - let condition_item = if old_json.is_some() { - &item - } else { - &std::collections::BTreeMap::new() - }; - match check_condition(condition, condition_item, maps) { - Ok(()) => {} - Err(StorageError::ConditionFailed(_)) => { - if old_json.is_some() { - return Err(StorageError::ConditionFailed(Some(item))); + // Save pre-mutation item for index sync and stream capture. + let pre_mutation_item = + if (!indexes.is_empty() || stream.is_some()) && old_json.is_some() { + Some(item.clone()) + } else { + None + }; + + let old_item = if return_old && old_json.is_some() { + Some(item.clone()) + } else { + None + }; + + // Evaluate condition against the existing item (empty if non-existent). + // DynamoDB treats a non-existent item as having no attributes at all. + let empty = std::collections::BTreeMap::new(); + let condition_item = if old_json.is_some() { &item } else { &empty }; + match check_condition(condition, condition_item, maps) { + Ok(()) => {} + Err(StorageError::ConditionFailed(_)) => { + if old_json.is_some() { + return Err(StorageError::ConditionFailed(Some(item))); + } + return Err(StorageError::ConditionFailed(None)); } - return Err(StorageError::ConditionFailed(None)); + Err(e) => return Err(e), } - Err(e) => return Err(e), - } - // Apply update actions - expression::apply_update(actions, &mut item, maps) - .map_err(|e| StorageError::Validation(e.to_string()))?; + // Apply update actions + expression::apply_update(actions, &mut item, maps) + .map_err(|e| StorageError::Validation(e.to_string()))?; - // Validate post-update item size (400 KB limit) - validation::validate_item_size(&item, self.max_item_size_bytes) - .map_err(|e| StorageError::Validation(e.to_string()))?; + // Validate post-update item size (400 KB limit) + validation::validate_item_size(&item, self.max_item_size_bytes) + .map_err(|e| StorageError::Validation(e.to_string()))?; - let new_item = if return_new { Some(item.clone()) } else { None }; + let new_item = if return_new { Some(item.clone()) } else { None }; - // Write the updated item back - let item_json = - serde_json::to_value(&item).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = + serde_json::to_value(&item).map_err(|e| StorageError::Internal(e.to_string()))?; - if let Some((_, sk_type)) = sk_info(&key_info.key_schema, &key_info.attribute_definitions) { - let sk_name_ref = key_info - .key_schema - .iter() - .find(|ks| ks.key_type == KeyType::Range) - .map(|ks| ks.attribute_name.as_str()) - .ok_or_else(|| StorageError::Internal("missing sort key schema".to_owned()))?; - let sk_value = key - .get(sk_name_ref) - .ok_or_else(|| StorageError::Internal("missing sort key".to_owned()))?; - let sk = parse_sk(sk_value, sk_type)?; - let sk_col = sk_column(sk_type); if old_json.is_some() { - // Row existed — update in place. - let update_sql = format!( - "UPDATE {ddb_table} SET item_data = $3 WHERE pk = $1 AND {sk_col} = $2" - ); - bind_sk_execute!(&update_sql, pk_text.as_ref(), &sk, &item_json, &mut *tx)?; - } else { - // Row didn't exist (upsert) — atomic insert, fail if someone beat us. - let insert_sql = format!( - "INSERT INTO {ddb_table} (pk, {sk_col}, item_data) VALUES ($1, $2, $3) \ - ON CONFLICT (pk, {sk_col}) DO NOTHING" - ); - let result = - bind_sk_execute!(&insert_sql, pk_text.as_ref(), &sk, &item_json, &mut *tx)?; - if result.rows_affected() == 0 { - // Another transaction inserted between our SELECT and INSERT. - // Fetch the winner to return with ConditionFailed. - let winner_sql = format!( - "SELECT item_data FROM {ddb_table} WHERE pk = $1 AND {sk_col} = $2" + // Row existed and is locked by the read above, so update in place. + match &sk_parts { + Some((sk, sk_col)) => { + let update_sql = format!( + "UPDATE {ddb_table} SET item_data = $3 WHERE pk = $1 AND {sk_col} = $2" + ); + bind_sk_execute!(&update_sql, pk_text.as_ref(), sk, &item_json, &mut *tx)?; + } + None => { + let update_sql = + format!("UPDATE {ddb_table} SET item_data = $2 WHERE pk = $1"); + sqlx::query(&update_sql) + .bind(pk_text.as_ref()) + .bind(&item_json) + .execute(&mut *tx) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + } + } + break (old_item, new_item, item, pre_mutation_item); + } + + // Row absent, so insert atomically; someone may beat us to it. + let inserted = match &sk_parts { + Some((sk, sk_col)) => { + let insert_sql = format!( + "INSERT INTO {ddb_table} (pk, {sk_col}, item_data) VALUES ($1, $2, $3) \ + ON CONFLICT (pk, {sk_col}) DO NOTHING" + ); + bind_sk_execute!(&insert_sql, pk_text.as_ref(), sk, &item_json, &mut *tx)? + .rows_affected() + == 1 + } + None => { + let insert_sql = format!( + "INSERT INTO {ddb_table} (pk, item_data) VALUES ($1, $2) \ + ON CONFLICT (pk) DO NOTHING" ); - let winner: Option<(serde_json::Value,)> = - bind_sk_fetch_optional!(&winner_sql, pk_text.as_ref(), &sk, &mut *tx)?; - let winner_item = winner.map(|(v,)| json_to_item(v)).transpose()?; - return Err(StorageError::ConditionFailed(winner_item)); + sqlx::query(&insert_sql) + .bind(pk_text.as_ref()) + .bind(&item_json) + .execute(&mut *tx) + .await + .map_err(|e| StorageError::Internal(e.to_string()))? + .rows_affected() + == 1 } + }; + if inserted { + break (old_item, new_item, item, pre_mutation_item); } - } else if old_json.is_some() { - let update_sql = format!("UPDATE {ddb_table} SET item_data = $2 WHERE pk = $1"); - sqlx::query(&update_sql) - .bind(pk_text.as_ref()) - .bind(&item_json) - .execute(&mut *tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - } else { - let insert_sql = format!( - "INSERT INTO {ddb_table} (pk, item_data) VALUES ($1, $2) \ - ON CONFLICT (pk) DO NOTHING" - ); - let result = sqlx::query(&insert_sql) - .bind(pk_text.as_ref()) - .bind(&item_json) - .execute(&mut *tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - if result.rows_affected() == 0 { - // Another transaction inserted between our SELECT and INSERT. - // Fetch the winner to return with ConditionFailed. - let winner_sql = format!("SELECT item_data FROM {ddb_table} WHERE pk = $1"); - let winner: Option<(serde_json::Value,)> = sqlx::query_as(&winner_sql) - .bind(pk_text.as_ref()) - .fetch_optional(&mut *tx) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - let winner_item = winner.map(|(v,)| json_to_item(v)).transpose()?; - return Err(StorageError::ConditionFailed(winner_item)); + + // Lost the create race. Loop round: the locking read at the top waits for + // the winner, then this item's update expression is applied on top of it. + if attempt >= MAX_CREATE_RACE_ATTEMPTS { + return Err(StorageError::Internal(format!( + "UpdateItem could not create the item after {attempt} attempts: a \ + concurrent writer repeatedly created and rolled back the row" + ))); } - } + }; // Sync GSI/LSI update within transaction (D-4). if !indexes.is_empty() { diff --git a/tests/rust/src/concurrency.rs b/tests/rust/src/concurrency.rs index 82c35262..09b3616d 100644 --- a/tests/rust/src/concurrency.rs +++ b/tests/rust/src/concurrency.rs @@ -109,3 +109,244 @@ async fn parallel_writes_then_deletes_no_internal_error() { let _ = c.delete_table().table_name(&table).send().await; } + +/// Concurrent unconditional `UpdateItem` upserts creating the SAME new item must +/// all succeed, and each update expression must apply on top of whatever the +/// previous writer committed. +/// +/// The row lock taken by the read side serialises writers to an item that already +/// exists, but it cannot lock a row that does not exist yet, so two writers could +/// both take the insert path and the loser was answered with +/// `ConditionalCheckFailedException`. An `UpdateItem` carrying no condition is an +/// upsert and must never produce that error. +/// +/// Asserting the merge, not merely the absence of the error, is deliberate and is +/// what makes this test discriminating. Measured against real DynamoDB on +/// 2026-08-10: four writers each setting a different attribute on one brand-new key +/// yield an item holding all four. A fix that resolved the conflict by overwriting +/// with the loser's own computed item would keep every call succeeding while +/// silently dropping the other writers' attributes, and would pass a test that only +/// counted errors. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_unconditional_upserts_to_one_new_item_all_succeed_and_merge() { + let c = client(); + let table = format!("ConcurrentUpsert_{}", ts()); + make_simple_table(c, &table).await; + + // Several trials: the race window is small, so one trial can miss it. + for trial in 0..5 { + let key = format!("same-key-{trial}"); + let attrs = ["a", "b", "c", "d"]; + let mut handles = Vec::new(); + for (i, attr) in attrs.iter().enumerate() { + let c = c.clone(); + let table = table.clone(); + let key = key.clone(); + let attr = (*attr).to_string(); + handles.push(tokio::spawn(async move { + c.update_item() + .table_name(&table) + .key("pk", s(&key)) + .update_expression(format!("SET #{attr} = :v")) + .expression_attribute_names(format!("#{attr}"), &attr) + .expression_attribute_values(":v", s(&format!("writer-{i}"))) + .send() + .await + .map_err(|e| format!("{:?}", e.into_service_error())) + })); + } + + for h in handles { + let r = h.await.expect("task panicked"); + assert!( + r.is_ok(), + "trial {trial}: an unconditional UpdateItem creating a new item failed: {:?}", + r.err() + ); + } + + let got = c + .get_item() + .table_name(&table) + .key("pk", s(&key)) + .consistent_read(true) + .send() + .await + .expect("get_item") + .item + .expect("item must exist after four successful upserts"); + for attr in attrs { + assert!( + got.contains_key(attr), + "trial {trial}: attribute '{attr}' was lost; every writer's expression must be \ + applied on top of the previous winner, got keys {:?}", + got.keys().collect::>() + ); + } + } +} + +/// Conditional creates keep their existing semantics under the create-race retry. +/// +/// Exactly one `attribute_not_exists(pk)` writer may win a race to create the same +/// key, and `attribute_exists(pk)` must never win against a key that never existed. +/// +/// This is a regression guard, not a bug demonstration, and the distinction is worth +/// recording: it passes both before and after the retry was introduced. Re-evaluating +/// the condition against the race winner produces the same answers, because an +/// `attribute_exists` writer fails against its empty base before ever reaching the +/// insert, and an `attribute_not_exists` writer that loses the race still fails when +/// re-evaluated against the winner. The value of this test is proving the retry did +/// not perturb either outcome. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn only_one_racing_conditional_create_wins() { + let c = client(); + let table = format!("ConcurrentCondUpsert_{}", ts()); + make_simple_table(c, &table).await; + + for (expr, expected_winners) in [("attribute_not_exists(pk)", 1), ("attribute_exists(pk)", 0)] { + let key = format!("cond-{}-{}", expected_winners, ts()); + let mut handles = Vec::new(); + for i in 0..4 { + let c = c.clone(); + let table = table.clone(); + let key = key.clone(); + let expr = expr.to_string(); + handles.push(tokio::spawn(async move { + c.update_item() + .table_name(&table) + .key("pk", s(&key)) + .update_expression("SET #v = :v") + .condition_expression(&expr) + .expression_attribute_names("#v", "value") + .expression_attribute_values(":v", s(&format!("w{i}"))) + .send() + .await + .is_ok() + })); + } + let mut winners = 0; + for h in handles { + if h.await.expect("task panicked") { + winners += 1; + } + } + assert_eq!( + winners, expected_winners, + "'{expr}' racing to create one new key: expected {expected_winners} writer(s) to \ + succeed, got {winners}" + ); + } +} + +/// Concurrent writers to the SAME attribute serialize into a linear history. +/// +/// "All succeed" is necessary but not sufficient: six writers each doing +/// `SET v = :their-value` on one new key could all return 200 while updates were +/// applied on stale bases and silently lost. This test proves a total order +/// exists, which is the actual contract (DynamoDB serializes writes per item; +/// which writer ends up last is scheduling-dependent, but the history is linear). +/// +/// The proof uses `ReturnValues: ALL_OLD`. Each writer observes the committed +/// item it replaced, so a clean serialization of N writers over one new key must +/// produce: +/// * exactly one writer that observed no previous item (it created the key), +/// * every other writer observing some OTHER writer's value, each observed at +/// most once (two writers seeing the same old value means one applied to a +/// stale base and the earlier write was lost), +/// * a final value equal to the one written value nobody observed as old. +/// Together these force the observations into one chain: +/// none -> w_a -> w_b -> ... -> final. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_same_attribute_updates_serialize_into_a_linear_history() { + let c = client(); + let table = format!("SameAttrChain_{}", ts()); + make_simple_table(c, &table).await; + + const WRITERS: usize = 6; + // Several trials: the race window is small, so one trial can miss it. + for trial in 0..5 { + let key = format!("chain-key-{trial}"); + let mut handles = Vec::new(); + for i in 0..WRITERS { + let c = c.clone(); + let table = table.clone(); + let key = key.clone(); + handles.push(tokio::spawn(async move { + c.update_item() + .table_name(&table) + .key("pk", s(&key)) + .update_expression("SET v = :v") + .expression_attribute_values(":v", s(&format!("writer-{i}"))) + .return_values(aws_sdk_dynamodb::types::ReturnValue::AllOld) + .send() + .await + .map(|out| { + out.attributes + .and_then(|a| a.get("v").and_then(|v| v.as_s().ok().cloned())) + }) + .map_err(|e| format!("{:?}", e.into_service_error())) + })); + } + + let mut observed_old: Vec> = Vec::new(); + for h in handles { + let r = h.await.expect("task panicked"); + observed_old.push(r.unwrap_or_else(|e| { + panic!("trial {trial}: a same-attribute UpdateItem failed: {e}") + })); + } + + // Exactly one writer created the key (observed no prior item). + let creators = observed_old.iter().filter(|o| o.is_none()).count(); + assert_eq!( + creators, 1, + "trial {trial}: exactly one writer must observe an absent item; {creators} did. \ + More than one means a write was applied as a fresh create over an existing \ + committed value, losing it. Observed: {observed_old:?}" + ); + + // No two writers observed the same predecessor. + let mut seen = std::collections::HashSet::new(); + for old in observed_old.iter().flatten() { + assert!( + seen.insert(old.clone()), + "trial {trial}: two writers observed the same old value '{old}', so one of \ + them applied to a stale base and the write between them was lost. \ + Observed: {observed_old:?}" + ); + } + + // The final value is the one written value that nobody observed as old. + let final_v = c + .get_item() + .table_name(&table) + .key("pk", s(&key)) + .consistent_read(true) + .send() + .await + .expect("get_item") + .item + .and_then(|i| i.get("v").and_then(|v| v.as_s().ok().cloned())) + .expect("item must hold v after successful updates"); + assert!( + final_v.starts_with("writer-"), + "trial {trial}: final value '{final_v}' is not any writer's value" + ); + assert!( + !seen.contains(&final_v), + "trial {trial}: final value '{final_v}' was also observed as someone's old \ + value, so the history has a fork rather than a single chain. \ + Observed: {observed_old:?}" + ); + // Chain accounting: 1 creator + (N-1) observed predecessors + 1 final tail + // must cover all N written values exactly once. + assert_eq!( + seen.len() + 1, + WRITERS, + "trial {trial}: {} distinct predecessors + final tail != {WRITERS} writers; \ + a value vanished from the history. Observed: {observed_old:?}, final: {final_v}", + seen.len() + ); + } +}