fix(storage-postgres): apply a TransactWriteItems once when two sends share a ClientRequestToken

crates/storage-postgres/src/data/tx_helpers.rs

@@ -395,17 +395,60 @@ pub(super) async fn write_stream_record_in_tx(
 ///
 /// The token is scoped to `account_id`: the same token value from different
 /// accounts is stored and matched independently. Returns `Ok(())` for new
-/// tokens (inserted), `Err(IdempotentReplay)` for matching replays,
-/// `Err(IdempotentMismatch)` for fingerprint conflicts.
+/// tokens (inserted, or a stored row taken for this run),
+/// `Err(IdempotentReplay)` for matching replays, `Err(IdempotentMismatch)`
+/// for fingerprint conflicts.
+///
+/// Two sends of one request with one token can reach this check at the same
+/// time (#851). The second send's insert waits on the first send's token row
+/// and, once the first send commits, locks the committed row. The select of
+/// the same statement reads with the snapshot taken when the statement
+/// started, before that commit, so it returns no row, or the `cancelled:` row
+/// that the committed row replaced. The locked row is neither: the insert's
+/// `WHERE` is false on it, so it is less than 10 minutes old and is not this
+/// request's cancelled row. In those two cases the statement runs once more.
+/// At READ COMMITTED the second statement takes a new snapshot, which sees
+/// the locked row. This transaction holds the lock until it ends, so the row
+/// cannot change between the two statements and the second one does not wait
+/// on it. A committed row of this request then answers as a replay, the
+/// answer the first send got.
 pub(super) async fn check_idempotency_token_in_tx(
     tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
     account_id: &str,
     token: &str,
     fingerprint: &str,
 ) -> Result<(), StorageError> {
-    // A token older than 10 minutes, or one the same request left cancelled
-    // (see `StoredToken::of`), is taken for this run.
-    let row: Option<(String, bool)> = sqlx::query_as(
+    let mut row = run_token_statement(tx, account_id, token, fingerprint).await?;
+    let stale_read = match &row {
+        None => true,
+        Some((stored, false)) => StoredToken::of(stored, fingerprint) == StoredToken::Rerun,
+        Some((_, true)) => false,
+    };
+    if stale_read {
+        row = run_token_statement(tx, account_id, token, fingerprint).await?;
+    }
+
+    match row {
+        Some((_, true)) | None => Ok(()),
+        Some((stored, false)) => match StoredToken::of(&stored, fingerprint) {
+            StoredToken::Replay => Err(StorageError::IdempotentReplay),
+            StoredToken::Rerun | StoredToken::Mismatch => Err(StorageError::IdempotentMismatch),
+        },
+    }
+}
+
+/// One statement of [`check_idempotency_token_in_tx`]: insert the token row,
+/// or take a stored row that is older than 10 minutes or that this request
+/// left cancelled (see `StoredToken::of`), and return it with `true`.
+/// Otherwise return the stored row with `false`, read with the statement's
+/// snapshot, or no row when that snapshot holds none less than 10 minutes old.
+async fn run_token_statement(
+    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
+    account_id: &str,
+    token: &str,
+    fingerprint: &str,
+) -> Result<Option<(String, bool)>, StorageError> {
+    sqlx::query_as(
         r"WITH ins AS (
             INSERT INTO idempotency_tokens (account_id, token, fingerprint)
             VALUES ($1, $2, $3)
@@ -431,15 +474,7 @@ pub(super) async fn check_idempotency_token_in_tx(
     .bind(cancelled_token_fingerprint(fingerprint))
     .fetch_optional(&mut **tx)
     .await
-    .map_err(db_error)?;
-
-    match row {
-        Some((_, true)) | None => Ok(()),
-        Some((stored, false)) => match StoredToken::of(&stored, fingerprint) {
-            StoredToken::Replay => Err(StorageError::IdempotentReplay),
-            StoredToken::Rerun | StoredToken::Mismatch => Err(StorageError::IdempotentMismatch),
-        },
-    }
+    .map_err(db_error)
 }
 
 /// Keep a token used after its request was cancelled past the request-level

crates/storage-postgres/tests/twi_conflict.rs

@@ -18,10 +18,10 @@
 //! 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::collections::{BTreeMap, HashMap};
 use std::time::Duration;
 
-use extenddb_core::expression::{self, Expr, ExpressionMaps};
+use extenddb_core::expression::{self, Expr, ExpressionMaps, UpdateAction};
 use extenddb_core::types::{
     AttributeDefinition, AttributeValue, BillingMode, CreateTableInput, GsiInput, Item,
     KeySchemaElement, KeyType, LsiInput, Projection, ProjectionType,
@@ -1268,3 +1268,223 @@ async fn a_cancelled_token_is_used_for_ten_minutes_from_its_last_cancellation()
 
     s.cleanup().await;
 }
+
+/// Rounds of the two tests below. Most rounds send the second request while
+/// the first is still running.
+const RACE_ROUNDS: usize = 20;
+
+/// The update `SET n = n + :one` and its maps.
+fn increment() -> (Vec<UpdateAction>, ExpressionMaps) {
+    let tokens = expression::tokenize("SET n = n + :one").expect("tokenize");
+    let actions = expression::parse_update(&tokens).expect("parse");
+    let maps = ExpressionMaps::new(
+        HashMap::new(),
+        HashMap::from([("one".to_owned(), AttributeValue::N("1".to_owned()))]),
+    );
+    (actions, maps)
+}
+
+/// Seed the item `cnt` with `n` = 0.
+async fn seed_counter(s: &Scratch, key_info: &TableKeyInfo) {
+    let cnt = BTreeMap::from([
+        ("pk".to_owned(), AttributeValue::S("cnt".to_owned())),
+        ("n".to_owned(), AttributeValue::N("0".to_owned())),
+    ]);
+    s.engine
+        .put_item(key_info, cnt, false, None, &ExpressionMaps::default(), None)
+        .await
+        .expect("seed cnt");
+}
+
+/// The stored `n` of `cnt`.
+async fn counter(db: &PgPool, table: &str) -> i64 {
+    let n: String = sqlx::query_scalar(&format!(
+        "SELECT item_data->'n'->>'N' FROM {table} WHERE pk = 'cnt'"
+    ))
+    .fetch_one(db)
+    .await
+    .expect("read n");
+    n.parse().expect("n is an integer")
+}
+
+/// Send `ops` with `key` twice at the same moment, released by a barrier, and
+/// return both answers. Bounded, so a send that never finishes fails the test
+/// instead of hanging it.
+async fn send_twice_at_once(
+    engine: &PostgresEngine,
+    ops: &[TransactWriteOp<'_>],
+    key: IdempotencyKey<'_>,
+) -> [Result<(), StorageError>; 2] {
+    let barrier = &tokio::sync::Barrier::new(2);
+    let send = move || async move {
+        barrier.wait().await;
+        engine.transact_write_items(ops, Some(key)).await
+    };
+    let (first, second) = tokio::time::timeout(Duration::from_secs(30), async {
+        tokio::join!(send(), send())
+    })
+    .await
+    .expect("both sends finished");
+    [first, second]
+}
+
+/// One send applied the request and the other answered the replay or
+/// TransactionInProgress.
+fn applied_by_one(answers: &[Result<(), StorageError>; 2]) -> bool {
+    let applied = answers.iter().filter(|a| a.is_ok()).count();
+    let waited = answers
+        .iter()
+        .filter(|a| {
+            matches!(
+                a,
+                Err(StorageError::IdempotentReplay | StorageError::TransactionInProgress)
+            )
+        })
+        .count();
+    applied == 1 && waited == 1
+}
+
+#[tokio::test]
+async fn two_sends_of_one_request_with_one_token_apply_it_once() {
+    // Measured on the service (2026-11-05): two sends of one request with one
+    // fresh token at the same moment applied it once in all 20 rounds. One
+    // send answered OK and the other TransactionInProgressException, or both
+    // answered OK. Here the second send's token insert waits on the first
+    // send's token row, and the second send answers the replay once the first
+    // commits (#851).
+    let test = "two_sends_of_one_request_with_one_token_apply_it_once";
+    if base_conn().is_none() {
+        return skip(test);
+    }
+    let s = scratch().await;
+    let plain = ExpressionMaps::default();
+    let key_info = seeded_table(&s, &plain).await;
+    seed_counter(&s, &key_info).await;
+    let table = data_table(&s.db).await;
+    let (actions, maps) = increment();
+    let cnt = key("cnt");
+    let ops = [TransactWriteOp::Update {
+        key_info: &key_info,
+        key: &cnt,
+        actions: &actions,
+        condition: None,
+        maps: &maps,
+        return_values_on_ccf: ReturnValuesOnConditionCheckFailure::None,
+        stream: None,
+    }];
+
+    let mut failed = Vec::new();
+    for round in 0..RACE_ROUNDS {
+        let token = format!("twi-race-{round}");
+        let key = IdempotencyKey {
+            account_id: ACCOUNT,
+            token: &token,
+            fingerprint: "fp-increment",
+        };
+        let before = counter(&s.db, &table).await;
+        let answers = send_twice_at_once(&s.engine, &ops, key).await;
+        let grew = counter(&s.db, &table).await - before;
+        if grew != 1 || !applied_by_one(&answers) {
+            failed.push(format!(
+                "round {round}: n grew by {grew}, answers {answers:?}"
+            ));
+        }
+    }
+    assert!(
+        failed.is_empty(),
+        "{} of {RACE_ROUNDS} rounds:\n{}",
+        failed.len(),
+        failed.join("\n")
+    );
+
+    s.cleanup().await;
+}
+
+#[tokio::test]
+async fn two_sends_that_run_a_cancelled_request_again_apply_it_once() {
+    // A request cancelled after the request-level key checks keeps its token,
+    // and the same request sent again with it runs again (measured
+    // 2026-11-05). Here the second run commits, and it is sent twice at the
+    // same moment. The second send waits on the first send's token row and
+    // answers the replay once the first commits. It does not answer the
+    // mismatch, which is the answer for a different request.
+    let test = "two_sends_that_run_a_cancelled_request_again_apply_it_once";
+    if base_conn().is_none() {
+        return skip(test);
+    }
+    let s = scratch().await;
+    let plain = ExpressionMaps::default();
+    let key_info = seeded_table(&s, &plain).await;
+    seed_counter(&s, &key_info).await;
+    let table = data_table(&s.db).await;
+    let (actions, maps) = increment();
+    let tokens = expression::tokenize("attribute_exists(go)").expect("tokenize");
+    let gate = expression::parse_condition(&tokens).expect("parse");
+    let cnt = key("cnt");
+    let ops = [TransactWriteOp::Update {
+        key_info: &key_info,
+        key: &cnt,
+        actions: &actions,
+        condition: Some(&gate),
+        maps: &maps,
+        return_values_on_ccf: ReturnValuesOnConditionCheckFailure::None,
+        stream: None,
+    }];
+    let set_go = |sql: String| {
+        let db = s.db.clone();
+        async move {
+            sqlx::query(&sql)
+                .execute(&db)
+                .await
+                .expect("change go on cnt");
+        }
+    };
+
+    let mut failed = Vec::new();
+    for round in 0..RACE_ROUNDS {
+        let token = format!("twi-rerun-race-{round}");
+        let key = IdempotencyKey {
+            account_id: ACCOUNT,
+            token: &token,
+            fingerprint: "fp-gated-increment",
+        };
+        // Without `go` the condition fails, and the token is kept cancelled.
+        set_go(format!(
+            "UPDATE {table} SET item_data = item_data - 'go' WHERE pk = 'cnt'"
+        ))
+        .await;
+        let first = tokio::time::timeout(
+            Duration::from_secs(30),
+            s.engine.transact_write_items(&ops, Some(key)),
+        )
+        .await
+        .expect("the first send finished");
+        assert!(
+            matches!(&first, Err(StorageError::TransactionCanceled(r))
+                if r.len() == 1 && r[0].code == "ConditionalCheckFailed"),
+            "round {round}: {first:?}"
+        );
+        // With `go` the same request commits when it runs again.
+        set_go(format!(
+            "UPDATE {table} SET item_data = item_data || \
+             jsonb_build_object('go', jsonb_build_object('S', 'y')) WHERE pk = 'cnt'"
+        ))
+        .await;
+        let before = counter(&s.db, &table).await;
+        let answers = send_twice_at_once(&s.engine, &ops, key).await;
+        let grew = counter(&s.db, &table).await - before;
+        if grew != 1 || !applied_by_one(&answers) {
+            failed.push(format!(
+                "round {round}: n grew by {grew}, answers {answers:?}"
+            ));
+        }
+    }
+    assert!(
+        failed.is_empty(),
+        "{} of {RACE_ROUNDS} rounds:\n{}",
+        failed.len(),
+        failed.join("\n")
+    );
+
+    s.cleanup().await;
+}