fix(mongodb): retry a write conflict on the stream sequence counter

crates/engine/src/batch_write_item/scripted.rs

@@ -0,0 +1,259 @@
+// Copyright 2026 ExtendDB contributors
+// SPDX-License-Identifier: Apache-2.0
+
+//! A `DataEngine` double for the `write_batch` tests, and the fixtures that run
+//! `write_batch` on it: it answers each write from a script and records the
+//! writes it received.
+
+use std::collections::VecDeque;
+use std::sync::Mutex;
+
+use extenddb_core::expression::{Expr, ExpressionMaps, KeyCondition, UpdateAction};
+use extenddb_core::types::{
+    AttributeDefinition, AttributeValue, Item, KeySchemaElement, KeyType, ScalarAttributeType,
+};
+use extenddb_storage::{
+    BoxedFuture, IdempotencyKey, ItemPairResult, QueryResult, StreamCapture, TransactGetOp,
+    TransactWriteOp,
+};
+use serde_json::{Value, json};
+
+use super::*;
+
+/// Answers the writes it receives from `script`, in order, and succeeds once
+/// the script is used up. Records `table/pk` for every write received.
+pub(super) struct Scripted {
+    script: Mutex<VecDeque<Result<(), StorageError>>>,
+    writes: Mutex<Vec<String>>,
+}
+
+impl Scripted {
+    pub(super) fn new(script: Vec<Result<(), StorageError>>) -> Self {
+        Self {
+            script: Mutex::new(script.into()),
+            writes: Mutex::new(Vec::new()),
+        }
+    }
+
+    /// `table/pk` of every write received, in order.
+    pub(super) fn writes(&self) -> Vec<String> {
+        self.writes.lock().unwrap().clone()
+    }
+
+    fn answer(&self, key_info: &TableKeyInfo, item: &Item) -> Result<Option<Item>, StorageError> {
+        let pk = match item.get("pk") {
+            Some(AttributeValue::S(s)) => s.clone(),
+            other => panic!("the double only takes string partition keys, got {other:?}"),
+        };
+        self.writes
+            .lock()
+            .unwrap()
+            .push(format!("{}/{pk}", key_info.table_name));
+        let next = self.script.lock().unwrap().pop_front();
+        next.unwrap_or(Ok(())).map(|()| None)
+    }
+}
+
+fn unused<T: Send + 'static>() -> BoxedFuture<'static, Result<T, StorageError>> {
+    Box::pin(async { Err(StorageError::Internal("not used by these tests".to_owned())) })
+}
+
+impl DataEngine for Scripted {
+    fn put_item(
+        &self,
+        key_info: &TableKeyInfo,
+        item: Item,
+        _return_old: bool,
+        _condition: Option<&Expr>,
+        _maps: &ExpressionMaps,
+        _stream: Option<&StreamCapture>,
+    ) -> BoxedFuture<'_, Result<Option<Item>, StorageError>> {
+        Box::pin(std::future::ready(self.answer(key_info, &item)))
+    }
+
+    fn delete_item(
+        &self,
+        key_info: &TableKeyInfo,
+        key: &Item,
+        _return_old: bool,
+        _condition: Option<&Expr>,
+        _maps: &ExpressionMaps,
+        _stream: Option<&StreamCapture>,
+    ) -> BoxedFuture<'_, Result<Option<Item>, StorageError>> {
+        Box::pin(std::future::ready(self.answer(key_info, key)))
+    }
+
+    fn get_item(
+        &self,
+        _key_info: &TableKeyInfo,
+        _key: &Item,
+    ) -> BoxedFuture<'_, Result<Option<Item>, StorageError>> {
+        unused()
+    }
+
+    fn update_item(
+        &self,
+        _key_info: &TableKeyInfo,
+        _key: &Item,
+        _actions: &[UpdateAction],
+        _return_old: bool,
+        _return_new: bool,
+        _condition: Option<&Expr>,
+        _maps: &ExpressionMaps,
+        _stream: Option<&StreamCapture>,
+    ) -> BoxedFuture<'_, ItemPairResult> {
+        unused()
+    }
+
+    fn query(
+        &self,
+        _key_info: &TableKeyInfo,
+        _key_condition: &KeyCondition,
+        _maps: &ExpressionMaps,
+        _forward: bool,
+        _limit: Option<i64>,
+        _exclusive_start_key: Option<&Item>,
+        _index_name: Option<&str>,
+    ) -> BoxedFuture<'_, QueryResult> {
+        unused()
+    }
+
+    fn scan(
+        &self,
+        _key_info: &TableKeyInfo,
+        _limit: Option<i64>,
+        _exclusive_start_key: Option<&Item>,
+        _segment: Option<i64>,
+        _total_segments: Option<i64>,
+        _index_name: Option<&str>,
+    ) -> BoxedFuture<'_, QueryResult> {
+        unused()
+    }
+
+    fn scan_key_in_segment(
+        &self,
+        _key_info: &TableKeyInfo,
+        _key: &Item,
+        _segment: i64,
+        _total_segments: i64,
+        _index_name: Option<&str>,
+    ) -> BoxedFuture<'_, Result<bool, StorageError>> {
+        unused()
+    }
+
+    fn transact_get_items(
+        &self,
+        _ops: &[TransactGetOp<'_>],
+    ) -> BoxedFuture<'_, Result<Vec<Option<Item>>, StorageError>> {
+        unused()
+    }
+
+    fn transact_write_items(
+        &self,
+        _ops: &[TransactWriteOp<'_>],
+        _idempotency: Option<IdempotencyKey<'_>>,
+    ) -> BoxedFuture<'_, Result<(), StorageError>> {
+        unused()
+    }
+
+    fn cleanup_expired_idempotency_tokens(
+        &self,
+        _max_age_seconds: i64,
+    ) -> BoxedFuture<'_, Result<u64, StorageError>> {
+        unused()
+    }
+}
+
+// --- fixtures ---------------------------------------------------------------
+
+/// A hash-only table keyed by a string `pk`.
+pub(super) fn table(name: &str) -> TableKeyInfo {
+    let key_schema = vec![KeySchemaElement {
+        attribute_name: "pk".to_owned(),
+        key_type: KeyType::Hash,
+    }];
+    TableKeyInfo {
+        table_name: name.to_owned(),
+        account_id: "123456789012".to_owned(),
+        table_id: format!("id-{name}"),
+        key_schema: key_schema.clone(),
+        base_key_schema: key_schema,
+        attribute_definitions: vec![AttributeDefinition {
+            attribute_name: "pk".to_owned(),
+            attribute_type: ScalarAttributeType::S,
+        }],
+        has_lsi: false,
+        global_secondary_indexes: vec![],
+        local_secondary_indexes: vec![],
+        stream_specification: None,
+        vector_indexes: vec![],
+    }
+}
+
+/// One successful lookup for each of `names`.
+pub(super) fn found(names: &[&str]) -> Vec<Result<TableKeyInfo, StorageError>> {
+    names.iter().map(|n| Ok(table(n))).collect()
+}
+
+pub(super) fn put(pk: &str) -> Value {
+    json!({"PutRequest": {"Item": {"pk": {"S": pk}}}})
+}
+
+/// A put whose partition key is a number, against a table keyed by a string.
+pub(super) fn put_with_wrong_key_type() -> Value {
+    json!({"PutRequest": {"Item": {"pk": {"N": "1"}}}})
+}
+
+/// Run `write_batch` on `request_items`, taking the tables in the order of
+/// `order` with one entry of `lookups` for each.
+pub(super) async fn run(
+    data: &Scripted,
+    request_items: Value,
+    order: &[&str],
+    lookups: Vec<Result<TableKeyInfo, StorageError>>,
+) -> Result<DispatchResult, DynamoDbError> {
+    run_body(
+        data,
+        json!({ "RequestItems": request_items }),
+        order,
+        lookups,
+    )
+    .await
+}
+
+/// Like [`run`], for a request body that carries more than `RequestItems`.
+pub(super) async fn run_body(
+    data: &Scripted,
+    body: Value,
+    order: &[&str],
+    lookups: Vec<Result<TableKeyInfo, StorageError>>,
+) -> Result<DispatchResult, DynamoDbError> {
+    let input: BatchWriteItemInput = serde_json::from_value(body).unwrap();
+    let tables: Vec<(&String, &Vec<WriteRequest>)> = order
+        .iter()
+        .map(|name| input.request_items.get_key_value(*name).unwrap())
+        .collect();
+    let region: Arc<str> = Arc::from("us-east-1");
+    write_batch(
+        data,
+        &LimitsConfig::default(),
+        &region,
+        &input,
+        &tables,
+        lookups,
+    )
+    .await
+}
+
+/// The error of a batch that must have been rejected.
+pub(super) fn rejected(result: Result<DispatchResult, DynamoDbError>) -> DynamoDbError {
+    match result {
+        Err(e) => e,
+        Ok(r) => panic!("the batch must be rejected, got {}", r.body),
+    }
+}
+
+/// The body of a batch that must have been answered with a 200.
+pub(super) fn answered(result: Result<DispatchResult, DynamoDbError>) -> DispatchResult {
+    result.unwrap_or_else(|e| panic!("the batch must answer 200: {e:?}"))
+}

crates/engine/src/batch_write_item/write_batch_deferral_tests.rs

@@ -0,0 +1,134 @@
+// Copyright 2026 ExtendDB contributors
+// SPDX-License-Identifier: Apache-2.0
+
+//! `write_batch` against a `DataEngine` double: what a write that cannot be
+//! applied now leaves behind, and the answer when no request could be applied.
+
+use serde_json::json;
+
+use super::scripted::{Scripted, answered, found, put, rejected, run, run_body};
+use super::*;
+
+fn transient() -> StorageError {
+    StorageError::Transient("lost every attempt".to_owned())
+}
+
+fn conflict() -> StorageError {
+    StorageError::TransactionConflict("held".to_owned())
+}
+
+/// A failed request does not stop the ones after it. The service answered 200
+/// with only the held item's request unprocessed and the other request applied
+/// whichever came first (the neighbouring test has the held item second).
+#[tokio::test]
+async fn a_failed_first_request_does_not_stop_the_rest() {
+    let data = Scripted::new(vec![Err(conflict()), Ok(())]);
+    let items = json!({"t": [put("held"), put("side")]});
+
+    let out = answered(run(&data, items, &["t"], found(&["t"])).await);
+
+    assert_eq!(out.body["UnprocessedItems"], json!({"t": [put("held")]}));
+    assert_eq!(data.writes(), ["t/held", "t/side"]);
+}
+
+/// UnprocessedItems holds exactly the requests that failed, per table, and the
+/// consumed capacity counts the requests that applied.
+#[tokio::test]
+async fn unprocessed_items_hold_only_the_failed_requests_and_capacity_the_applied_ones() {
+    let data = Scripted::new(vec![Ok(()), Err(transient()), Err(conflict()), Ok(())]);
+    let body = json!({
+        "RequestItems": {"t1": [put("a"), put("b")], "t2": [put("c"), put("d")]},
+        "ReturnConsumedCapacity": "TOTAL",
+    });
+
+    let out = answered(run_body(&data, body, &["t1", "t2"], found(&["t1", "t2"])).await);
+
+    assert_eq!(
+        out.body["UnprocessedItems"],
+        json!({"t1": [put("b")], "t2": [put("c")]})
+    );
+    assert_eq!(data.writes(), ["t1/a", "t1/b", "t2/c", "t2/d"]);
+    let mut units: Vec<(String, f64)> = out.body["ConsumedCapacity"]
+        .as_array()
+        .expect("consumed capacity was requested")
+        .iter()
+        .map(|c| {
+            (
+                c["TableName"].as_str().unwrap().to_owned(),
+                c["CapacityUnits"].as_f64().unwrap(),
+            )
+        })
+        .collect();
+    units.sort_by(|a, b| a.0.cmp(&b.0));
+    assert_eq!(units, [("t1".to_owned(), 1.0), ("t2".to_owned(), 1.0)]);
+    assert!((out.metrics.write_capacity_units - 2.0).abs() < f64::EPSILON);
+}
+
+/// A batch of which no request applied is an error, not a 200 whose client
+/// replays it at full speed. Every request is tried first. The service
+/// answered 500 InternalServerError for such a batch; this server answers 503,
+/// which the SDKs retry with backoff.
+#[tokio::test]
+async fn a_batch_where_nothing_applied_answers_503_after_trying_every_request() {
+    for failure in [transient(), conflict()] {
+        let data = Scripted::new(vec![Err(failure.clone()), Err(failure.clone())]);
+        let items = json!({"t": [put("a"), put("b")]});
+
+        let err = rejected(run(&data, items, &["t"], found(&["t"])).await);
+
+        assert!(
+            matches!(err, DynamoDbError::ServiceUnavailable(_)),
+            "{failure:?} gave {err:?}"
+        );
+        assert_eq!(err.status_code(), 503);
+        assert_eq!(data.writes(), ["t/a", "t/b"], "{failure:?}");
+    }
+}
+
+/// A connection error ends the attempt: the requests after it would wait on
+/// the same unreachable backend, so they are left for the client.
+#[tokio::test]
+async fn a_connection_error_leaves_the_rest_of_the_batch_unattempted() {
+    let down = StorageError::Connection("refused".to_owned());
+    let data = Scripted::new(vec![Ok(()), Err(down)]);
+    let items = json!({"t1": [put("a"), put("b"), put("c")], "t2": [put("d")]});
+
+    let out = answered(run(&data, items, &["t1", "t2"], found(&["t1", "t2"])).await);
+
+    assert_eq!(
+        out.body["UnprocessedItems"],
+        json!({"t1": [put("b"), put("c")], "t2": [put("d")]})
+    );
+    assert_eq!(data.writes(), ["t1/a", "t1/b"]);
+}
+
+/// A connection error with nothing applied answers 503 after one attempt.
+#[tokio::test]
+async fn a_connection_error_with_nothing_applied_answers_503_after_one_attempt() {
+    let down = StorageError::Connection("refused".to_owned());
+    let data = Scripted::new(vec![Err(down)]);
+    let items = json!({"t": [put("a"), put("b")]});
+
+    let err = rejected(run(&data, items, &["t"], found(&["t"])).await);
+
+    assert_eq!(err.status_code(), 503, "{err:?}");
+    assert_eq!(data.writes(), ["t/a"]);
+}
+
+/// An error a replay would not clear fails the batch with 500, also after an
+/// earlier request was left unprocessed.
+#[tokio::test]
+async fn an_internal_error_after_a_failed_request_fails_the_batch() {
+    let broken = StorageError::Internal("corrupt row".to_owned());
+    let data = Scripted::new(vec![Err(transient()), Err(broken)]);
+    let items = json!({"t": [put("a"), put("b")]});
+
+    let err = rejected(run(&data, items, &["t"], found(&["t"])).await);
+
+    assert!(
+        matches!(err, DynamoDbError::InternalServerError(_)),
+        "{err:?}"
+    );
+    assert_eq!(err.status_code(), 500);
+    assert_eq!(data.writes(), ["t/a", "t/b"]);
+}