From 2b892ef192f3162b1906f5b487733a3baee2568c Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 01/14] fix(mongodb): store the key when UpdateItem creates an item An unconditional UpdateItem on a table without an index runs as one findAndModify with upsert on the MongoDB backend. When the item did not exist, MongoDB built the inserted document from the filter fields and the update operators, so the document got a generated ObjectId as `_id` and an `item_data` without the key attributes. GetItem then returned the item without its key, and every PutItem on that key failed with code 66 (ImmutableField) and HTTP 500, because PutItem writes `_id` as the key text. A REMOVE-only update on a missing item produced a document with no `item_data` at all and answered 500 itself. The upsert now carries a `$setOnInsert` with `_id` as the key text and the key attributes under `item_data`, so a created document is the one PutItem writes, with `_v` 1 as on the transaction path. Two writers creating the same item now collide on `_id` instead of the key index, which the server does not retry for this filter, so the engine retries the upsert within the transient retry budget. The UpdateItem transaction path treats a duplicate key on its own insert as a transient conflict and runs again. Documents written without their key by earlier builds are repaired at startup: each table's data collection is scanned for an ObjectId `_id`, and each such document is rebuilt from its `pk` and `sk_*` fields and `item_data` and replaced in one transaction. The server logs a warning with the counts when it changed anything. New modules: data/key_fields.rs (insert-only key fields, repair of a document), data_engine_update_create.rs (the upsert and its retry), data_engine_key_repair.rs (startup repair), data_engine_live_tests.rs (live MongoDB tests, ignored without EXTENDDB_TEST_MONGODB_URI). The wire test tests/test_update_item_creates_key.py covers creation through SET, ADD on a number and on a set, REMOVE only, a condition, every ReturnValues, the later GetItem, PutItem, UpdateItem, Query, Scan and DeleteItem on the key, and the issue's race of UpdateItem and PutItem writers creating the same keys. Fixes #436 --- crates/storage-mongodb/src/data/key_fields.rs | 328 +++++++++++ crates/storage-mongodb/src/data/mod.rs | 2 + crates/storage-mongodb/src/data_engine.rs | 38 +- .../src/data_engine_key_repair.rs | 222 ++++++++ .../src/data_engine_live_tests.rs | 507 ++++++++++++++++++ .../src/data_engine_update_create.rs | 116 ++++ crates/storage-mongodb/src/lib.rs | 12 + tests/test_update_item_creates_key.py | 363 +++++++++++++ 8 files changed, 1577 insertions(+), 11 deletions(-) create mode 100644 crates/storage-mongodb/src/data/key_fields.rs create mode 100644 crates/storage-mongodb/src/data_engine_key_repair.rs create mode 100644 crates/storage-mongodb/src/data_engine_live_tests.rs create mode 100644 crates/storage-mongodb/src/data_engine_update_create.rs create mode 100644 tests/test_update_item_creates_key.py diff --git a/crates/storage-mongodb/src/data/key_fields.rs b/crates/storage-mongodb/src/data/key_fields.rs new file mode 100644 index 000000000..3eea9fb5b --- /dev/null +++ b/crates/storage-mongodb/src/data/key_fields.rs @@ -0,0 +1,328 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Key attributes of a base-table document, for the two places that build a +//! document from something other than a full item. +//! +//! [`item_to_document`] writes a base-table document as `{ _id, pk, sk_*, +//! item_data, _v }`, with `_id` the key text and the key attributes inside +//! `item_data`. An UpdateItem that creates an item through a native MongoDB +//! upsert starts from the key filter (`pk` and `sk_*`) and the update +//! operators, so it needs the `_id` and the `item_data` key attributes added +//! with `$setOnInsert`: [`insert_only_key_fields`] builds that document. +//! +//! Before that was done, such an upsert stored a document with a generated +//! ObjectId `_id` and no key attributes in `item_data` (issue #436). +//! [`repaired_document`] rebuilds a document of that shape into the shape +//! `item_to_document` writes, reading the key values back from the top-level +//! `pk` and `sk_*` fields. + +use base64::Engine as _; +use bson::{Bson, Document}; +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, Item, KeySchemaElement, KeyType, ScalarAttributeType, +}; +use extenddb_storage::error::StorageError; +use extenddb_storage::util::sk_info; + +use super::{document_to_item, item_to_document, sk_suffix}; + +/// The `$setOnInsert` document for an upsert whose filter is [`super::pk_filter`] +/// on `key`: `_id` as `item_to_document` writes it, and every key attribute +/// under `item_data.`. +/// +/// The filter's own equality fields (`pk`, `sk_*`) are not repeated; MongoDB +/// copies them into the inserted document itself. +pub fn insert_only_key_fields( + key: &Item, + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], +) -> Result { + let key_doc = item_to_document(key, key_schema, attribute_definitions)?; + let mut fields = Document::new(); + let id = key_doc + .get("_id") + .cloned() + .ok_or_else(|| StorageError::Internal("key document has no _id".to_owned()))?; + fields.insert("_id", id); + for ks in key_schema { + let value = key.get(&ks.attribute_name).ok_or_else(|| { + StorageError::Internal(format!("missing key attribute {}", ks.attribute_name)) + })?; + let value_json = + serde_json::to_value(value).map_err(|e| StorageError::Internal(e.to_string()))?; + let value_bson = + bson::to_bson(&value_json).map_err(|e| StorageError::Internal(e.to_string()))?; + fields.insert(format!("item_data.{}", ks.attribute_name), value_bson); + } + Ok(fields) +} + +/// Whether `doc` was stored by the upsert described in the module comment: an +/// `_id` that is not the key text and key attributes missing from `item_data`. +#[must_use] +pub fn is_stored_without_key(doc: &Document) -> bool { + matches!(doc.get("_id"), Some(Bson::ObjectId(_))) +} + +/// Rebuild a document stored without its key into the shape +/// `item_to_document` writes: the same `item_data` with the key attributes +/// added, `_id` as the key text, and `_v` carried over. +/// +/// The key values come from the top-level `pk` text and `sk_s`/`sk_n`/`sk_b` +/// field, typed by `attribute_definitions`. A document whose top-level key +/// fields are missing or of the wrong type is an error; the caller leaves it +/// in place. +pub fn repaired_document( + doc: &Document, + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], +) -> Result { + let mut item = document_to_item(doc)?; + + let pk_name = key_schema + .iter() + .find(|ks| ks.key_type == KeyType::Hash) + .map(|ks| ks.attribute_name.as_str()) + .ok_or_else(|| StorageError::Internal("key schema has no partition key".to_owned()))?; + let pk_type = attribute_type(attribute_definitions, pk_name)?; + let pk_text = doc + .get_str("pk") + .map_err(|_| StorageError::Internal("document has no pk text".to_owned()))?; + item.insert(pk_name.to_owned(), pk_value(pk_type, pk_text)?); + + if let Some((sk_name, sk_type)) = sk_info(key_schema, attribute_definitions) { + let field = format!("sk_{}", sk_suffix(sk_type)); + let value = doc + .get(&field) + .ok_or_else(|| StorageError::Internal(format!("document has no {field}")))?; + item.insert(sk_name.to_owned(), sk_value(sk_type, value, &field)?); + } + + let mut repaired = item_to_document(&item, key_schema, attribute_definitions)?; + if let Some(version) = doc.get("_v") { + repaired.insert("_v", version.clone()); + } + Ok(repaired) +} + +fn attribute_type( + attribute_definitions: &[AttributeDefinition], + name: &str, +) -> Result { + attribute_definitions + .iter() + .find(|d| d.attribute_name == name) + .map(|d| d.attribute_type) + .ok_or_else(|| StorageError::Internal(format!("no attribute definition for {name}"))) +} + +/// The partition key attribute from its `pk` text, inverting +/// `extenddb_storage::util::pk_to_text`. +fn pk_value(pk_type: ScalarAttributeType, text: &str) -> Result { + Ok(match pk_type { + ScalarAttributeType::S => AttributeValue::S(text.to_owned()), + ScalarAttributeType::N => AttributeValue::N(text.to_owned()), + ScalarAttributeType::B => AttributeValue::B( + base64::engine::general_purpose::STANDARD + .decode(text) + .map_err(|e| StorageError::Internal(format!("pk text is not base64: {e}")))?, + ), + }) +} + +/// The sort key attribute from its typed `sk_*` field, inverting +/// `insert_typed_sk`. +fn sk_value( + sk_type: ScalarAttributeType, + value: &Bson, + field: &str, +) -> Result { + match (sk_type, value) { + (ScalarAttributeType::S, Bson::String(s)) => Ok(AttributeValue::S(s.clone())), + (ScalarAttributeType::N, Bson::Decimal128(d)) => Ok(AttributeValue::N(d.to_string())), + (ScalarAttributeType::B, Bson::String(hex)) => Ok(AttributeValue::B(hex_to_bytes(hex)?)), + _ => Err(StorageError::Internal(format!( + "{field} holds a {} value, not the sort key type", + value.element_type() as u8 + ))), + } +} + +fn hex_to_bytes(hex: &str) -> Result, StorageError> { + let bytes = hex.as_bytes(); + if !bytes.len().is_multiple_of(2) { + return Err(StorageError::Internal("sk_b hex has odd length".to_owned())); + } + bytes + .chunks(2) + .map(|pair| { + let s = std::str::from_utf8(pair).map_err(|e| StorageError::Internal(e.to_string()))?; + u8::from_str_radix(s, 16) + .map_err(|e| StorageError::Internal(format!("sk_b is not hex: {e}"))) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use bson::doc; + + fn key(name: &str, key_type: KeyType) -> KeySchemaElement { + KeySchemaElement { + attribute_name: name.to_owned(), + key_type, + } + } + + fn def(name: &str, attribute_type: ScalarAttributeType) -> AttributeDefinition { + AttributeDefinition { + attribute_name: name.to_owned(), + attribute_type, + } + } + + fn hash_schema() -> (Vec, Vec) { + ( + vec![key("pk", KeyType::Hash)], + vec![def("pk", ScalarAttributeType::S)], + ) + } + + fn range_schema( + sk_type: ScalarAttributeType, + ) -> (Vec, Vec) { + ( + vec![key("pk", KeyType::Hash), key("sk", KeyType::Range)], + vec![def("pk", ScalarAttributeType::S), def("sk", sk_type)], + ) + } + + fn item(pairs: &[(&str, AttributeValue)]) -> Item { + pairs + .iter() + .map(|(k, v)| ((*k).to_owned(), v.clone())) + .collect() + } + + #[test] + fn insert_only_fields_carry_the_id_and_the_key_attributes_for_a_hash_key() { + let (schema, defs) = hash_schema(); + let key = item(&[("pk", AttributeValue::S("u1".to_owned()))]); + let fields = insert_only_key_fields(&key, &schema, &defs).expect("fields"); + assert_eq!(fields, doc! { "_id": "u1", "item_data.pk": { "S": "u1" } }); + } + + #[test] + fn insert_only_fields_carry_both_key_attributes_for_a_range_key() { + let (schema, defs) = range_schema(ScalarAttributeType::N); + let key = item(&[ + ("pk", AttributeValue::S("u1".to_owned())), + ("sk", AttributeValue::N("42".to_owned())), + ]); + let fields = insert_only_key_fields(&key, &schema, &defs).expect("fields"); + let full = item_to_document(&key, &schema, &defs).expect("document"); + assert_eq!(fields.get("_id"), full.get("_id")); + assert_eq!( + fields.get("item_data.pk"), + Some(&bson::bson!({ "S": "u1" })) + ); + assert_eq!( + fields.get("item_data.sk"), + Some(&bson::bson!({ "N": "42" })) + ); + assert_eq!( + fields.len(), + 3, + "no pk or sk_n field: the filter supplies those" + ); + } + + /// The document MongoDB inserts for the pre-fix upsert on a hash table, + /// as measured on MongoDB 7.0: the filter's `pk`, the `$set` attribute, + /// the `$inc` version, and a generated `_id`. + fn stored_without_key_hash() -> Document { + doc! { + "_id": bson::oid::ObjectId::new(), + "pk": "u1", + "item_data": { "w": { "S": "upd" } }, + "_v": 3_i64, + } + } + + #[test] + fn a_generated_object_id_marks_a_document_stored_without_its_key() { + let (schema, defs) = hash_schema(); + assert!(is_stored_without_key(&stored_without_key_hash())); + let good = item_to_document( + &item(&[("pk", AttributeValue::S("u1".to_owned()))]), + &schema, + &defs, + ) + .expect("document"); + assert!(!is_stored_without_key(&good)); + } + + #[test] + fn repair_restores_a_hash_key_and_keeps_attributes_and_version() { + let (schema, defs) = hash_schema(); + let repaired = + repaired_document(&stored_without_key_hash(), &schema, &defs).expect("repair"); + let want_item = item(&[ + ("pk", AttributeValue::S("u1".to_owned())), + ("w", AttributeValue::S("upd".to_owned())), + ]); + let mut want = item_to_document(&want_item, &schema, &defs).expect("document"); + want.insert("_v", 3_i64); + assert_eq!(repaired, want); + } + + #[test] + fn repair_restores_string_numeric_and_binary_sort_keys() { + let cases = [ + (ScalarAttributeType::S, AttributeValue::S("s1".to_owned())), + (ScalarAttributeType::N, AttributeValue::N("12.5".to_owned())), + ( + ScalarAttributeType::B, + AttributeValue::B(vec![0x01, 0xff, 0x10]), + ), + ]; + for (sk_type, sk) in cases { + let (schema, defs) = range_schema(sk_type); + let want_item = item(&[ + ("pk", AttributeValue::S("u1".to_owned())), + ("sk", sk.clone()), + ("w", AttributeValue::S("upd".to_owned())), + ]); + let want = item_to_document(&want_item, &schema, &defs).expect("document"); + // The pre-fix upsert copied the filter's pk and sk_* fields and + // set only the non-key attribute. + let sk_field = format!("sk_{}", sk_suffix(sk_type)); + let bad = doc! { + "_id": bson::oid::ObjectId::new(), + "pk": "u1", + &sk_field: want.get(&sk_field).cloned().expect("sort key field"), + "item_data": { "w": { "S": "upd" } }, + }; + let repaired = repaired_document(&bad, &schema, &defs).expect("repair"); + assert_eq!(repaired, want, "{sk_type:?} sort key"); + assert_eq!( + document_to_item(&repaired).expect("item"), + want_item, + "{sk_type:?} sort key item" + ); + } + } + + #[test] + fn repair_refuses_a_document_whose_sort_key_field_has_the_wrong_type() { + let (schema, defs) = range_schema(ScalarAttributeType::N); + let bad = doc! { + "_id": bson::oid::ObjectId::new(), + "pk": "u1", + "sk_n": "not a decimal", + "item_data": { "w": { "S": "upd" } }, + }; + assert!(repaired_document(&bad, &schema, &defs).is_err()); + } +} diff --git a/crates/storage-mongodb/src/data/mod.rs b/crates/storage-mongodb/src/data/mod.rs index 45465ba61..d1678862a 100644 --- a/crates/storage-mongodb/src/data/mod.rs +++ b/crates/storage-mongodb/src/data/mod.rs @@ -5,6 +5,8 @@ //! //! Contains document conversion, collection naming, and key extraction utilities. +pub mod key_fields; + use bson::{Document, doc}; #[cfg(test)] diff --git a/crates/storage-mongodb/src/data_engine.rs b/crates/storage-mongodb/src/data_engine.rs index dd8f4cc9f..2183dc7d2 100644 --- a/crates/storage-mongodb/src/data_engine.rs +++ b/crates/storage-mongodb/src/data_engine.rs @@ -809,15 +809,12 @@ impl MongoEngine { { let (fast_filter, took_fast) = match mongo_update { NativeUpdate::Doc(d) => { - let opts = mongodb::options::FindOneAndUpdateOptions::builder() - .upsert(true) - .return_document(ReturnDocument::After) - .build(); - let result = coll - .find_one_and_update(key_filter.clone(), d) - .with_options(opts) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + // The upsert's insert branch must write the document + // `item_to_document` would write (issue #436); see + // `update_create::upsert_native_update`. + let result = self + .upsert_native_update(&coll, &key_filter, d, key, key_info) + .await?; (Some(result), true) } NativeUpdate::Pipeline { @@ -994,11 +991,20 @@ impl MongoEngine { let opts = mongodb::options::ReplaceOptions::builder() .upsert(true) .build(); - coll.replace_one(key_filter.clone(), new_doc) + // A duplicate-key error here means another writer + // created the item after the find_one above. Run the + // transaction again so that it reads that item and + // applies the update to it. + match coll + .replace_one(key_filter.clone(), new_doc) .with_options(opts) .session(&mut session) .await - .map_err(TxErr::from)?; + { + Ok(_) => {} + Err(e) if is_duplicate_key(&e) => return Err(TxErr::Transient), + Err(e) => return Err(TxErr::from(e)), + } } self.sync_indexes_in_session( @@ -4059,6 +4065,16 @@ fn project_item( } } +#[path = "data_engine_key_repair.rs"] +pub(crate) mod key_repair; + +#[path = "data_engine_update_create.rs"] +mod update_create; + +#[cfg(test)] +#[path = "data_engine_live_tests.rs"] +mod live_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/crates/storage-mongodb/src/data_engine_key_repair.rs b/crates/storage-mongodb/src/data_engine_key_repair.rs new file mode 100644 index 000000000..ebc2cbeef --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_key_repair.rs @@ -0,0 +1,222 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Startup repair of base-table documents stored without their key. +//! +//! Before issue #436 was fixed, an UpdateItem that created an item on the +//! native fast path upserted a document with a generated ObjectId `_id` and +//! an `item_data` without the key attributes. Such a document answers +//! GetItem without the key, and every PutItem on the key fails on the +//! immutable `_id`. `_id` cannot be changed in place, so the repair replaces +//! each such document with the one `item_to_document` writes, in one +//! transaction: delete the old document by its ObjectId, insert the rebuilt +//! one. The scan is a bracket query on the `_id` index +//! (`{_id: {$type: "objectId"}}`), so a collection with no such document costs +//! one index probe. +//! +//! `serve` runs this once at startup, before the server accepts requests +//! (see `server_components_factory` in `lib.rs`). A failure on one document +//! is logged and skipped; the summary reports it. + +use bson::{Document, doc}; +use extenddb_core::types::{AttributeDefinition, KeySchemaElement}; +use extenddb_storage::error::StorageError; +use futures::TryStreamExt; + +use super::{CommitOutcome, TRANSIENT_RETRY_ATTEMPTS, backoff_sleep, is_transient_write_conflict}; +use crate::MongoEngine; +use crate::data::data_collection_name; +use crate::data::key_fields::{is_stored_without_key, repaired_document}; + +/// What a repair run did. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub struct KeyRepairSummary { + /// Tables whose data collection was scanned. + pub tables_scanned: u64, + /// Documents rewritten with their key. + pub documents_repaired: u64, + /// Documents left as they were because they could not be rebuilt or + /// written; each is logged. + pub documents_failed: u64, +} + +impl KeyRepairSummary { + #[must_use] + pub fn is_clean(&self) -> bool { + self.documents_repaired == 0 && self.documents_failed == 0 + } +} + +impl std::fmt::Display for KeyRepairSummary { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{} tables scanned, {} documents stored without their key repaired, {} failed", + self.tables_scanned, self.documents_repaired, self.documents_failed + ) + } +} + +impl MongoEngine { + /// Repair every base-table document stored without its key, across all + /// tables in the catalog. See the module comment. + pub async fn repair_items_stored_without_key(&self) -> Result { + let tables = self.catalog_db.collection::("tables"); + let mut cursor = tables + .find(doc! {}) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + let mut summary = KeyRepairSummary::default(); + while let Some(table_doc) = cursor + .try_next() + .await + .map_err(|e| StorageError::Internal(e.to_string()))? + { + let (table_id, key_schema, attribute_definitions) = match table_key(&table_doc) { + Ok(parsed) => parsed, + Err(e) => { + tracing::warn!("key repair: skipping a catalog table row: {e}"); + continue; + } + }; + summary.tables_scanned += 1; + self.repair_table(&table_id, &key_schema, &attribute_definitions, &mut summary) + .await?; + } + Ok(summary) + } + + /// Repair the documents of one table's data collection. + pub(crate) async fn repair_table( + &self, + table_id: &str, + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], + summary: &mut KeyRepairSummary, + ) -> Result<(), StorageError> { + let coll = self + .data_db + .collection::(&data_collection_name(table_id)); + let mut cursor = coll + .find(doc! { "_id": { "$type": "objectId" } }) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + while let Some(bad) = cursor + .try_next() + .await + .map_err(|e| StorageError::Internal(e.to_string()))? + { + if !is_stored_without_key(&bad) { + continue; + } + let repaired = match repaired_document(&bad, key_schema, attribute_definitions) { + Ok(d) => d, + Err(e) => { + tracing::error!( + "key repair: table {table_id}: cannot rebuild document {:?}: {e}", + bad.get("_id") + ); + summary.documents_failed += 1; + continue; + } + }; + match self.replace_document(&coll, &bad, repaired).await { + Ok(true) => summary.documents_repaired += 1, + Ok(false) => {} + Err(e) => { + tracing::error!( + "key repair: table {table_id}: cannot rewrite document {:?}: {e}", + bad.get("_id") + ); + summary.documents_failed += 1; + } + } + } + Ok(()) + } + + /// Delete `bad` by its `_id` and insert `repaired` in one transaction. + /// Returns false when `bad` was gone already (another process repaired or + /// deleted it), in which case nothing is inserted. + async fn replace_document( + &self, + coll: &mongodb::Collection, + bad: &Document, + repaired: Document, + ) -> Result { + let bad_id = bad + .get("_id") + .cloned() + .ok_or_else(|| StorageError::Internal("document has no _id".to_owned()))?; + let mut session = self + .client + .start_session() + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + let tx_options = self.write_transaction_options(); + for attempt in 0..TRANSIENT_RETRY_ATTEMPTS { + session + .start_transaction() + .with_options(tx_options.clone()) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + let body = async { + let deleted = coll + .delete_one(doc! { "_id": bad_id.clone() }) + .session(&mut session) + .await?; + if deleted.deleted_count == 0 { + return Ok::(false); + } + coll.insert_one(repaired.clone()) + .session(&mut session) + .await?; + Ok(true) + } + .await; + match body { + Ok(false) => { + let _ = session.abort_transaction().await; + return Ok(false); + } + Ok(true) => match self.commit_write_transaction(&mut session, None).await? { + CommitOutcome::Committed => return Ok(true), + CommitOutcome::RerunBody => backoff_sleep(attempt).await, + }, + Err(e) if is_transient_write_conflict(&e) => { + let _ = session.abort_transaction().await; + backoff_sleep(attempt).await; + } + Err(e) => { + let _ = session.abort_transaction().await; + return Err(StorageError::Internal(e.to_string())); + } + } + } + Err(StorageError::TransactionConflict( + "key repair: too many concurrent write conflicts, giving up".to_owned(), + )) + } +} + +/// `table_id`, key schema and attribute definitions of a catalog `tables` row. +fn table_key( + table_doc: &Document, +) -> Result<(String, Vec, Vec), StorageError> { + let table_id = table_doc + .get_str("table_id") + .map_err(|_| StorageError::Internal("missing table_id".to_owned()))? + .to_owned(); + let key_schema = table_doc + .get("key_schema") + .ok_or_else(|| StorageError::Internal(format!("table {table_id}: missing key_schema")))?; + let key_schema: Vec = bson::from_bson(key_schema.clone()) + .map_err(|e| StorageError::Internal(format!("table {table_id}: key_schema: {e}")))?; + let attribute_definitions = table_doc.get("attribute_definitions").ok_or_else(|| { + StorageError::Internal(format!("table {table_id}: missing attribute_definitions")) + })?; + let attribute_definitions: Vec = + bson::from_bson(attribute_definitions.clone()).map_err(|e| { + StorageError::Internal(format!("table {table_id}: attribute_definitions: {e}")) + })?; + Ok((table_id, key_schema, attribute_definitions)) +} diff --git a/crates/storage-mongodb/src/data_engine_live_tests.rs b/crates/storage-mongodb/src/data_engine_live_tests.rs new file mode 100644 index 000000000..1c8429bb3 --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_live_tests.rs @@ -0,0 +1,507 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for the shape of the document an UpdateItem writes when +//! it creates an item (issue #436), and for the startup repair of documents an +//! earlier build stored without their key. +//! +//! They read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are +//! `#[ignore]`d so `cargo test --lib` passes without one; run them with +//! `cargo test -p extenddb-storage-mongodb --lib -- --ignored live_`. Each test +//! uses its own catalog and data databases and drops them at the end. + +use std::collections::HashMap; + +use bson::{Document, doc}; +use extenddb_core::expression::{Expr, ExpressionMaps, PathElement, UpdateAction}; +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, Item, KeySchemaElement, KeyType, ScalarAttributeType, + TableKeyInfo, +}; + +use extenddb_storage::error::StorageError; + +use super::key_repair::KeyRepairSummary; +use crate::MongoEngine; +use crate::data::{data_collection_name, document_to_item, item_to_document, pk_filter}; + +fn uri() -> String { + std::env::var("EXTENDDB_TEST_MONGODB_URI").expect( + "set EXTENDDB_TEST_MONGODB_URI to a MongoDB replica set, e.g. \ + mongodb://127.0.0.1:27741/?directConnection=true", + ) +} + +/// An engine on fresh catalog and data databases. +async fn scratch_engine() -> MongoEngine { + let base = MongoEngine::new( + &uri(), + "us-east-1", + None, + mongodb::options::ReadConcern::snapshot(), + ) + .await + .expect("engine"); + let tag = uuid::Uuid::new_v4().simple().to_string(); + MongoEngine { + catalog_db: base.client.database(&format!("k{}_cat", &tag[..12])), + data_db: base.client.database(&format!("k{}_data", &tag[..12])), + ..base + } +} + +async fn finish(engine: MongoEngine) { + engine.catalog_db.drop().await.expect("drop catalog"); + engine.data_db.drop().await.expect("drop data"); +} + +fn key(name: &str, key_type: KeyType) -> KeySchemaElement { + KeySchemaElement { + attribute_name: name.to_owned(), + key_type, + } +} + +fn def(name: &str, attribute_type: ScalarAttributeType) -> AttributeDefinition { + AttributeDefinition { + attribute_name: name.to_owned(), + attribute_type, + } +} + +/// Key info for a table with a string hash key `pk` and, with `range`, a +/// string sort key `sk`. The GSI cache is set to "no indexes" so an +/// unconditional UpdateItem takes the native fast path, as it does on a +/// served table without indexes once the first write has filled the cache. +fn table(engine: &MongoEngine, table_id: &str, range: bool) -> TableKeyInfo { + engine.gsi_cache_set(table_id, false); + let mut key_schema = vec![key("pk", KeyType::Hash)]; + let mut attribute_definitions = vec![def("pk", ScalarAttributeType::S)]; + if range { + key_schema.push(key("sk", KeyType::Range)); + attribute_definitions.push(def("sk", ScalarAttributeType::S)); + } + TableKeyInfo { + table_id: table_id.to_owned(), + table_name: table_id.to_owned(), + key_schema: key_schema.clone(), + base_key_schema: key_schema, + attribute_definitions, + ..Default::default() + } +} + +fn item_key(range: bool) -> Item { + let mut key = Item::new(); + key.insert("pk".to_owned(), AttributeValue::S("u1".to_owned())); + if range { + key.insert("sk".to_owned(), AttributeValue::S("s1".to_owned())); + } + key +} + +fn with(mut item: Item, name: &str, value: AttributeValue) -> Item { + item.insert(name.to_owned(), value); + item +} + +fn s(v: &str) -> AttributeValue { + AttributeValue::S(v.to_owned()) +} + +fn n(v: &str) -> AttributeValue { + AttributeValue::N(v.to_owned()) +} + +fn path(name: &str) -> Vec { + vec![PathElement::Attribute(name.to_owned())] +} + +fn maps(value: AttributeValue) -> ExpressionMaps { + let mut values = HashMap::new(); + values.insert(":v".to_owned(), value); + ExpressionMaps::new(HashMap::new(), values) +} + +async fn stored(engine: &MongoEngine, key_info: &TableKeyInfo, key: &Item) -> Document { + let coll = engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)); + let filter = + pk_filter(key, &key_info.key_schema, &key_info.attribute_definitions).expect("filter"); + coll.find_one(filter) + .await + .expect("read back") + .expect("the item's document exists") +} + +/// The document `item_to_document` writes for `item`, with version `v`. +fn expected_document(item: &Item, key_info: &TableKeyInfo, v: i64) -> Document { + let mut doc = item_to_document(item, &key_info.key_schema, &key_info.attribute_definitions) + .expect("document"); + doc.insert("_v", v); + doc +} + +/// An UpdateItem that creates the item, on the native fast path (SET with a +/// placeholder, REMOVE) and on the transaction path (numeric ADD), stores the +/// same document PutItem would store: `_id` as the key text, the key +/// attributes in `item_data`, version 1. ALL_NEW returns the item with its +/// key. A PutItem on the key then replaces the item and returns it as the old +/// image. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_update_item_that_creates_an_item_stores_its_key_like_put_item() { + let engine = scratch_engine().await; + let cases: [(&str, Vec, AttributeValue, Item); 3] = [ + ( + "set", + vec![UpdateAction::Set { + path: path("w"), + value: Expr::Placeholder(":v".to_owned()), + }], + s("upd"), + with(Item::new(), "w", s("upd")), + ), + ( + "remove", + vec![UpdateAction::Remove { path: path("w") }], + s("unused"), + Item::new(), + ), + ( + "add", + vec![UpdateAction::Add { + path: path("c"), + value: Expr::Placeholder(":v".to_owned()), + }], + n("5"), + with(Item::new(), "c", n("5")), + ), + ]; + for range in [false, true] { + for (name, actions, value, extra) in &cases { + let table_id = format!("{name}-{}", if range { "range" } else { "hash" }); + let key_info = table(&engine, &table_id, range); + let key = item_key(range); + let mut want = key.clone(); + want.extend(extra.clone()); + + let (old, new) = engine + .update_item_impl( + &key_info, + &key, + actions, + false, + true, + None, + &maps(value.clone()), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: update: {e:?}")); + assert_eq!(old, None, "{table_id}: no old image on creation"); + assert_eq!( + new, + Some(want.clone()), + "{table_id}: ALL_NEW carries the key" + ); + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&want, &key_info, 1), + "{table_id}: stored document" + ); + + let put = with(key.clone(), "w", s("put")); + let old = engine + .put_item_impl( + &key_info, + put.clone(), + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: put after update: {e:?}")); + assert_eq!(old, Some(want), "{table_id}: put returns the created item"); + let after = stored(&engine, &key_info, &key).await; + assert_eq!( + document_to_item(&after).expect("item"), + put, + "{table_id}: put replaced the item" + ); + } + } + finish(engine).await; +} + +/// Sixty-four UpdateItem calls creating the same item at once, on the fast +/// path, from eight runtime threads: every call succeeds, the stored +/// document has the shape PutItem writes, and `_v` counts every call, so +/// each one applied to the one document after losing the create race. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_concurrent_creating_updates_on_one_key_all_apply() { + const WRITERS: i64 = 64; + let engine = std::sync::Arc::new(scratch_engine().await); + for range in [false, true] { + let table_id = format!("race-{}", if range { "range" } else { "hash" }); + let key_info = std::sync::Arc::new(table(&engine, &table_id, range)); + let key = item_key(range); + let mut tasks = Vec::new(); + for i in 0..WRITERS { + let engine = engine.clone(); + let key_info = key_info.clone(); + let key = key.clone(); + tasks.push(tokio::spawn(async move { + let actions = vec![UpdateAction::Set { + path: path("w"), + value: Expr::Placeholder(":v".to_owned()), + }]; + engine + .update_item_impl( + &key_info, + &key, + &actions, + false, + true, + None, + &maps(s(&format!("w{i}"))), + None, + ) + .await + })); + } + for task in tasks { + let (old, new) = task + .await + .expect("task") + .unwrap_or_else(|e| panic!("{table_id}: a creating update failed: {e:?}")); + assert_eq!(old, None, "{table_id}"); + let new = new.expect("ALL_NEW"); + for (name, value) in &key { + assert_eq!( + new.get(name), + Some(value), + "{table_id}: ALL_NEW has the key" + ); + } + } + let doc = stored(&engine, &key_info, &key).await; + let item = document_to_item(&doc).expect("item"); + let w = item.get("w").cloned().expect("w was set"); + let mut want = key.clone(); + want.insert("w".to_owned(), w); + assert_eq!( + doc, + expected_document(&want, &key_info, WRITERS), + "{table_id}" + ); + } + let engine = std::sync::Arc::try_unwrap(engine) + .ok() + .expect("all tasks done"); + finish(engine).await; +} + +/// A creating upsert whose insert always collides on `_id` (another document +/// holds the key text as its `_id` under a different `pk`) exhausts the +/// retry budget and answers `TransactionConflict`, which the server maps to +/// `TransactionConflictException`, not an internal error: the duplicate-key +/// command error `findAndModify` returns is recognised. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_creating_update_that_always_loses_on_id_answers_transaction_conflict() { + let engine = scratch_engine().await; + let key_info = table(&engine, "collide-hash", false); + let key = item_key(false); + engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)) + .insert_one(doc! { "_id": "u1", "pk": "other", "item_data": { "pk": { "S": "other" } } }) + .await + .expect("colliding document"); + let actions = vec![UpdateAction::Set { + path: path("w"), + value: Expr::Placeholder(":v".to_owned()), + }]; + let err = engine + .update_item_impl( + &key_info, + &key, + &actions, + false, + true, + None, + &maps(s("upd")), + None, + ) + .await + .expect_err("the insert cannot succeed"); + assert!( + matches!(err, StorageError::TransactionConflict(_)), + "expected TransactionConflict, got {err:?}" + ); + finish(engine).await; +} + +/// Store a document the way the fast path did before the fix: an upsert whose +/// filter is the key and whose update sets one attribute. MongoDB copies the +/// filter's `pk` and `sk_*` fields, generates an ObjectId `_id`, and writes +/// no key attribute into `item_data`. +async fn store_without_key(engine: &MongoEngine, key_info: &TableKeyInfo, key: &Item, w: &str) { + let coll = engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)); + let filter = + pk_filter(key, &key_info.key_schema, &key_info.attribute_definitions).expect("filter"); + coll.find_one_and_update( + filter, + doc! { "$set": { "item_data.w": { "S": w } }, "$inc": { "_v": 1_i64 } }, + ) + .upsert(true) + .await + .expect("upsert"); +} + +/// The repair rewrites a document stored without its key into the document +/// PutItem would have written, keeping the attributes and the version. The +/// item then reads with its key and a PutItem on it succeeds. Documents with +/// their key are left alone, and a second run finds nothing. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_rewrites_documents_stored_without_their_key() { + let engine = scratch_engine().await; + for range in [false, true] { + let table_id = format!("repair-{}", if range { "range" } else { "hash" }); + let key_info = table(&engine, &table_id, range); + let key = item_key(range); + store_without_key(&engine, &key_info, &key, "upd").await; + store_without_key(&engine, &key_info, &key, "upd2").await; + let bad = stored(&engine, &key_info, &key).await; + assert!( + bad.get_object_id("_id").is_ok(), + "{table_id}: the pre-fix shape has an ObjectId _id" + ); + assert!( + !document_to_item(&bad).expect("item").contains_key("pk"), + "{table_id}: and no key" + ); + + // A healthy item in the same collection is not touched. + let other = with(with(item_key(range), "w", s("fine")), "pk", s("u2")); + engine + .put_item_impl( + &key_info, + other.clone(), + false, + None, + &ExpressionMaps::default(), + None, + ) + .await + .expect("put the healthy item"); + + let mut summary = KeyRepairSummary::default(); + engine + .repair_table( + &table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut summary, + ) + .await + .expect("repair"); + assert_eq!( + summary, + KeyRepairSummary { + tables_scanned: 0, + documents_repaired: 1, + documents_failed: 0, + }, + "{table_id}" + ); + + let want = with(key.clone(), "w", s("upd2")); + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&want, &key_info, 2), + "{table_id}: repaired document" + ); + assert_eq!( + document_to_item(&stored(&engine, &key_info, &other).await).expect("item"), + other, + "{table_id}: the healthy item is unchanged" + ); + + let put = with(key.clone(), "w", s("put")); + let old = engine + .put_item_impl( + &key_info, + put.clone(), + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: put after repair: {e:?}")); + assert_eq!(old, Some(want), "{table_id}: put returns the repaired item"); + + let mut again = KeyRepairSummary::default(); + engine + .repair_table( + &table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut again, + ) + .await + .expect("second repair"); + assert!( + again.is_clean(), + "{table_id}: nothing left to repair: {again}" + ); + } + finish(engine).await; +} + +/// The whole-catalog entry point finds the table through its catalog row. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_reads_the_key_schema_from_the_catalog() { + let engine = scratch_engine().await; + let key_info = table(&engine, "catalog-hash", false); + engine + .catalog_db + .collection::("tables") + .insert_one(doc! { + "_id": { "account_id": "123456789012", "table_name": "t" }, + "table_id": "catalog-hash", + "table_status": "ACTIVE", + "key_schema": bson::to_bson(&key_info.key_schema).expect("bson"), + "attribute_definitions": bson::to_bson(&key_info.attribute_definitions).expect("bson"), + }) + .await + .expect("catalog row"); + let key = item_key(false); + store_without_key(&engine, &key_info, &key, "upd").await; + + let summary = engine + .repair_items_stored_without_key() + .await + .expect("repair"); + assert_eq!( + summary, + KeyRepairSummary { + tables_scanned: 1, + documents_repaired: 1, + documents_failed: 0, + } + ); + assert_eq!( + document_to_item(&stored(&engine, &key_info, &key).await).expect("item"), + with(key, "w", s("upd")) + ); + finish(engine).await; +} diff --git a/crates/storage-mongodb/src/data_engine_update_create.rs b/crates/storage-mongodb/src/data_engine_update_create.rs new file mode 100644 index 000000000..b3b73fecb --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_update_create.rs @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! The native UpdateItem upsert, which creates the item when it does not +//! exist yet. +//! +//! `update_item_impl` runs an unconditional single-attribute update on a +//! table without indexes as one `findAndModify` with `upsert: true`, filtered +//! on the key fields `pk` and `sk_*`. When no document matches, MongoDB builds +//! the inserted document from those filter fields and the update operators. +//! Without more, that document gets a generated ObjectId `_id` and an +//! `item_data` without the key attributes, while PutItem writes `_id` as the +//! key text and the whole item (`item_to_document`). Every later PutItem on +//! the key then fails on the immutable `_id` (issue #436). +//! [`MongoEngine::upsert_native_update`] adds a `$setOnInsert` with the `_id` +//! and the key attributes, so the created document is the one PutItem would +//! have written, with `_v` 1 from the update's `$inc`. +//! +//! Two writers creating the same item at once each find no document and each +//! insert; the second insert fails with a duplicate-key error on `_id`. The +//! data collection also has a unique index on the key fields, and MongoDB +//! retries an upsert itself only when the error names an index whose fields +//! the filter matches by equality. The `_id` collision is reported first and +//! the filter does not name `_id`, so the server does not retry it and this +//! module does: the item now exists, so the next attempt updates it. + +use bson::Document; +use extenddb_core::types::{Item, TableKeyInfo}; +use extenddb_storage::error::StorageError; +use mongodb::options::{FindOneAndUpdateOptions, ReturnDocument}; + +use super::{TRANSIENT_RETRY_ATTEMPTS, backoff_sleep}; +use crate::MongoEngine; +use crate::data::key_fields::insert_only_key_fields; + +/// Whether `e` is a duplicate-key error (E11000) as `findAndModify` reports +/// it: a command error, where `insert_one` reports a write error. +fn is_duplicate_key_command_error(e: &mongodb::error::Error) -> bool { + matches!(*e.kind, mongodb::error::ErrorKind::Command(ref c) if c.code == 11000) +} + +impl MongoEngine { + /// Apply the native `update` (a `$set`/`$unset`/`$inc` operator document) + /// to the item at `key_filter`, creating the item when it does not exist, + /// and return the document after the update. + /// + /// `key` and `key_info` supply the `$setOnInsert` fields that make a + /// created document match what PutItem writes. Attempts that lose the + /// create race share the transient retry budget and back off between + /// them; when every attempt loses, the error is the same + /// `TransactionConflict` the transaction path returns. + pub(crate) async fn upsert_native_update( + &self, + coll: &mongodb::Collection, + key_filter: &Document, + mut update: Document, + key: &Item, + key_info: &TableKeyInfo, + ) -> Result, StorageError> { + update.insert( + "$setOnInsert", + insert_only_key_fields(key, &key_info.key_schema, &key_info.attribute_definitions)?, + ); + let opts = FindOneAndUpdateOptions::builder() + .upsert(true) + .return_document(ReturnDocument::After) + .build(); + for attempt in 0..TRANSIENT_RETRY_ATTEMPTS { + match coll + .find_one_and_update(key_filter.clone(), update.clone()) + .with_options(opts.clone()) + .await + { + Ok(doc) => return Ok(doc), + Err(e) if is_duplicate_key_command_error(&e) => backoff_sleep(attempt).await, + Err(e) => return Err(StorageError::Internal(e.to_string())), + } + } + Err(StorageError::TransactionConflict( + "UpdateItem: too many concurrent write conflicts, giving up".to_owned(), + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use bson::doc; + + fn command_error(code: i32) -> mongodb::error::Error { + let command_error: mongodb::error::CommandError = bson::from_document(doc! { + "code": code, + "codeName": "SyntheticTestError", + "errmsg": "synthetic command error", + }) + .expect("static CommandError document deserializes"); + mongodb::error::ErrorKind::Command(command_error).into() + } + + fn write_error(code: i32) -> mongodb::error::Error { + let write_error: mongodb::error::WriteError = bson::from_document(doc! { + "code": code, + "codeName": "SyntheticTestError", + "errmsg": "synthetic write error", + }) + .expect("static WriteError document deserializes"); + mongodb::error::ErrorKind::Write(mongodb::error::WriteFailure::WriteError(write_error)) + .into() + } + + #[test] + fn only_the_command_form_of_e11000_is_a_lost_create_race() { + assert!(is_duplicate_key_command_error(&command_error(11000))); + assert!(!is_duplicate_key_command_error(&command_error(112))); + assert!(!is_duplicate_key_command_error(&write_error(11000))); + } +} diff --git a/crates/storage-mongodb/src/lib.rs b/crates/storage-mongodb/src/lib.rs index b27cc0baf..8c1de6c6c 100644 --- a/crates/storage-mongodb/src/lib.rs +++ b/crates/storage-mongodb/src/lib.rs @@ -137,6 +137,18 @@ fn server_components_factory( Err(e) => tracing::error!("Failed to sweep stale backup claims: {e}"), } + // An UpdateItem on a build before the fix for issue #436 could create + // an item whose document has no key. Rewrite any such document before + // serving, so that reads return the key and PutItem on it succeeds + // (see `data_engine::key_repair`). Each rewrite is its own + // transaction, so running this while another process serves the same + // data is safe. + match engine.repair_items_stored_without_key().await { + Ok(repair) if repair.is_clean() => {} + Ok(repair) => tracing::warn!("startup key repair: {repair}"), + Err(e) => tracing::error!("Failed to repair items stored without their key: {e}"), + } + let engine = Arc::new(engine); // Create one shared catalog/management client. MongoDB Client clones diff --git a/tests/test_update_item_creates_key.py b/tests/test_update_item_creates_key.py new file mode 100644 index 000000000..89ebc50ad --- /dev/null +++ b/tests/test_update_item_creates_key.py @@ -0,0 +1,363 @@ +# Copyright 2026 ExtendDB contributors +# SPDX-License-Identifier: Apache-2.0 + +"""An UpdateItem that creates an item stores the item with its key. + +Amazon DynamoDB creates the item from the key and the update expression: +GetItem returns the key attributes with the updated ones, ALL_NEW on creation +returns the whole new item including the key, UPDATED_NEW returns only the +updated attributes, ALL_OLD and UPDATED_OLD on creation return no Attributes, +and a later PutItem, UpdateItem, Query, Scan and DeleteItem on the key treat +the item as any other. On the MongoDB backend an unconditional UpdateItem on a +table without an index used to store the created item without its key +attributes, after which GetItem returned the item without the key and every +PutItem on the key answered HTTP 500 (issue #436). The race tests run the +issue's shape: writers creating the same keys with UpdateItem and PutItem at +once. + +Expected responses were measured against the service on 2026-10-23 (probe +W4-1 in the fix report). +""" + +from __future__ import annotations + +import os +import threading +import uuid +from dataclasses import dataclass, field + +import boto3 +import pytest +from botocore.config import Config +from botocore.exceptions import ClientError + +from conftest import scoped_table + +WRITERS = 8 +KEYS = 30 +RETURN_VALUES = ["ALL_NEW", "UPDATED_NEW", "ALL_OLD", "UPDATED_OLD", "NONE"] + + +# TEMPORARY: on MongoDB, a put without a condition on a table with no index and +# no stream returns HTTP 500 when it loses the create race to another writer's +# insert. The MongoDB write-race fix repairs it and lands separately. +# TODO: remove this marker once the MongoDB write-race fix is on main. +# Only an assertion counts as the expected failure. The marker is not strict, +# because the race does not fire in every run. 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 returns HTTP 500 for a lost create race, fixed by the MongoDB write-race fix", + raises=AssertionError, + strict=False, +) + + +def _s(name: str) -> dict: + return {"AttributeName": name, "AttributeType": "S"} + + +def _k(name: str, kind: str) -> dict: + return {"AttributeName": name, "KeyType": kind} + + +@pytest.fixture(scope="module") +def hash_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=[_s("pk"), _s("sk")], + key_schema=[_k("pk", "HASH"), _k("sk", "RANGE")], + ) as name: + yield name + + +@pytest.fixture(scope="module") +def raw_client(endpoint_url): + """A client that never retries, so every failure 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=WRITERS * 4, + ), + } + if endpoint_url: + kwargs["endpoint_url"] = endpoint_url + if endpoint_url.startswith("https://"): + kwargs["verify"] = False + return boto3.client(**kwargs) + + +def _key(table_kind: str, pk: str) -> dict: + key = {"pk": {"S": pk}} + if table_kind == "range": + key["sk"] = {"S": "s"} + return key + + +def _get(client, table: str, key: dict) -> dict | None: + return client.get_item(TableName=table, Key=key, ConsistentRead=True).get("Item") + + +def _warm(client, table: str, table_kind: str) -> None: + """Write another key first. + + The MongoDB backend takes its native UpdateItem path only once a write has + recorded that the table has no index; the first write on a fresh table + takes the transaction path, which did not have the defect. The issue + reproduced it the same way. + """ + client.put_item(TableName=table, Item={**_key(table_kind, f"warm-{uuid.uuid4().hex}")}) + + +def _create(client, table: str, key: dict, **extra) -> dict: + """UpdateItem SET w = upd on a key that does not exist yet.""" + resp = client.update_item( + TableName=table, + Key=key, + UpdateExpression="SET w = :w", + ExpressionAttributeValues={":w": {"S": "upd"}}, + **extra, + ) + resp.pop("ResponseMetadata", None) + return resp + + +def _put_then_get(client, table: str, key: dict) -> None: + """PutItem on a key created by UpdateItem succeeds and replaces the item.""" + put = {**key, "w": {"S": "put"}} + client.put_item(TableName=table, Item=put) + assert _get(client, table, key) == put + + +@pytest.mark.parametrize("table_kind", ["hash", "range"]) +def test_set_creates_the_item_with_its_key(request, dynamodb_client, table_kind): + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + key = _key(table_kind, f"set-{uuid.uuid4().hex}") + assert _create(dynamodb_client, table, key) == {} + assert _get(dynamodb_client, table, key) == {**key, "w": {"S": "upd"}} + _put_then_get(dynamodb_client, table, key) + + +@pytest.mark.parametrize("table_kind", ["hash", "range"]) +def test_add_number_creates_the_item_with_its_key(request, dynamodb_client, table_kind): + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + key = _key(table_kind, f"add-{uuid.uuid4().hex}") + dynamodb_client.update_item( + TableName=table, + Key=key, + UpdateExpression="ADD c :one", + ExpressionAttributeValues={":one": {"N": "1"}}, + ) + assert _get(dynamodb_client, table, key) == {**key, "c": {"N": "1"}} + _put_then_get(dynamodb_client, table, key) + + +def test_add_set_creates_the_item_with_its_key(dynamodb_client, hash_table): + _warm(dynamodb_client, hash_table, "hash") + key = _key("hash", f"add-set-{uuid.uuid4().hex}") + dynamodb_client.update_item( + TableName=hash_table, + Key=key, + UpdateExpression="ADD tags :t", + ExpressionAttributeValues={":t": {"SS": ["a", "b"]}}, + ) + item = _get(dynamodb_client, hash_table, key) + assert item is not None + assert item["pk"] == key["pk"] + assert sorted(item["tags"]["SS"]) == ["a", "b"] + _put_then_get(dynamodb_client, hash_table, key) + + +@pytest.mark.parametrize("table_kind", ["hash", "range"]) +def test_remove_only_on_a_missing_item_succeeds_and_the_key_stays_writable( + request, dynamodb_client, table_kind +): + """The service answers 200 and creates no item from a REMOVE-only update. + + ExtendDB creates an item holding the key (measured on SQLite and on + MongoDB); that difference is older than this test and is not asserted + either way here. + On MongoDB the call itself answered 500 before the fix, because the + created document had no item_data at all. + """ + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + key = _key(table_kind, f"remove-{uuid.uuid4().hex}") + resp = dynamodb_client.update_item(TableName=table, Key=key, UpdateExpression="REMOVE w") + assert "Attributes" not in resp + assert _get(dynamodb_client, table, key) in (None, key) + _put_then_get(dynamodb_client, table, key) + + +def test_attribute_not_exists_condition_creates_the_item_once(dynamodb_client, hash_table): + _warm(dynamodb_client, hash_table, "hash") + key = _key("hash", f"cond-{uuid.uuid4().hex}") + resp = _create(dynamodb_client, hash_table, key, ConditionExpression="attribute_not_exists(pk)") + assert resp == {} + assert _get(dynamodb_client, hash_table, key) == {**key, "w": {"S": "upd"}} + with pytest.raises(ClientError) as exc: + _create(dynamodb_client, hash_table, key, ConditionExpression="attribute_not_exists(pk)") + assert exc.value.response["Error"]["Code"] == "ConditionalCheckFailedException" + assert exc.value.response["Error"]["Message"] == "The conditional request failed" + _put_then_get(dynamodb_client, hash_table, key) + + +@pytest.mark.parametrize("return_values", RETURN_VALUES) +@pytest.mark.parametrize("table_kind", ["hash", "range"]) +def test_return_values_on_creation(request, dynamodb_client, table_kind, return_values): + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + key = _key(table_kind, f"rv-{return_values.lower()}-{uuid.uuid4().hex}") + resp = _create(dynamodb_client, table, key, ReturnValues=return_values) + created = {**key, "w": {"S": "upd"}} + if return_values == "ALL_NEW": + assert resp == {"Attributes": created} + elif return_values == "UPDATED_NEW": + assert resp == {"Attributes": {"w": {"S": "upd"}}} + else: + assert resp == {} + assert _get(dynamodb_client, table, key) == created + _put_then_get(dynamodb_client, table, key) + + +@pytest.mark.parametrize("table_kind", ["hash", "range"]) +def test_later_operations_on_a_key_created_by_update_item(request, dynamodb_client, table_kind): + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + key = _key(table_kind, f"later-{uuid.uuid4().hex}") + created = {**key, "w": {"S": "upd"}} + _create(dynamodb_client, table, key) + assert _get(dynamodb_client, table, key) == created + + put = {**key, "w": {"S": "put"}} + resp = dynamodb_client.put_item(TableName=table, Item=put, ReturnValues="ALL_OLD") + assert resp["Attributes"] == created + assert _get(dynamodb_client, table, key) == put + + resp = dynamodb_client.update_item( + TableName=table, + Key=key, + UpdateExpression="SET x = :x", + ExpressionAttributeValues={":x": {"N": "1"}}, + ReturnValues="ALL_NEW", + ) + final = {**put, "x": {"N": "1"}} + assert resp["Attributes"] == final + + resp = dynamodb_client.query( + TableName=table, + KeyConditionExpression="pk = :p", + ExpressionAttributeValues={":p": key["pk"]}, + ConsistentRead=True, + ) + assert resp["Count"] == 1 + assert resp["Items"] == [final] + + scanned = dynamodb_client.scan(TableName=table, ConsistentRead=True)["Items"] + assert final in scanned + assert all("pk" in item for item in scanned) + + resp = dynamodb_client.delete_item(TableName=table, Key=key, ReturnValues="ALL_OLD") + assert resp["Attributes"] == final + assert _get(dynamodb_client, table, key) is None + + resp = dynamodb_client.put_item(TableName=table, Item=put, ReturnValues="ALL_OLD") + assert "Attributes" not in resp + assert _get(dynamodb_client, table, key) == put + + +@dataclass +class RaceResult: + table: str + table_kind: str + keys: list[dict] + failures: list[tuple[str, str, str]] = field(default_factory=list) + + def failed(self, op: str) -> list[tuple[str, str, str]]: + return [f for f in self.failures if f[0] == op] + + +@pytest.fixture(scope="module", params=["hash", "range"]) +def race(request, dynamodb_client, raw_client) -> RaceResult: + """WRITERS UpdateItem and WRITERS PutItem per key, KEYS keys, all at once. + + The issue measured 128 of 240 PutItem failing with HTTP 500 and 16 of 30 + final items without `pk` on the MongoDB backend. + """ + table_kind = request.param + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + result = RaceResult( + table, table_kind, [_key(table_kind, f"race-{uuid.uuid4().hex}") for _ in range(KEYS)] + ) + start = threading.Barrier(2 * WRITERS, timeout=60) + errors: list[BaseException] = [] + lock = threading.Lock() + + def run(op: str, i: int): + try: + start.wait() + for key in result.keys: + try: + if op == "put": + raw_client.put_item(TableName=table, Item={**key, "w": {"S": f"put{i}"}}) + else: + raw_client.update_item( + TableName=table, + Key=key, + UpdateExpression="SET w = :w", + ExpressionAttributeValues={":w": {"S": f"upd{i}"}}, + ) + except ClientError as e: + with lock: + result.failures.append((op, key["pk"]["S"], e.response["Error"]["Code"])) + except BaseException as e: # noqa: BLE001 - surfaced below + errors.append(e) + + threads = [ + threading.Thread(target=run, args=(op, i)) + for op in ("put", "update") + for i in range(WRITERS) + ] + for t in threads: + t.start() + for t in threads: + t.join() + if errors: + raise errors[0] + return result + + +def test_racing_update_item_calls_all_succeed(race): + failed = race.failed("update") + assert not failed, f"{len(failed)} of {WRITERS * KEYS} UpdateItem failed: {failed[:5]}" + + +def test_racing_writers_leave_every_item_with_its_key(dynamodb_client, race): + without_key = [] + for key in race.keys: + item = _get(dynamodb_client, race.table, key) + if item is None or any(item.get(name) != value for name, value in key.items()): + without_key.append(key["pk"]["S"]) + assert not without_key, f"{len(without_key)} of {KEYS} items lack a key attribute" + + +def test_put_item_after_the_race_replaces_every_item(dynamodb_client, race): + for key in race.keys: + _put_then_get(dynamodb_client, race.table, key) + + +@XFAIL_UNTIL_MONGODB_FIX +def test_racing_put_item_calls_all_succeed(race): + failed = race.failed("put") + assert not failed, f"{len(failed)} of {WRITERS * KEYS} PutItem failed: {failed[:5]}" From 5c5ce7e03b4bc812dbb1c7f4dfce025b37c955e2 Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 02/14] fix(mongodb): repair a pre-fix document that has no item_data A REMOVE-only UpdateItem that created an item on the native fast path before the fix for issue #436 upserted a document with no item_data at all: the filter supplied pk and sk_*, the $unset contributed nothing and the $inc gave _v. The startup repair could not rebuild it, because repaired_document went through document_to_item, which fails with "Document missing item_data field". The document was counted failed and left in place, so its key answered 500 on GetItem, PutItem and UpdateItem on that start and every later one. Treat a missing item_data as an empty item, which rebuilds the document into the key-only item the fixed fast path writes for the same request. Two unit tests cover the shape for a hash key and a range key, and a live test plants it, reads the key as an error before the repair and as the key-only item after it. The live tests' store_without_key helper takes Some(attribute) for the SET form and None for the REMOVE-only form. --- crates/storage-mongodb/src/data/key_fields.rs | 51 +++++++- .../src/data_engine_live_tests.rs | 118 +++++++++++++++--- 2 files changed, 152 insertions(+), 17 deletions(-) diff --git a/crates/storage-mongodb/src/data/key_fields.rs b/crates/storage-mongodb/src/data/key_fields.rs index 3eea9fb5b..d28c37c27 100644 --- a/crates/storage-mongodb/src/data/key_fields.rs +++ b/crates/storage-mongodb/src/data/key_fields.rs @@ -11,8 +11,9 @@ //! with `$setOnInsert`: [`insert_only_key_fields`] builds that document. //! //! Before that was done, such an upsert stored a document with a generated -//! ObjectId `_id` and no key attributes in `item_data` (issue #436). -//! [`repaired_document`] rebuilds a document of that shape into the shape +//! ObjectId `_id` and no key attributes in `item_data` (issue #436), and for +//! a REMOVE-only update no `item_data` at all. +//! [`repaired_document`] rebuilds a document of either shape into the shape //! `item_to_document` writes, reading the key values back from the top-level //! `pk` and `sk_*` fields. @@ -68,6 +69,12 @@ pub fn is_stored_without_key(doc: &Document) -> bool { /// `item_to_document` writes: the same `item_data` with the key attributes /// added, `_id` as the key text, and `_v` carried over. /// +/// A document with no `item_data` at all rebuilds into the key-only item. +/// That is the shape the pre-fix upsert stored for a REMOVE-only creating +/// UpdateItem, where the filter supplied `pk` and `sk_*`, the `$unset` +/// contributed nothing and the `$inc` gave `_v`; it is also what the fixed +/// fast path now writes for the same request. +/// /// The key values come from the top-level `pk` text and `sk_s`/`sk_n`/`sk_b` /// field, typed by `attribute_definitions`. A document whose top-level key /// fields are missing or of the wrong type is an error; the caller leaves it @@ -77,7 +84,11 @@ pub fn repaired_document( key_schema: &[KeySchemaElement], attribute_definitions: &[AttributeDefinition], ) -> Result { - let mut item = document_to_item(doc)?; + let mut item = if doc.contains_key("item_data") { + document_to_item(doc)? + } else { + Item::new() + }; let pk_name = key_schema .iter() @@ -314,6 +325,40 @@ mod tests { } } + #[test] + fn repair_rebuilds_a_document_with_no_item_data_as_the_key_only_item() { + let (schema, defs) = hash_schema(); + let bad = doc! { + "_id": bson::oid::ObjectId::new(), + "pk": "u1", + "_v": 1_i64, + }; + assert!(is_stored_without_key(&bad)); + let repaired = repaired_document(&bad, &schema, &defs).expect("repair"); + let want_item = item(&[("pk", AttributeValue::S("u1".to_owned()))]); + let mut want = item_to_document(&want_item, &schema, &defs).expect("document"); + want.insert("_v", 1_i64); + assert_eq!(repaired, want); + assert_eq!(document_to_item(&repaired).expect("item"), want_item); + } + + #[test] + fn repair_rebuilds_a_range_key_document_with_no_item_data() { + let (schema, defs) = range_schema(ScalarAttributeType::S); + let want_item = item(&[ + ("pk", AttributeValue::S("u1".to_owned())), + ("sk", AttributeValue::S("s1".to_owned())), + ]); + let want = item_to_document(&want_item, &schema, &defs).expect("document"); + let bad = doc! { + "_id": bson::oid::ObjectId::new(), + "pk": "u1", + "sk_s": "s1", + }; + let repaired = repaired_document(&bad, &schema, &defs).expect("repair"); + assert_eq!(repaired, want); + } + #[test] fn repair_refuses_a_document_whose_sort_key_field_has_the_wrong_type() { let (schema, defs) = range_schema(ScalarAttributeType::N); diff --git a/crates/storage-mongodb/src/data_engine_live_tests.rs b/crates/storage-mongodb/src/data_engine_live_tests.rs index 1c8429bb3..1df0763cd 100644 --- a/crates/storage-mongodb/src/data_engine_live_tests.rs +++ b/crates/storage-mongodb/src/data_engine_live_tests.rs @@ -346,22 +346,31 @@ async fn live_creating_update_that_always_loses_on_id_answers_transaction_confli } /// Store a document the way the fast path did before the fix: an upsert whose -/// filter is the key and whose update sets one attribute. MongoDB copies the -/// filter's `pk` and `sk_*` fields, generates an ObjectId `_id`, and writes -/// no key attribute into `item_data`. -async fn store_without_key(engine: &MongoEngine, key_info: &TableKeyInfo, key: &Item, w: &str) { +/// filter is the key and whose update sets one attribute, or, when `w` is +/// `None`, removes one. MongoDB copies the filter's `pk` and `sk_*` fields, +/// generates an ObjectId `_id`, and writes no key attribute into `item_data`. +/// The `None` form is main's REMOVE-only creating update: `$unset` +/// contributes nothing to the inserted document, so it has no `item_data` at +/// all. +async fn store_without_key( + engine: &MongoEngine, + key_info: &TableKeyInfo, + key: &Item, + w: Option<&str>, +) { let coll = engine .data_db .collection::(&data_collection_name(&key_info.table_id)); let filter = pk_filter(key, &key_info.key_schema, &key_info.attribute_definitions).expect("filter"); - coll.find_one_and_update( - filter, - doc! { "$set": { "item_data.w": { "S": w } }, "$inc": { "_v": 1_i64 } }, - ) - .upsert(true) - .await - .expect("upsert"); + let update = match w { + Some(w) => doc! { "$set": { "item_data.w": { "S": w } }, "$inc": { "_v": 1_i64 } }, + None => doc! { "$unset": { "item_data.w": "" }, "$inc": { "_v": 1_i64 } }, + }; + coll.find_one_and_update(filter, update) + .upsert(true) + .await + .expect("upsert"); } /// The repair rewrites a document stored without its key into the document @@ -376,8 +385,8 @@ async fn live_repair_rewrites_documents_stored_without_their_key() { let table_id = format!("repair-{}", if range { "range" } else { "hash" }); let key_info = table(&engine, &table_id, range); let key = item_key(range); - store_without_key(&engine, &key_info, &key, "upd").await; - store_without_key(&engine, &key_info, &key, "upd2").await; + store_without_key(&engine, &key_info, &key, Some("upd")).await; + store_without_key(&engine, &key_info, &key, Some("upd2")).await; let bad = stored(&engine, &key_info, &key).await; assert!( bad.get_object_id("_id").is_ok(), @@ -466,6 +475,87 @@ async fn live_repair_rewrites_documents_stored_without_their_key() { finish(engine).await; } +/// A pre-fix document with no `item_data` at all, which is what main's fast +/// path stored for a REMOVE-only creating UpdateItem, reads as an internal +/// error before the repair and as the key-only item after it. The repaired +/// document is the one the fixed fast path writes for the same request, and a +/// PutItem on the key then succeeds. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_rewrites_a_document_stored_with_no_item_data() { + let engine = scratch_engine().await; + for range in [false, true] { + let table_id = format!("repair-noitem-{}", if range { "range" } else { "hash" }); + let key_info = table(&engine, &table_id, range); + let key = item_key(range); + store_without_key(&engine, &key_info, &key, None).await; + + let bad = stored(&engine, &key_info, &key).await; + assert!( + bad.get_object_id("_id").is_ok(), + "{table_id}: the pre-fix shape has an ObjectId _id" + ); + assert!( + !bad.contains_key("item_data"), + "{table_id}: a REMOVE-only creating update stored no item_data: {bad:?}" + ); + assert!( + engine.get_item_impl(&key_info, &key).await.is_err(), + "{table_id}: the key reads as an error until the repair runs" + ); + + let mut summary = KeyRepairSummary::default(); + engine + .repair_table( + &table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut summary, + ) + .await + .expect("repair"); + assert_eq!( + summary, + KeyRepairSummary { + tables_scanned: 0, + documents_repaired: 1, + documents_failed: 0, + }, + "{table_id}" + ); + + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&key, &key_info, 1), + "{table_id}: repaired into the key-only item" + ); + assert_eq!( + engine.get_item_impl(&key_info, &key).await.expect("get"), + Some(key.clone()), + "{table_id}: the key reads back" + ); + + let put = with(key.clone(), "w", s("put")); + let old = engine + .put_item_impl( + &key_info, + put.clone(), + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: put after repair: {e:?}")); + assert_eq!( + old, + Some(key.clone()), + "{table_id}: put returns the repaired key-only item" + ); + } + finish(engine).await; +} + /// The whole-catalog entry point finds the table through its catalog row. #[tokio::test] #[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] @@ -485,7 +575,7 @@ async fn live_repair_reads_the_key_schema_from_the_catalog() { .await .expect("catalog row"); let key = item_key(false); - store_without_key(&engine, &key_info, &key, "upd").await; + store_without_key(&engine, &key_info, &key, Some("upd")).await; let summary = engine .repair_items_stored_without_key() From 9972b26f2ee263778b93d4d22a91d2207675f431 Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 03/14] docs(mongodb): say what is_stored_without_key checks The comment described a check on both the _id and the item_data key attributes; the function tests only that _id is an ObjectId. State that, and why it is the whole test for a base-table document. --- crates/storage-mongodb/src/data/key_fields.rs | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/crates/storage-mongodb/src/data/key_fields.rs b/crates/storage-mongodb/src/data/key_fields.rs index d28c37c27..a6df7aaac 100644 --- a/crates/storage-mongodb/src/data/key_fields.rs +++ b/crates/storage-mongodb/src/data/key_fields.rs @@ -58,8 +58,13 @@ pub fn insert_only_key_fields( Ok(fields) } -/// Whether `doc` was stored by the upsert described in the module comment: an -/// `_id` that is not the key text and key attributes missing from `item_data`. +/// Whether `doc`'s `_id` is an ObjectId. +/// +/// That is the mark of the upsert described in the module comment, and the +/// whole test: `item_to_document` writes `_id` as the key text for every +/// base-table document, so a base-table document whose `_id` is an ObjectId +/// was generated by MongoDB on an insert that supplied no `_id`, which is +/// the pre-fix upsert. The `item_data` contents are not examined. #[must_use] pub fn is_stored_without_key(doc: &Document) -> bool { matches!(doc.get("_id"), Some(Bson::ObjectId(_))) From 16efffc1d0a2aaca75c5b8f91b082c13c45ba84b Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 04/14] fix(mongodb): run the key repair in the background, in batches The repair held the listener closed: serve awaited the whole pass in server_components_factory, and the pass was one transaction per document. Measured on this host, 10,000 documents cost 25.64 s before anything answered on the port, and nothing reported progress while it ran. Run the pass as a background task, so the listener opens whatever the number of affected documents. A key the pass has not reached yet answers as it did before the repair existed, so a long pass delays the repair of that key rather than the start of the server. Rebuild up to 500 documents per transaction. A commit is one majority-acknowledged round trip and dominates the cost of a rebuild, so one commit per batch is what sets the rate: each batch is one read, one delete and one insert. The batch is atomic, so a process killed mid-pass still leaves every document either as it was or rebuilt. When a batch transaction cannot run, its documents are rebuilt one transaction at a time, so one document that cannot be written leaves the rest repaired. The batch reads its documents inside the transaction instead of taking them from the scan. A client write that lands between the scan and the rebuild is then either carried into the rebuilt document or turns the delete into a write conflict, which runs the attempt again. Progress is logged at INFO every 10,000 documents and once at the end with the elapsed time and the rate. A live test repairs 1,100 documents, which crosses two batch boundaries and ends on a partial batch. --- .../src/data_engine_key_repair.rs | 339 ++++++++++++++++-- .../src/data_engine_live_tests.rs | 91 +++++ crates/storage-mongodb/src/lib.rs | 23 +- 3 files changed, 415 insertions(+), 38 deletions(-) diff --git a/crates/storage-mongodb/src/data_engine_key_repair.rs b/crates/storage-mongodb/src/data_engine_key_repair.rs index ebc2cbeef..c5ada64f2 100644 --- a/crates/storage-mongodb/src/data_engine_key_repair.rs +++ b/crates/storage-mongodb/src/data_engine_key_repair.rs @@ -7,17 +7,34 @@ //! an `item_data` without the key attributes. Such a document answers //! GetItem without the key, and every PutItem on the key fails on the //! immutable `_id`. `_id` cannot be changed in place, so the repair replaces -//! each such document with the one `item_to_document` writes, in one -//! transaction: delete the old document by its ObjectId, insert the rebuilt -//! one. The scan is a bracket query on the `_id` index -//! (`{_id: {$type: "objectId"}}`), so a collection with no such document costs -//! one index probe. +//! each such document with the one `item_to_document` writes: delete the old +//! document by its ObjectId, insert the rebuilt one. The scan is a bracket +//! query on the `_id` index (`{_id: {$type: "objectId"}}`), so a collection +//! with no such document costs one index probe. //! -//! `serve` runs this once at startup, before the server accepts requests -//! (see `server_components_factory` in `lib.rs`). A failure on one document -//! is logged and skipped; the summary reports it. +//! Up to [`REPAIR_BATCH`] documents are rebuilt in one transaction. A +//! transaction's commit is one majority-acknowledged round trip and dominates +//! the cost of a rebuild, so batching the rebuilds behind one commit is what +//! sets the rate; each batch is one in-transaction read, one delete and one +//! insert. The batch is atomic, so a process killed mid-repair leaves every +//! document either as it was or rebuilt, and the next start finishes the +//! remainder. +//! +//! The documents are re-read inside the transaction rather than taken from +//! the scan, so a client that writes to a document between the scan and the +//! rebuild does not lose that write: the rebuild either sees it or hits a +//! write conflict and runs again. +//! +//! [`spawn_key_repair`] runs the whole pass as a background task, so `serve` +//! opens its listener without waiting for it (see `server_components_factory` +//! in `lib.rs`). Until the repair reaches a given document, requests on that +//! document's key answer as they did before this repair existed. Progress is +//! logged every [`PROGRESS_EVERY`] documents and once at the end. A failure +//! on one document is logged and skipped; the summary reports it. -use bson::{Document, doc}; +use std::sync::Arc; + +use bson::{Bson, Document, doc}; use extenddb_core::types::{AttributeDefinition, KeySchemaElement}; use extenddb_storage::error::StorageError; use futures::TryStreamExt; @@ -27,6 +44,12 @@ use crate::MongoEngine; use crate::data::data_collection_name; use crate::data::key_fields::{is_stored_without_key, repaired_document}; +/// Documents rebuilt in one transaction. +const REPAIR_BATCH: usize = 500; + +/// Repaired documents between progress lines. +const PROGRESS_EVERY: u64 = 10_000; + /// What a repair run did. #[derive(Debug, Default, Clone, PartialEq, Eq)] pub struct KeyRepairSummary { @@ -100,6 +123,10 @@ impl MongoEngine { .find(doc! { "_id": { "$type": "objectId" } }) .await .map_err(|e| StorageError::Internal(e.to_string()))?; + // The rebuilt documents have the key text as their `_id`, outside the + // ObjectId bracket this cursor reads, so inserting them cannot feed + // the scan its own output. + let mut ids: Vec = Vec::with_capacity(REPAIR_BATCH); while let Some(bad) = cursor .try_next() .await @@ -108,30 +135,155 @@ impl MongoEngine { if !is_stored_without_key(&bad) { continue; } - let repaired = match repaired_document(&bad, key_schema, attribute_definitions) { - Ok(d) => d, - Err(e) => { + if let Some(id) = bad.get("_id") { + ids.push(id.clone()); + } + if ids.len() >= REPAIR_BATCH { + self.repair_batch(&coll, table_id, key_schema, attribute_definitions, &ids, summary) + .await; + ids.clear(); + } + } + if !ids.is_empty() { + self.repair_batch(&coll, table_id, key_schema, attribute_definitions, &ids, summary) + .await; + } + Ok(()) + } + + /// Rebuild one batch, and when the batch transaction cannot run, rebuild + /// its documents one transaction at a time so that one document that + /// cannot be written leaves the rest of the batch repaired. + async fn repair_batch( + &self, + coll: &mongodb::Collection, + table_id: &str, + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], + ids: &[Bson], + summary: &mut KeyRepairSummary, + ) { + let before = summary.documents_repaired; + match self + .replace_batch(coll, ids, key_schema, attribute_definitions) + .await + { + Ok(outcome) => { + summary.documents_repaired += outcome.repaired; + summary.documents_failed += outcome.failed.len() as u64; + for (id, reason) in outcome.failed { tracing::error!( - "key repair: table {table_id}: cannot rebuild document {:?}: {e}", - bad.get("_id") + "key repair: table {table_id}: cannot rebuild document {id:?}: {reason}" ); - summary.documents_failed += 1; - continue; } - }; - match self.replace_document(&coll, &bad, repaired).await { - Ok(true) => summary.documents_repaired += 1, - Ok(false) => {} + } + Err(e) => { + tracing::warn!( + "key repair: table {table_id}: a batch of {} documents could not be \ + rewritten together ({e}); rewriting them one at a time", + ids.len() + ); + for id in ids { + self.repair_one(coll, table_id, id, key_schema, attribute_definitions, summary) + .await; + } + } + } + log_progress(before, summary); + } + + /// Rebuild the document at `id` in its own transaction. + async fn repair_one( + &self, + coll: &mongodb::Collection, + table_id: &str, + id: &Bson, + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], + summary: &mut KeyRepairSummary, + ) { + let bad = match coll.find_one(doc! { "_id": id.clone() }).await { + Ok(Some(bad)) => bad, + // Gone, so another process repaired or deleted it. + Ok(None) => return, + Err(e) => { + tracing::error!("key repair: table {table_id}: cannot read document {id:?}: {e}"); + summary.documents_failed += 1; + return; + } + }; + if !is_stored_without_key(&bad) { + return; + } + let repaired = match repaired_document(&bad, key_schema, attribute_definitions) { + Ok(d) => d, + Err(e) => { + tracing::error!( + "key repair: table {table_id}: cannot rebuild document {id:?}: {e}" + ); + summary.documents_failed += 1; + return; + } + }; + match self.replace_document(coll, &bad, repaired).await { + Ok(true) => summary.documents_repaired += 1, + Ok(false) => {} + Err(e) => { + tracing::error!( + "key repair: table {table_id}: cannot rewrite document {id:?}: {e}" + ); + summary.documents_failed += 1; + } + } + } + + /// Rebuild the documents at `ids` in one transaction: read them in the + /// transaction, delete them, insert the rebuilt documents. Ids whose + /// document is gone or already carries its key are skipped. Ids whose + /// document cannot be rebuilt are reported in `failed` and left in place. + async fn replace_batch( + &self, + coll: &mongodb::Collection, + ids: &[Bson], + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], + ) -> Result { + let mut session = self + .client + .start_session() + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + let tx_options = self.write_transaction_options(); + for attempt in 0..TRANSIENT_RETRY_ATTEMPTS { + session + .start_transaction() + .with_options(tx_options.clone()) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + match rebuild_in_session(coll, &mut session, ids, key_schema, attribute_definitions) + .await + { + Ok(outcome) if outcome.repaired == 0 => { + let _ = session.abort_transaction().await; + return Ok(outcome); + } + Ok(outcome) => match self.commit_write_transaction(&mut session, None).await? { + CommitOutcome::Committed => return Ok(outcome), + CommitOutcome::RerunBody => backoff_sleep(attempt).await, + }, + Err(BatchError::Mongo(e)) if is_transient_write_conflict(&e) => { + let _ = session.abort_transaction().await; + backoff_sleep(attempt).await; + } Err(e) => { - tracing::error!( - "key repair: table {table_id}: cannot rewrite document {:?}: {e}", - bad.get("_id") - ); - summary.documents_failed += 1; + let _ = session.abort_transaction().await; + return Err(StorageError::Internal(e.to_string())); } } } - Ok(()) + Err(StorageError::TransactionConflict( + "key repair: too many concurrent write conflicts, giving up".to_owned(), + )) } /// Delete `bad` by its `_id` and insert `repaired` in one transaction. @@ -220,3 +372,138 @@ fn table_key( })?; Ok((table_id, key_schema, attribute_definitions)) } + +/// What one batch transaction did. +#[derive(Debug, Default)] +struct BatchOutcome { + /// Documents the transaction rewrote. + repaired: u64, + /// Documents it left in place, with the reason. + failed: Vec<(Bson, String)>, +} + +/// Why a batch transaction could not run. +#[derive(Debug)] +enum BatchError { + /// A MongoDB error, which may be a transient write conflict. + Mongo(mongodb::error::Error), + /// The delete did not match the documents the transaction had just read, + /// so the batch cannot be inserted without risking a duplicate item. A + /// transaction holds its snapshot, so this does not happen; the caller + /// falls back to one transaction per document rather than trusting it. + Counts { deleted: u64, read: usize }, +} + +impl From for BatchError { + fn from(e: mongodb::error::Error) -> Self { + Self::Mongo(e) + } +} + +impl std::fmt::Display for BatchError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Mongo(e) => write!(f, "{e}"), + Self::Counts { deleted, read } => { + write!(f, "deleted {deleted} of {read} documents read in the transaction") + } + } + } +} + +/// Read the documents at `ids` in `session`'s transaction, delete them and +/// insert the rebuilt documents. The read is inside the transaction, so a +/// write that landed after the scan is either rebuilt with the document or +/// turns the delete into a write conflict. +async fn rebuild_in_session( + coll: &mongodb::Collection, + session: &mut mongodb::ClientSession, + ids: &[Bson], + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], +) -> Result { + let mut outcome = BatchOutcome::default(); + let mut stale: Vec = Vec::with_capacity(ids.len()); + let mut rebuilt: Vec = Vec::with_capacity(ids.len()); + let mut cursor = coll + .find(doc! { "_id": { "$in": ids.to_vec() } }) + .session(&mut *session) + .await?; + while let Some(bad) = cursor.next(&mut *session).await { + let bad = bad?; + let Some(id) = bad.get("_id").cloned() else { + continue; + }; + if !is_stored_without_key(&bad) { + continue; + } + match repaired_document(&bad, key_schema, attribute_definitions) { + Ok(d) => { + stale.push(id); + rebuilt.push(d); + } + Err(e) => outcome.failed.push((id, e.to_string())), + } + } + if rebuilt.is_empty() { + return Ok(outcome); + } + let deleted = coll + .delete_many(doc! { "_id": { "$in": stale } }) + .session(&mut *session) + .await?; + if deleted.deleted_count != rebuilt.len() as u64 { + return Err(BatchError::Counts { + deleted: deleted.deleted_count, + read: rebuilt.len(), + }); + } + coll.insert_many(&rebuilt).session(session).await?; + outcome.repaired = rebuilt.len() as u64; + Ok(outcome) +} + +/// Log a progress line when the repaired count has crossed a multiple of +/// `PROGRESS_EVERY`. A batch is smaller than that interval, so each crossing +/// is logged once. +fn log_progress(before: u64, summary: &KeyRepairSummary) { + if before / PROGRESS_EVERY != summary.documents_repaired / PROGRESS_EVERY { + tracing::info!("key repair in progress: {summary}"); + } +} + +/// Run a whole repair pass in the background and log the outcome. +/// +/// `serve` starts this instead of awaiting the pass, so the listener opens at +/// once however many documents need rebuilding. Requests on a key whose +/// document the pass has not reached yet answer as they did before the repair +/// existed, so a long pass delays the repair of a key rather than the start +/// of the server. +pub(crate) fn spawn_key_repair(engine: Arc) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + let started = std::time::Instant::now(); + match engine.repair_items_stored_without_key().await { + Ok(repair) if repair.is_clean() => { + tracing::debug!("startup key repair: {repair}"); + } + Ok(repair) => { + let seconds = started.elapsed().as_secs_f64(); + let rate = if seconds > 0.0 { + (repair.documents_repaired as f64) / seconds + } else { + 0.0 + }; + let line = format!( + "startup key repair finished: {repair}, {seconds:.2} s, \ + {rate:.0} documents a second" + ); + if repair.documents_failed == 0 { + tracing::info!("{line}"); + } else { + tracing::warn!("{line}"); + } + } + Err(e) => tracing::error!("Failed to repair items stored without their key: {e}"), + } + }) +} diff --git a/crates/storage-mongodb/src/data_engine_live_tests.rs b/crates/storage-mongodb/src/data_engine_live_tests.rs index 1df0763cd..f15761132 100644 --- a/crates/storage-mongodb/src/data_engine_live_tests.rs +++ b/crates/storage-mongodb/src/data_engine_live_tests.rs @@ -556,6 +556,97 @@ async fn live_repair_rewrites_a_document_stored_with_no_item_data() { finish(engine).await; } +/// Insert `count` documents in the pre-fix shape directly, one per key +/// `u0..u`, with no key attribute in `item_data`. The upsert path in +/// `store_without_key` is one round trip per document, which is too slow for +/// the counts the batched repair needs. +async fn store_many_without_key(engine: &MongoEngine, key_info: &TableKeyInfo, count: usize) { + let coll = engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)); + let docs: Vec = (0..count) + .map(|i| { + doc! { + "_id": bson::oid::ObjectId::new(), + "pk": format!("u{i}"), + "item_data": { "w": { "S": format!("value-{i}") } }, + "_v": 1_i64, + } + }) + .collect(); + coll.insert_many(docs).await.expect("plant"); +} + +/// A repair pass that spans more than one batch rebuilds every document. The +/// batch is 500 documents, so 1,100 crosses two boundaries and ends with a +/// partial batch. Every key reads back with its value, none is duplicated, +/// and a second pass finds nothing. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_spans_more_than_one_batch() { + const COUNT: usize = 1_100; + let engine = scratch_engine().await; + let key_info = table(&engine, "repair-batches", false); + store_many_without_key(&engine, &key_info, COUNT).await; + + let mut summary = KeyRepairSummary::default(); + engine + .repair_table( + &key_info.table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut summary, + ) + .await + .expect("repair"); + assert_eq!( + summary, + KeyRepairSummary { + tables_scanned: 0, + documents_repaired: COUNT as u64, + documents_failed: 0, + } + ); + + let coll = engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)); + assert_eq!( + coll.count_documents(doc! {}).await.expect("count"), + COUNT as u64, + "no document was duplicated or lost" + ); + assert_eq!( + coll.count_documents(doc! { "_id": { "$type": "objectId" } }) + .await + .expect("count"), + 0, + "every document carries its key text as _id" + ); + for i in [0_usize, 499, 500, 999, 1_000, COUNT - 1] { + let mut key = Item::new(); + key.insert("pk".to_owned(), s(&format!("u{i}"))); + assert_eq!( + engine.get_item_impl(&key_info, &key).await.expect("get"), + Some(with(key.clone(), "w", s(&format!("value-{i}")))), + "key u{i} reads back" + ); + } + + let mut again = KeyRepairSummary::default(); + engine + .repair_table( + &key_info.table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut again, + ) + .await + .expect("second repair"); + assert!(again.is_clean(), "nothing left to repair: {again}"); + finish(engine).await; +} + /// The whole-catalog entry point finds the table through its catalog row. #[tokio::test] #[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] diff --git a/crates/storage-mongodb/src/lib.rs b/crates/storage-mongodb/src/lib.rs index 8c1de6c6c..37d9ed246 100644 --- a/crates/storage-mongodb/src/lib.rs +++ b/crates/storage-mongodb/src/lib.rs @@ -137,20 +137,19 @@ fn server_components_factory( Err(e) => tracing::error!("Failed to sweep stale backup claims: {e}"), } - // An UpdateItem on a build before the fix for issue #436 could create - // an item whose document has no key. Rewrite any such document before - // serving, so that reads return the key and PutItem on it succeeds - // (see `data_engine::key_repair`). Each rewrite is its own - // transaction, so running this while another process serves the same - // data is safe. - match engine.repair_items_stored_without_key().await { - Ok(repair) if repair.is_clean() => {} - Ok(repair) => tracing::warn!("startup key repair: {repair}"), - Err(e) => tracing::error!("Failed to repair items stored without their key: {e}"), - } - let engine = Arc::new(engine); + // An UpdateItem on a build before the fix for issue #436 could create + // an item whose document has no key. Rewrite any such document, so + // that reads return the key and PutItem on it succeeds (see + // `data_engine::key_repair`). This runs in the background: the pass + // costs time in proportion to the number of affected documents, and + // the listener must not wait for it. A key the pass has not reached + // yet answers as it did before the repair existed, and each batch is + // its own transaction, so running this while another process serves + // the same data is safe. + data_engine::key_repair::spawn_key_repair(engine.clone()); + // Create one shared catalog/management client. MongoDB Client clones // share the underlying pool, so max_catalog_connections limits the // combined catalog and auth traffic rather than creating two From 816df478e988a6a4d033c1b1ba7949c0288e8d5c Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 05/14] test(mongodb): run the live tests of the key repair in CI The repair driver had no automated coverage. The scan, the batch transaction and the behaviour of two processes starting together are tested only by the live tests in crates/storage-mongodb, which are EXTENDDB_TEST_MONGODB_URI or passed --ignored. devtools/run-mongodb-tests already starts the replica set those tests need, so it now exports its URI and takes a --rust-live flag that runs them after the delegated suite. The MongoDB rust CI job passes that flag. Each live test works in its own pair of databases and drops them at the end, so it does not touch the served instance's catalog or data. --- .github/workflows/integration-mongodb.yml | 2 +- devtools/run-mongodb-tests | 31 ++++++++++++++++++++++- 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/.github/workflows/integration-mongodb.yml b/.github/workflows/integration-mongodb.yml index 03223f967..b35233402 100644 --- a/.github/workflows/integration-mongodb.yml +++ b/.github/workflows/integration-mongodb.yml @@ -117,7 +117,7 @@ jobs: - name: Build release (mongodb backend) run: cargo build --release --no-default-features --features mongodb-test-hooks - name: Run MongoDB rust integration suite - run: devtools/run-mongodb-tests -- --rust --rust-integration + run: devtools/run-mongodb-tests --rust-live -- --rust --rust-integration mongodb-production-build: runs-on: extenddb_ubuntu-2404_4-core diff --git a/devtools/run-mongodb-tests b/devtools/run-mongodb-tests index 2a287798a..4a3b03885 100755 --- a/devtools/run-mongodb-tests +++ b/devtools/run-mongodb-tests @@ -23,6 +23,10 @@ # For tests that talk to MongoDB directly (the storage # crate's snapshot tests). Tear down with # docker rm -f extenddb-runtests-mongo-. +# --rust-live Also run the Rust live tests of the storage-mongodb +# crate against this runner's replica set. They are +# `#[ignore]`d, because they need a replica set, and +# they are the only tests of the startup key repair. # -h, --help Show this help # # Arguments after `--` are passed to `devtools/run-tests` verbatim. @@ -32,6 +36,7 @@ # devtools/run-mongodb-tests # devtools/run-mongodb-tests -- --pytest --filter test_put_item # devtools/run-mongodb-tests --keep -- --pytest --filter test_rfc0003 +# devtools/run-mongodb-tests --rust-live -- --rust --rust-integration # # Prerequisites: # - Docker running @@ -54,6 +59,7 @@ MONGO_PORT=27021 OUTPUT_DIR="" KEEP=false MONGO_ONLY=false +RUST_LIVE=false RUN_TESTS_ARGS=() while [[ $# -gt 0 ]]; do @@ -63,8 +69,9 @@ while [[ $# -gt 0 ]]; do --output) OUTPUT_DIR="$2"; shift 2 ;; --keep) KEEP=true; shift ;; --mongo-only) MONGO_ONLY=true; KEEP=true; shift ;; + --rust-live) RUST_LIVE=true; shift ;; --) shift; RUN_TESTS_ARGS=("$@"); break ;; - -h|--help) sed -n '5,40p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;; + -h|--help) sed -n '5,45p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;; *) echo "error: unknown option: $1" >&2; exit 1 ;; esac done @@ -317,7 +324,29 @@ export EXTENDDB_CONFIG="$CONFIG" # tempfiles land in an allowed location on both Linux and macOS. export TMPDIR="$TMP_CANON" +# The Rust live tests in crates/storage-mongodb connect to the replica set +# directly, rather than through the server, so they read its URI from here. +export EXTENDDB_TEST_MONGODB_URI="mongodb://localhost:$MONGO_PORT/?replicaSet=rs0&directConnection=true" + RC=0 devtools/run-tests --extenddb --backend mongodb --config "$CONFIG" "${RUN_TESTS_ARGS[@]}" || RC=$? +# The live tests of the storage-mongodb crate are `#[ignore]`d, so +# `cargo test` does not run them: they need the replica set this runner +# started. They are the only tests of the startup key repair, which is why +# they run here rather than in the plain `cargo test --workspace` job. Each +# one works in its own pair of databases and drops them at the end, so they +# do not touch the served instance's catalog or data. +if $RUST_LIVE; then + echo "" + echo "=== MongoDB live Rust tests (storage-mongodb, --ignored) ===" + if cargo test -p extenddb-storage-mongodb --lib -- \ + --ignored --test-threads=1 data_engine::live_tests::; then + echo " ✓ MongoDB live Rust tests passed" + else + echo " ✗ MongoDB live Rust tests failed" + RC=1 + fi +fi + exit $RC From c71dc451a74a4c01d293a4aacd1f6dab585f39da Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 06/14] fix(mongodb): wait out a writer that holds a document being repaired Measuring the batched repair after a SIGKILL exposed a regression against the one-transaction-per-document form. The killed process leaves its last transaction holding that batch's documents until the server expires it, which takes up to transactionLifetimeLimitSeconds, 60 s by default. The next start found 500 locked documents, spent the transient retry budget of about 0.3 s on them, and reported 114 of 200,000 as failed, so those keys stayed broken until a third start. Measured before this change: "152386 documents stored without their key repaired, 114 failed". A batch that cannot run is already rebuilt one document per transaction. Each of those now waits for a conflicting writer, 40 waits of 2 s, which covers the lock expiry with margin. Nothing waits on the repair, so waiting costs the start nothing and finishes the pass instead of leaving documents for the next one. A document rewritten after a wait is logged with how long it waited. Measured after this change on the same scenario: "129000 documents stored without their key repaired, 0 failed", one document logged as rewritten after waiting 62 s. The single-document rebuild now goes through the batch path with one id rather than its own delete-and-insert, so it reads the document inside its transaction and cannot overwrite a client's write with a stale copy. That makes the old replace_document redundant. --- .../src/data_engine_key_repair.rs | 145 ++++++------------ 1 file changed, 50 insertions(+), 95 deletions(-) diff --git a/crates/storage-mongodb/src/data_engine_key_repair.rs b/crates/storage-mongodb/src/data_engine_key_repair.rs index c5ada64f2..370ed634f 100644 --- a/crates/storage-mongodb/src/data_engine_key_repair.rs +++ b/crates/storage-mongodb/src/data_engine_key_repair.rs @@ -18,7 +18,9 @@ //! sets the rate; each batch is one in-transaction read, one delete and one //! insert. The batch is atomic, so a process killed mid-repair leaves every //! document either as it was or rebuilt, and the next start finishes the -//! remainder. +//! remainder. A batch whose transaction cannot run is rebuilt one document +//! per transaction, so one document another writer holds, or one that cannot +//! be rebuilt at all, leaves the rest of the batch repaired. //! //! The documents are re-read inside the transaction rather than taken from //! the scan, so a client that writes to a document between the scan and the @@ -50,6 +52,17 @@ const REPAIR_BATCH: usize = 500; /// Repaired documents between progress lines. const PROGRESS_EVERY: u64 = 10_000; +/// How many times a single-document rebuild waits for a writer that holds its +/// document, and how long each wait is. +/// +/// A process killed mid-repair leaves its last transaction holding that +/// batch's documents until the server expires it, which takes up to +/// `transactionLifetimeLimitSeconds` (60 s by default). The repair runs in the +/// background and nothing waits on it, so it waits the locks out instead of +/// counting those documents failed and leaving them for the next start. +const CONFLICT_WAITS: u32 = 40; +const CONFLICT_WAIT: std::time::Duration = std::time::Duration::from_secs(2); + /// What a repair run did. #[derive(Debug, Default, Clone, PartialEq, Eq)] pub struct KeyRepairSummary { @@ -192,7 +205,8 @@ impl MongoEngine { log_progress(before, summary); } - /// Rebuild the document at `id` in its own transaction. + /// Rebuild the document at `id` in its own transaction, waiting out a + /// writer that holds it rather than leaving it for the next start. async fn repair_one( &self, coll: &mongodb::Collection, @@ -202,37 +216,40 @@ impl MongoEngine { attribute_definitions: &[AttributeDefinition], summary: &mut KeyRepairSummary, ) { - let bad = match coll.find_one(doc! { "_id": id.clone() }).await { - Ok(Some(bad)) => bad, - // Gone, so another process repaired or deleted it. - Ok(None) => return, - Err(e) => { - tracing::error!("key repair: table {table_id}: cannot read document {id:?}: {e}"); - summary.documents_failed += 1; - return; - } - }; - if !is_stored_without_key(&bad) { - return; - } - let repaired = match repaired_document(&bad, key_schema, attribute_definitions) { - Ok(d) => d, - Err(e) => { - tracing::error!( - "key repair: table {table_id}: cannot rebuild document {id:?}: {e}" - ); - summary.documents_failed += 1; - return; - } - }; - match self.replace_document(coll, &bad, repaired).await { - Ok(true) => summary.documents_repaired += 1, - Ok(false) => {} - Err(e) => { - tracing::error!( - "key repair: table {table_id}: cannot rewrite document {id:?}: {e}" - ); - summary.documents_failed += 1; + let one = std::slice::from_ref(id); + for waited in 0..=CONFLICT_WAITS { + match self + .replace_batch(coll, one, key_schema, attribute_definitions) + .await + { + Ok(outcome) => { + if waited > 0 && outcome.repaired > 0 { + tracing::info!( + "key repair: table {table_id}: document {id:?} rewritten after \ + waiting {} s for the writer holding it", + waited * CONFLICT_WAIT.as_secs() as u32 + ); + } + summary.documents_repaired += outcome.repaired; + summary.documents_failed += outcome.failed.len() as u64; + for (failed_id, reason) in outcome.failed { + tracing::error!( + "key repair: table {table_id}: cannot rebuild document \ + {failed_id:?}: {reason}" + ); + } + return; + } + Err(StorageError::TransactionConflict(_)) if waited < CONFLICT_WAITS => { + tokio::time::sleep(CONFLICT_WAIT).await; + } + Err(e) => { + tracing::error!( + "key repair: table {table_id}: cannot rewrite document {id:?}: {e}" + ); + summary.documents_failed += 1; + return; + } } } } @@ -286,68 +303,6 @@ impl MongoEngine { )) } - /// Delete `bad` by its `_id` and insert `repaired` in one transaction. - /// Returns false when `bad` was gone already (another process repaired or - /// deleted it), in which case nothing is inserted. - async fn replace_document( - &self, - coll: &mongodb::Collection, - bad: &Document, - repaired: Document, - ) -> Result { - let bad_id = bad - .get("_id") - .cloned() - .ok_or_else(|| StorageError::Internal("document has no _id".to_owned()))?; - let mut session = self - .client - .start_session() - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - let tx_options = self.write_transaction_options(); - for attempt in 0..TRANSIENT_RETRY_ATTEMPTS { - session - .start_transaction() - .with_options(tx_options.clone()) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; - let body = async { - let deleted = coll - .delete_one(doc! { "_id": bad_id.clone() }) - .session(&mut session) - .await?; - if deleted.deleted_count == 0 { - return Ok::(false); - } - coll.insert_one(repaired.clone()) - .session(&mut session) - .await?; - Ok(true) - } - .await; - match body { - Ok(false) => { - let _ = session.abort_transaction().await; - return Ok(false); - } - Ok(true) => match self.commit_write_transaction(&mut session, None).await? { - CommitOutcome::Committed => return Ok(true), - CommitOutcome::RerunBody => backoff_sleep(attempt).await, - }, - Err(e) if is_transient_write_conflict(&e) => { - let _ = session.abort_transaction().await; - backoff_sleep(attempt).await; - } - Err(e) => { - let _ = session.abort_transaction().await; - return Err(StorageError::Internal(e.to_string())); - } - } - } - Err(StorageError::TransactionConflict( - "key repair: too many concurrent write conflicts, giving up".to_owned(), - )) - } } /// `table_id`, key schema and attribute definitions of a catalog `tables` row. From c034aa278bb9eeb0839a2e36a28917e4f40f79e0 Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 07/14] style(mongodb): format the key repair module rustfmt 1.97.0 splits three calls in data_engine_key_repair.rs; no other change. --- .../src/data_engine_key_repair.rs | 39 +++++++++++++++---- 1 file changed, 31 insertions(+), 8 deletions(-) diff --git a/crates/storage-mongodb/src/data_engine_key_repair.rs b/crates/storage-mongodb/src/data_engine_key_repair.rs index 370ed634f..c5c9e714f 100644 --- a/crates/storage-mongodb/src/data_engine_key_repair.rs +++ b/crates/storage-mongodb/src/data_engine_key_repair.rs @@ -152,14 +152,28 @@ impl MongoEngine { ids.push(id.clone()); } if ids.len() >= REPAIR_BATCH { - self.repair_batch(&coll, table_id, key_schema, attribute_definitions, &ids, summary) - .await; + self.repair_batch( + &coll, + table_id, + key_schema, + attribute_definitions, + &ids, + summary, + ) + .await; ids.clear(); } } if !ids.is_empty() { - self.repair_batch(&coll, table_id, key_schema, attribute_definitions, &ids, summary) - .await; + self.repair_batch( + &coll, + table_id, + key_schema, + attribute_definitions, + &ids, + summary, + ) + .await; } Ok(()) } @@ -197,8 +211,15 @@ impl MongoEngine { ids.len() ); for id in ids { - self.repair_one(coll, table_id, id, key_schema, attribute_definitions, summary) - .await; + self.repair_one( + coll, + table_id, + id, + key_schema, + attribute_definitions, + summary, + ) + .await; } } } @@ -302,7 +323,6 @@ impl MongoEngine { "key repair: too many concurrent write conflicts, giving up".to_owned(), )) } - } /// `table_id`, key schema and attribute definitions of a catalog `tables` row. @@ -360,7 +380,10 @@ impl std::fmt::Display for BatchError { match self { Self::Mongo(e) => write!(f, "{e}"), Self::Counts { deleted, read } => { - write!(f, "deleted {deleted} of {read} documents read in the transaction") + write!( + f, + "deleted {deleted} of {read} documents read in the transaction" + ) } } } From 30537567c791bb2b4248be2918b51af3072eab8a Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 08/14] fix(mongodb): keep dotted key names off the native upsert The native UpdateItem upsert writes each key attribute with a $setOnInsert on the path item_data.. MongoDB reads a dot in that path as a step into a sub-document, so on a table with a key attribute named a.b the creating call stored item_data { a: { b: ... } } where PutItem stores the field "a.b". The call answered 500 and GetItem could not read the document back. The same path also conflicts with an update of an attribute named a, so SET a on an existing item of such a table answered 500. A table with a key attribute name that contains a dot or starts with $ now takes the session path, which writes the document item_to_document builds. That is the test native_attribute_name already applies to the updated attribute, and both gates now call one predicate. The live tests create an item on a table with hash key a.b and on one with range key s.k and compare the stored document with the one item_to_document builds. They then read it with GetItem, Query and Scan, set an attribute named a (or s) on it, and replace and delete it. A third live test sets that attribute on an item PutItem stored. The wire tests run the create test on tables with those key names and add a Query and a Scan over them. A unit test pins the predicate for the dot and the leading $. --- crates/storage-mongodb/src/data_engine.rs | 8 +- .../src/data_engine_live_tests.rs | 3 + .../src/data_engine_live_tests_key_names.rs | 266 ++++++++++++++++++ .../src/data_engine_update_create.rs | 61 +++- tests/test_update_item_creates_key.py | 60 +++- 5 files changed, 391 insertions(+), 7 deletions(-) create mode 100644 crates/storage-mongodb/src/data_engine_live_tests_key_names.rs diff --git a/crates/storage-mongodb/src/data_engine.rs b/crates/storage-mongodb/src/data_engine.rs index 2183dc7d2..c4d3b4f8c 100644 --- a/crates/storage-mongodb/src/data_engine.rs +++ b/crates/storage-mongodb/src/data_engine.rs @@ -47,7 +47,7 @@ fn native_attribute_name(path: &[PathElement], maps: &ExpressionMaps) -> Option< return None; }; let attr_name = resolve_name_ref(raw_name, maps).ok()?; - (!attr_name.contains('.') && !attr_name.starts_with('$')).then(|| attr_name.into_owned()) + update_create::is_plain_path_name(&attr_name).then(|| attr_name.into_owned()) } #[derive(Clone, Copy)] @@ -800,11 +800,15 @@ impl MongoEngine { // indexed attribute is $set to a new value or $unset (RFC-0003 // §2.2). The GSI cache lets us avoid a catalog query in the // common (no-GSI) case; when the cache is stale-or-unknown - // we fall through to the slow path which is authoritative. + // we fall through to the slow path which is authoritative. Key + // attribute names that are not plain path names (a dot, a leading + // `$`) also take the slow path, because the upsert writes the key + // attributes as `item_data.` (see `update_create`). if condition.is_none() && !return_old && stream.is_none() && self.gsi_cache_get_fresh(&key_info.table_id) == Some(false) + && update_create::key_names_are_plain_paths(&key_info.key_schema) && let Some(mongo_update) = Self::try_build_native_update(actions, maps) { let (fast_filter, took_fast) = match mongo_update { diff --git a/crates/storage-mongodb/src/data_engine_live_tests.rs b/crates/storage-mongodb/src/data_engine_live_tests.rs index f15761132..df3cc184a 100644 --- a/crates/storage-mongodb/src/data_engine_live_tests.rs +++ b/crates/storage-mongodb/src/data_engine_live_tests.rs @@ -24,6 +24,9 @@ use super::key_repair::KeyRepairSummary; use crate::MongoEngine; use crate::data::{data_collection_name, document_to_item, item_to_document, pk_filter}; +#[path = "data_engine_live_tests_key_names.rs"] +mod key_names; + fn uri() -> String { std::env::var("EXTENDDB_TEST_MONGODB_URI").expect( "set EXTENDDB_TEST_MONGODB_URI to a MongoDB replica set, e.g. \ diff --git a/crates/storage-mongodb/src/data_engine_live_tests_key_names.rs b/crates/storage-mongodb/src/data_engine_live_tests_key_names.rs new file mode 100644 index 000000000..ebf058fd6 --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_live_tests_key_names.rs @@ -0,0 +1,266 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for tables whose key attribute names are not plain +//! field names. +//! +//! The native UpdateItem upsert writes each key attribute as the path +//! `item_data.`, and MongoDB reads a dot in a path as a step into a +//! sub-document. A table with a key attribute named `a.b` or `s.k` must take +//! the session path instead, which stores the field `"a.b"` the way PutItem +//! does (see `data_engine_update_create.rs`). +//! +//! The tests read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are +//! `#[ignore]`d, like the tests in the parent module. + +use std::collections::HashMap; + +use extenddb_core::expression::{Expr, ExpressionMaps, KeyCondition, UpdateAction}; +use extenddb_core::types::{AttributeValue, Item, KeyType, ScalarAttributeType, TableKeyInfo}; + +use super::{def, expected_document, finish, key, maps, path, s, scratch_engine, stored, with}; +use crate::MongoEngine; + +/// Key info for a table with a string hash key named `hash` and, when given, +/// a string range key named `range`. The GSI cache is set to "no indexes", so +/// an unconditional UpdateItem would take the native fast path if the key +/// names allowed it. +fn named_table( + engine: &MongoEngine, + table_id: &str, + hash: &str, + range: Option<&str>, +) -> TableKeyInfo { + engine.gsi_cache_set(table_id, false); + let mut key_schema = vec![key(hash, KeyType::Hash)]; + let mut attribute_definitions = vec![def(hash, ScalarAttributeType::S)]; + if let Some(range) = range { + key_schema.push(key(range, KeyType::Range)); + attribute_definitions.push(def(range, ScalarAttributeType::S)); + } + TableKeyInfo { + table_id: table_id.to_owned(), + table_name: table_id.to_owned(), + key_schema: key_schema.clone(), + base_key_schema: key_schema, + attribute_definitions, + ..Default::default() + } +} + +fn named_key(hash: &str, range: Option<&str>) -> Item { + let mut key = Item::new(); + key.insert(hash.to_owned(), s("u1")); + if let Some(range) = range { + key.insert(range.to_owned(), s("s1")); + } + key +} + +/// `SET = :v`. +fn set(name: &str) -> Vec { + vec![UpdateAction::Set { + path: path(name), + value: Expr::Placeholder(":v".to_owned()), + }] +} + +/// The key condition ` = :p` and the map that supplies `:p`. +fn partition_condition(hash: &str) -> (KeyCondition, ExpressionMaps) { + let condition = KeyCondition { + pk_path: path(hash), + pk_value: Expr::Placeholder(":p".to_owned()), + extra_pk_conditions: Vec::new(), + sk_condition: None, + extra_sk_conditions: Vec::new(), + }; + let values = HashMap::from([(":p".to_owned(), s("u1"))]); + (condition, ExpressionMaps::new(HashMap::new(), values)) +} + +/// An UpdateItem that creates an item on a table with these key names +/// answers Ok, returns the item with its key under ALL_NEW and stores the +/// document PutItem stores. GetItem, Query and Scan return the item. An +/// UpdateItem that sets an attribute named `prefix`, which is a path prefix +/// of a key name, succeeds on the created item, and PutItem and DeleteItem +/// then read it as any other item. +async fn creation_matches_put_item( + engine: &MongoEngine, + table_id: &str, + hash: &str, + range: Option<&str>, + prefix: &str, +) { + let key_info = named_table(engine, table_id, hash, range); + let key = named_key(hash, range); + let created = with(key.clone(), "w", s("upd")); + + let (old, new) = engine + .update_item_impl( + &key_info, + &key, + &set("w"), + false, + true, + None, + &maps(s("upd")), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: creating update: {e:?}")); + assert_eq!(old, None, "{table_id}: no old image on creation"); + assert_eq!( + new, + Some(created.clone()), + "{table_id}: ALL_NEW has the key" + ); + assert_eq!( + stored(engine, &key_info, &key).await, + expected_document(&created, &key_info, 1), + "{table_id}: stored document" + ); + assert_eq!( + engine + .get_item_impl(&key_info, &key) + .await + .unwrap_or_else(|e| panic!("{table_id}: get: {e:?}")), + Some(created.clone()), + "{table_id}: GetItem returns the created item" + ); + + let (condition, query_maps) = partition_condition(hash); + let (query, _) = engine + .query_impl(&key_info, &condition, &query_maps, true, None, None, None) + .await + .unwrap_or_else(|e| panic!("{table_id}: query: {e:?}")); + assert_eq!(query, vec![created.clone()], "{table_id}: Query"); + let (scan, _) = engine + .scan_impl(&key_info, None, None, None, None, None) + .await + .unwrap_or_else(|e| panic!("{table_id}: scan: {e:?}")); + assert_eq!(scan, vec![created.clone()], "{table_id}: Scan"); + + let updated = with(created, prefix, s("p")); + let (_, new) = engine + .update_item_impl( + &key_info, + &key, + &set(prefix), + false, + true, + None, + &maps(s("p")), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: SET {prefix} on the created item: {e:?}")); + assert_eq!(new, Some(updated.clone()), "{table_id}: SET {prefix}"); + assert_eq!( + stored(engine, &key_info, &key).await, + expected_document(&updated, &key_info, 2), + "{table_id}: stored document after SET {prefix}" + ); + + let put = with(key.clone(), "w", s("put")); + let old = engine + .put_item_impl( + &key_info, + put.clone(), + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: put: {e:?}")); + assert_eq!( + old, + Some(updated), + "{table_id}: PutItem returns the old item" + ); + let deleted = engine + .delete_item_impl( + &key_info, + &key, + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: delete: {e:?}")); + assert_eq!( + deleted, + Some(put), + "{table_id}: DeleteItem returns the item" + ); +} + +/// A hash key named `a.b`: the creating UpdateItem would have stored +/// `item_data: { a: { b: ... } }`, which no read parses. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_update_item_creates_an_item_when_the_hash_key_name_has_a_dot() { + let engine = scratch_engine().await; + creation_matches_put_item(&engine, "dot-hash", "a.b", None, "a").await; + finish(engine).await; +} + +/// A range key named `s.k`, with a plain hash key. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_update_item_creates_an_item_when_the_range_key_name_has_a_dot() { + let engine = scratch_engine().await; + creation_matches_put_item(&engine, "dot-range", "pk", Some("s.k"), "s").await; + finish(engine).await; +} + +/// An UpdateItem on an item that already exists may set an attribute named +/// like a path prefix of a key name (`a` for the key `a.b`). Through the +/// native upsert the `$set` of `item_data.a` and the `$setOnInsert` of +/// `item_data.a.b` conflict, and MongoDB refuses the update before it looks +/// for the document, so this worked on main and must keep working. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_update_item_may_set_an_attribute_named_like_a_prefix_of_a_key_name() { + let engine = scratch_engine().await; + let cases: [(&str, Option<&str>, &str); 2] = [("a.b", None, "a"), ("pk", Some("s.k"), "s")]; + for (hash, range, prefix) in cases { + let table_id = format!("prefix-{prefix}"); + let key_info = named_table(&engine, &table_id, hash, range); + let key = named_key(hash, range); + let existing = with(key.clone(), "w", s("put")); + engine + .put_item_impl( + &key_info, + existing.clone(), + false, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: put: {e:?}")); + + let (_, new) = engine + .update_item_impl( + &key_info, + &key, + &set(prefix), + false, + true, + None, + &maps(AttributeValue::S("p".to_owned())), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: SET {prefix} on a stored item: {e:?}")); + let updated = with(existing, prefix, s("p")); + assert_eq!(new, Some(updated.clone()), "{table_id}: ALL_NEW"); + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&updated, &key_info, 1), + "{table_id}: stored document" + ); + } + finish(engine).await; +} diff --git a/crates/storage-mongodb/src/data_engine_update_create.rs b/crates/storage-mongodb/src/data_engine_update_create.rs index b3b73fecb..f4317e9be 100644 --- a/crates/storage-mongodb/src/data_engine_update_create.rs +++ b/crates/storage-mongodb/src/data_engine_update_create.rs @@ -22,9 +22,18 @@ //! the filter matches by equality. The `_id` collision is reported first and //! the filter does not name `_id`, so the server does not retry it and this //! module does: the item now exists, so the next attempt updates it. +//! +//! The `$setOnInsert` names each key attribute as the path +//! `item_data.`, and MongoDB reads a dot in that path as a +//! step into a sub-document. A key attribute named `a.b` would be inserted as +//! `item_data: { a: { b: ... } }`, where PutItem stores the field `"a.b"`, +//! and the same path conflicts with an update of an attribute named `a`. +//! A table with a key attribute name that `key_names_are_plain_paths` +//! rejects therefore does not take the native upsert; it takes the session +//! path, which writes the document `item_to_document` builds. use bson::Document; -use extenddb_core::types::{Item, TableKeyInfo}; +use extenddb_core::types::{Item, KeySchemaElement, TableKeyInfo}; use extenddb_storage::error::StorageError; use mongodb::options::{FindOneAndUpdateOptions, ReturnDocument}; @@ -32,6 +41,21 @@ use super::{TRANSIENT_RETRY_ATTEMPTS, backoff_sleep}; use crate::MongoEngine; use crate::data::key_fields::insert_only_key_fields; +/// Whether `name` is usable as one field name in a MongoDB update path. A dot +/// would split it into nested fields and a leading `$` would read as an +/// operator or positional path. +pub(super) fn is_plain_path_name(name: &str) -> bool { + !name.contains('.') && !name.starts_with('$') +} + +/// Whether every key attribute of `key_schema` is a plain path name, so that +/// [`MongoEngine::upsert_native_update`] can write it as `item_data.`. +pub(super) fn key_names_are_plain_paths(key_schema: &[KeySchemaElement]) -> bool { + key_schema + .iter() + .all(|ks| is_plain_path_name(&ks.attribute_name)) +} + /// Whether `e` is a duplicate-key error (E11000) as `findAndModify` reports /// it: a command error, where `insert_one` reports a write error. fn is_duplicate_key_command_error(e: &mongodb::error::Error) -> bool { @@ -85,6 +109,7 @@ impl MongoEngine { mod tests { use super::*; use bson::doc; + use extenddb_core::types::KeyType; fn command_error(code: i32) -> mongodb::error::Error { let command_error: mongodb::error::CommandError = bson::from_document(doc! { @@ -113,4 +138,38 @@ mod tests { assert!(!is_duplicate_key_command_error(&command_error(112))); assert!(!is_duplicate_key_command_error(&write_error(11000))); } + + fn schema(names: &[(&str, KeyType)]) -> Vec { + names + .iter() + .map(|(name, key_type)| KeySchemaElement { + attribute_name: (*name).to_owned(), + key_type: *key_type, + }) + .collect() + } + + #[test] + fn a_dot_or_a_leading_dollar_in_a_key_name_leaves_the_native_upsert() { + let plain = [("pk", KeyType::Hash), ("sk", KeyType::Range)]; + assert!(key_names_are_plain_paths(&schema(&plain))); + assert!(key_names_are_plain_paths(&schema(&[( + "a$b", + KeyType::Hash + )]))); + + for hash in ["a.b", "$k", ".", "pk."] { + assert!( + !key_names_are_plain_paths(&schema(&[(hash, KeyType::Hash)])), + "hash key {hash:?}" + ); + } + for range in ["s.k", "$s"] { + let keys = [("pk", KeyType::Hash), (range, KeyType::Range)]; + assert!( + !key_names_are_plain_paths(&schema(&keys)), + "range key {range:?}" + ); + } + } } diff --git a/tests/test_update_item_creates_key.py b/tests/test_update_item_creates_key.py index 89ebc50ad..cee23776e 100644 --- a/tests/test_update_item_creates_key.py +++ b/tests/test_update_item_creates_key.py @@ -77,6 +77,28 @@ def range_table(dynamodb_client): yield name +@pytest.fixture(scope="module") +def dotted_hash_table(dynamodb_client): + """A table whose hash key name contains a dot.""" + with scoped_table( + dynamodb_client, + attribute_definitions=[_s("a.b")], + key_schema=[_k("a.b", "HASH")], + ) as name: + yield name + + +@pytest.fixture(scope="module") +def dotted_range_table(dynamodb_client): + """A table whose range key name contains a dot.""" + with scoped_table( + dynamodb_client, + attribute_definitions=[_s("pk"), _s("s.k")], + key_schema=[_k("pk", "HASH"), _k("s.k", "RANGE")], + ) as name: + yield name + + @pytest.fixture(scope="module") def raw_client(endpoint_url): """A client that never retries, so every failure is seen.""" @@ -95,10 +117,19 @@ def raw_client(endpoint_url): return boto3.client(**kwargs) +KEY_NAMES = { + "hash": ("pk", None), + "range": ("pk", "sk"), + "dotted_hash": ("a.b", None), + "dotted_range": ("pk", "s.k"), +} + + def _key(table_kind: str, pk: str) -> dict: - key = {"pk": {"S": pk}} - if table_kind == "range": - key["sk"] = {"S": "s"} + hash_name, range_name = KEY_NAMES[table_kind] + key = {hash_name: {"S": pk}} + if range_name: + key[range_name] = {"S": "s"} return key @@ -137,7 +168,7 @@ def _put_then_get(client, table: str, key: dict) -> None: assert _get(client, table, key) == put -@pytest.mark.parametrize("table_kind", ["hash", "range"]) +@pytest.mark.parametrize("table_kind", ["hash", "range", "dotted_hash", "dotted_range"]) def test_set_creates_the_item_with_its_key(request, dynamodb_client, table_kind): table = request.getfixturevalue(f"{table_kind}_table") _warm(dynamodb_client, table, table_kind) @@ -147,6 +178,27 @@ def test_set_creates_the_item_with_its_key(request, dynamodb_client, table_kind) _put_then_get(dynamodb_client, table, key) +@pytest.mark.parametrize("table_kind", ["dotted_hash", "dotted_range"]) +def test_query_and_scan_return_an_item_created_on_a_dotted_key_table( + request, dynamodb_client, table_kind +): + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + key = _key(table_kind, f"qs-{uuid.uuid4().hex}") + _create(dynamodb_client, table, key) + created = {**key, "w": {"S": "upd"}} + hash_name = next(iter(key)) + resp = dynamodb_client.query( + TableName=table, + KeyConditionExpression="#h = :p", + ExpressionAttributeNames={"#h": hash_name}, + ExpressionAttributeValues={":p": key[hash_name]}, + ConsistentRead=True, + ) + assert resp["Items"] == [created] + assert created in dynamodb_client.scan(TableName=table, ConsistentRead=True)["Items"] + + @pytest.mark.parametrize("table_kind", ["hash", "range"]) def test_add_number_creates_the_item_with_its_key(request, dynamodb_client, table_kind): table = request.getfixturevalue(f"{table_kind}_table") From 2b8eac99eb7126ddd471b215617abcc8f30bdc70 Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 09/14] fix(mongodb): rebuild numeric sort keys as the plain decimal The startup repair took a numeric sort key from its Decimal128 and used the display form as the number. Decimal128 displays a value below 1e-6 in exponent form (0.0000001 is 1E-7), while a request carries, and PutItem builds the _id from, the plain decimal that validate_and_normalize_number produces. A document repaired from 0.0000001 got the _id 2:u1,4:1E-7, where PutItem builds 2:u1,9:0.0000001,. PutItem on that key then matched the document by sk_n, tried to replace the _id and failed on the immutable field. The repair had counted the document as repaired. The repair now passes the display text through validate_and_normalize_number. A value at or above 1e-6 displays as the plain decimal already and is unchanged. The unit test adds 0.0000001, -0.00000012345, 1E-30 and 1E-7 (the last two spelled as a request spells them) to the numeric cases, with 0.000001, which Decimal128 displays as written, and compares the whole repaired document including the _id. The live test plants the pre-fix document for four numeric sort keys, runs the repair, and runs a PutItem on each key. --- crates/storage-mongodb/src/data/key_fields.rs | 47 ++++++- .../src/data_engine_live_tests.rs | 3 + .../src/data_engine_live_tests_repair_keys.rs | 116 ++++++++++++++++++ 3 files changed, 163 insertions(+), 3 deletions(-) create mode 100644 crates/storage-mongodb/src/data_engine_live_tests_repair_keys.rs diff --git a/crates/storage-mongodb/src/data/key_fields.rs b/crates/storage-mongodb/src/data/key_fields.rs index a6df7aaac..4d645a5aa 100644 --- a/crates/storage-mongodb/src/data/key_fields.rs +++ b/crates/storage-mongodb/src/data/key_fields.rs @@ -22,6 +22,7 @@ use bson::{Bson, Document}; use extenddb_core::types::{ AttributeDefinition, AttributeValue, Item, KeySchemaElement, KeyType, ScalarAttributeType, }; +use extenddb_core::validation::number::validate_and_normalize_number; use extenddb_storage::error::StorageError; use extenddb_storage::util::sk_info; @@ -155,7 +156,9 @@ fn sk_value( ) -> Result { match (sk_type, value) { (ScalarAttributeType::S, Bson::String(s)) => Ok(AttributeValue::S(s.clone())), - (ScalarAttributeType::N, Bson::Decimal128(d)) => Ok(AttributeValue::N(d.to_string())), + (ScalarAttributeType::N, Bson::Decimal128(d)) => { + Ok(AttributeValue::N(plain_number_text(&d.to_string())?)) + } (ScalarAttributeType::B, Bson::String(hex)) => Ok(AttributeValue::B(hex_to_bytes(hex)?)), _ => Err(StorageError::Internal(format!( "{field} holds a {} value, not the sort key type", @@ -164,6 +167,19 @@ fn sk_value( } } +/// The text a request carries for a numeric sort key, from the display form +/// of its stored `Decimal128`. +/// +/// `Decimal128` displays a value below 1e-6 in exponent form (0.0000001 is +/// `1E-7`), while a request's number is the plain decimal that +/// `validate_and_normalize_number` produces, and PutItem builds the `_id` +/// from that text. A document rebuilt with the exponent form in its `_id` +/// would make every later PutItem on the key fail on the immutable `_id`. +fn plain_number_text(display: &str) -> Result { + validate_and_normalize_number(display) + .map_err(|e| StorageError::Internal(format!("sort key is not a number: {e}"))) +} + fn hex_to_bytes(hex: &str) -> Result, StorageError> { let bytes = hex.as_bytes(); if !bytes.len().is_multiple_of(2) { @@ -293,11 +309,36 @@ mod tests { assert_eq!(repaired, want); } + /// A number as a request carries it: `{"N": raw}` deserialised, which + /// normalises the text to the plain decimal that PutItem stores. + fn request_number(raw: &str) -> AttributeValue { + serde_json::from_value(serde_json::json!({ "N": raw })).expect("attribute value") + } + #[test] fn repair_restores_string_numeric_and_binary_sort_keys() { let cases = [ (ScalarAttributeType::S, AttributeValue::S("s1".to_owned())), (ScalarAttributeType::N, AttributeValue::N("12.5".to_owned())), + // `Decimal128` displays the next three in exponent form + // (`1E-7`, `-1.2345E-7`, `1E-30`), and the `_id` PutItem builds + // from them is the plain decimal. + ( + ScalarAttributeType::N, + AttributeValue::N("0.0000001".to_owned()), + ), + ( + ScalarAttributeType::N, + AttributeValue::N("-0.00000012345".to_owned()), + ), + (ScalarAttributeType::N, request_number("1E-30")), + // The same value as a request spells it. + (ScalarAttributeType::N, request_number("1E-7")), + // Plain below the switch to exponent form. + ( + ScalarAttributeType::N, + AttributeValue::N("0.000001".to_owned()), + ), ( ScalarAttributeType::B, AttributeValue::B(vec![0x01, 0xff, 0x10]), @@ -321,11 +362,11 @@ mod tests { "item_data": { "w": { "S": "upd" } }, }; let repaired = repaired_document(&bad, &schema, &defs).expect("repair"); - assert_eq!(repaired, want, "{sk_type:?} sort key"); + assert_eq!(repaired, want, "{sk_type:?} sort key {sk:?}"); assert_eq!( document_to_item(&repaired).expect("item"), want_item, - "{sk_type:?} sort key item" + "{sk_type:?} sort key {sk:?}: item" ); } } diff --git a/crates/storage-mongodb/src/data_engine_live_tests.rs b/crates/storage-mongodb/src/data_engine_live_tests.rs index df3cc184a..0f253df21 100644 --- a/crates/storage-mongodb/src/data_engine_live_tests.rs +++ b/crates/storage-mongodb/src/data_engine_live_tests.rs @@ -27,6 +27,9 @@ use crate::data::{data_collection_name, document_to_item, item_to_document, pk_f #[path = "data_engine_live_tests_key_names.rs"] mod key_names; +#[path = "data_engine_live_tests_repair_keys.rs"] +mod repair_keys; + fn uri() -> String { std::env::var("EXTENDDB_TEST_MONGODB_URI").expect( "set EXTENDDB_TEST_MONGODB_URI to a MongoDB replica set, e.g. \ diff --git a/crates/storage-mongodb/src/data_engine_live_tests_repair_keys.rs b/crates/storage-mongodb/src/data_engine_live_tests_repair_keys.rs new file mode 100644 index 000000000..a0e56672d --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_live_tests_repair_keys.rs @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for the startup repair of documents stored without their +//! key, for key shapes the tests in the parent module do not cover. +//! +//! The tests read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are +//! `#[ignore]`d, like the tests in the parent module. + +use extenddb_core::expression::ExpressionMaps; +use extenddb_core::types::{AttributeValue, Item, KeyType, ScalarAttributeType, TableKeyInfo}; + +use super::{ + KeyRepairSummary, def, finish, key, s, scratch_engine, store_without_key, stored, with, +}; +use crate::MongoEngine; +use crate::data::{document_to_item, item_to_document}; + +/// Key info for a table with a string hash key `pk` and a numeric sort key +/// `sk`, with the GSI cache set to "no indexes". +fn numeric_range_table(engine: &MongoEngine, table_id: &str) -> TableKeyInfo { + engine.gsi_cache_set(table_id, false); + let key_schema = vec![key("pk", KeyType::Hash), key("sk", KeyType::Range)]; + TableKeyInfo { + table_id: table_id.to_owned(), + table_name: table_id.to_owned(), + key_schema: key_schema.clone(), + base_key_schema: key_schema, + attribute_definitions: vec![ + def("pk", ScalarAttributeType::S), + def("sk", ScalarAttributeType::N), + ], + ..Default::default() + } +} + +/// The key `pk = "u1"`, `sk = `, with the number in the plain decimal +/// form a request carries. +fn number_key(number: &str) -> Item { + let mut key = Item::new(); + key.insert("pk".to_owned(), s("u1")); + key.insert("sk".to_owned(), AttributeValue::N(number.to_owned())); + key +} + +/// A numeric sort key below 1e-6, which `Decimal128` displays in exponent +/// form (`1E-7`), is repaired into the `_id` that PutItem builds from the +/// plain decimal. A PutItem on each repaired key then succeeds and returns +/// the repaired item as the old image; built on the exponent form, the repair +/// would leave an `_id` that PutItem cannot replace. The values at and above +/// 1e-6 are repaired the same way and show the other side of the switch. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_keeps_a_small_numeric_sort_key_writable() { + const NUMBERS: [&str; 4] = ["0.0000001", "-0.00000012345", "0.000001", "12.5"]; + let engine = scratch_engine().await; + let key_info = numeric_range_table(&engine, "repair-small-numbers"); + for number in NUMBERS { + store_without_key(&engine, &key_info, &number_key(number), Some("upd")).await; + } + + let mut summary = KeyRepairSummary::default(); + engine + .repair_table( + &key_info.table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut summary, + ) + .await + .expect("repair"); + assert_eq!( + summary, + KeyRepairSummary { + tables_scanned: 0, + documents_repaired: NUMBERS.len() as u64, + documents_failed: 0, + } + ); + + for number in NUMBERS { + let key = number_key(number); + let repaired = with(key.clone(), "w", s("upd")); + let put = with(key.clone(), "w", s("put")); + let old = engine + .put_item_impl( + &key_info, + put.clone(), + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{number}: PutItem on the repaired key: {e:?}")); + assert_eq!( + old, + Some(repaired), + "{number}: PutItem returns the repaired item" + ); + + let after = stored(&engine, &key_info, &key).await; + let want = item_to_document(&put, &key_info.key_schema, &key_info.attribute_definitions) + .expect("document"); + assert_eq!( + after.get("_id"), + want.get("_id"), + "{number}: the _id is the one PutItem builds" + ); + assert_eq!( + document_to_item(&after).expect("item"), + put, + "{number}: PutItem replaced the item" + ); + } + finish(engine).await; +} From 60a5af8d7e29f738e350921afd30782a6c4bd35b Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 10/14] fix(mongodb): drop the duplicate-key rerun on the transaction path The UpdateItem transaction path ran its body again when the replace_one upsert that creates an item failed with a duplicate-key error, on the reasoning that another writer had created the item after the find_one. On MongoDB 7.0 that race does not surface as a duplicate-key error. An insert that loses to one committed after the transaction's snapshot, or still uncommitted in another transaction, fails with WriteConflict (code 112, TransientTransactionError), and the transient retry that was already there reruns the body and reads the winner's item. The duplicate-key error appears only when the colliding _id is visible in the snapshot, and every rerun meets it again, so the branch used all 50 attempts and ended in TransactionConflictException. The statement is back to the form main has: the error maps to Internal at once, with the server's duplicate-key text. A live test creates an item with eight concurrent numeric ADD calls, which take the transaction path, and checks that all apply: c is the sum and _v counts every call. It passes with and without the branch, because the write conflict is retried either way. A second live test plants a document that holds the key text as its _id under a different pk and checks that a creating ADD answers Internal with the duplicate-key error; with the branch it answers TransactionConflict after the 50 attempts. --- crates/storage-mongodb/src/data_engine.rs | 13 +- .../src/data_engine_live_tests.rs | 3 + .../src/data_engine_live_tests_txn_create.rs | 113 ++++++++++++++++++ 3 files changed, 118 insertions(+), 11 deletions(-) create mode 100644 crates/storage-mongodb/src/data_engine_live_tests_txn_create.rs diff --git a/crates/storage-mongodb/src/data_engine.rs b/crates/storage-mongodb/src/data_engine.rs index c4d3b4f8c..871a04254 100644 --- a/crates/storage-mongodb/src/data_engine.rs +++ b/crates/storage-mongodb/src/data_engine.rs @@ -995,20 +995,11 @@ impl MongoEngine { let opts = mongodb::options::ReplaceOptions::builder() .upsert(true) .build(); - // A duplicate-key error here means another writer - // created the item after the find_one above. Run the - // transaction again so that it reads that item and - // applies the update to it. - match coll - .replace_one(key_filter.clone(), new_doc) + coll.replace_one(key_filter.clone(), new_doc) .with_options(opts) .session(&mut session) .await - { - Ok(_) => {} - Err(e) if is_duplicate_key(&e) => return Err(TxErr::Transient), - Err(e) => return Err(TxErr::from(e)), - } + .map_err(TxErr::from)?; } self.sync_indexes_in_session( diff --git a/crates/storage-mongodb/src/data_engine_live_tests.rs b/crates/storage-mongodb/src/data_engine_live_tests.rs index 0f253df21..46d8fa180 100644 --- a/crates/storage-mongodb/src/data_engine_live_tests.rs +++ b/crates/storage-mongodb/src/data_engine_live_tests.rs @@ -30,6 +30,9 @@ mod key_names; #[path = "data_engine_live_tests_repair_keys.rs"] mod repair_keys; +#[path = "data_engine_live_tests_txn_create.rs"] +mod txn_create; + fn uri() -> String { std::env::var("EXTENDDB_TEST_MONGODB_URI").expect( "set EXTENDDB_TEST_MONGODB_URI to a MongoDB replica set, e.g. \ diff --git a/crates/storage-mongodb/src/data_engine_live_tests_txn_create.rs b/crates/storage-mongodb/src/data_engine_live_tests_txn_create.rs new file mode 100644 index 000000000..c057786aa --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_live_tests_txn_create.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for UpdateItem calls that create an item on the +//! transaction path. A numeric ADD takes it: the call finds no item inside a +//! snapshot transaction and inserts with a `replace_one` upsert. +//! +//! The tests read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are +//! `#[ignore]`d, like the tests in the parent module. + +use std::sync::Arc; + +use bson::{Document, doc}; +use extenddb_core::expression::{Expr, UpdateAction}; +use extenddb_storage::error::StorageError; + +use super::{ + expected_document, finish, item_key, maps, n, path, scratch_engine, stored, table, with, +}; +use crate::data::data_collection_name; + +fn add_one() -> Vec { + vec![UpdateAction::Add { + path: path("c"), + value: Expr::Placeholder(":v".to_owned()), + }] +} + +/// Eight numeric ADD calls creating the same item at once all apply. Each +/// finds no item, and the creators that lose to the first insert get a write +/// conflict from MongoDB, which the engine retries: the retry reads the +/// winner's item and adds to it. The stored `c` is the sum and `_v` counts +/// every call. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_concurrent_numeric_adds_that_create_one_item_all_apply() { + const WRITERS: i64 = 8; + let engine = Arc::new(scratch_engine().await); + let key_info = Arc::new(table(&engine, "race-add", false)); + let key = item_key(false); + let mut tasks = Vec::new(); + for _ in 0..WRITERS { + let engine = engine.clone(); + let key_info = key_info.clone(); + let key = key.clone(); + tasks.push(tokio::spawn(async move { + engine + .update_item_impl( + &key_info, + &key, + &add_one(), + false, + true, + None, + &maps(n("1")), + None, + ) + .await + })); + } + for task in tasks { + task.await + .expect("task") + .unwrap_or_else(|e| panic!("a creating ADD failed: {e:?}")); + } + let want = with(key.clone(), "c", n(&WRITERS.to_string())); + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&want, &key_info, WRITERS) + ); + let engine = Arc::try_unwrap(engine).ok().expect("all tasks done"); + finish(engine).await; +} + +/// A creating numeric ADD whose `_id` is already held by a document with a +/// different `pk` fails with the duplicate-key error, and the call does not +/// run again. The other document is visible in the transaction's snapshot, +/// so every rerun would meet it again; a creator that loses a race for the +/// key is a write conflict instead, which the test above covers. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_creating_add_that_collides_on_id_answers_the_duplicate_key_error() { + let engine = scratch_engine().await; + let key_info = table(&engine, "collide-add", false); + let key = item_key(false); + engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)) + .insert_one(doc! { "_id": "u1", "pk": "other", "item_data": { "pk": { "S": "other" } } }) + .await + .expect("colliding document"); + + let err = engine + .update_item_impl( + &key_info, + &key, + &add_one(), + false, + true, + None, + &maps(n("1")), + None, + ) + .await + .expect_err("the insert cannot succeed"); + match err { + StorageError::Internal(message) => assert!( + message.contains("E11000 duplicate key error"), + "expected the server's duplicate-key error, got {message}" + ), + other => panic!("expected Internal with the duplicate-key error, got {other:?}"), + } + finish(engine).await; +} From d8fb10deb227550729398bd04bfbe345d94a2f54 Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 11/14] fix(mongodb): continue the key repair past a table it cannot scan The startup repair ended the whole pass at the first table-level error. A cursor error in repair_table, for instance from a table dropped while the pass ran, came back through the ? in repair_items_stored_without_key, so every table after it kept its keyless documents until the next start, and one error line was the only sign. The pass now logs the table, counts it in the summary as stopped by an error, and goes on to the next table. The final line is logged at warn when a table or a document failed. The scan returned whole documents although only the ObjectId _id is used. It projects _id now, so the _id index answers it. In a profiled run on MongoDB 7.0 the same find over five keyless documents examines five documents without the projection and none with it. The module comment states how a client write behaves while a repair transaction holds its document. The write waits for the transaction and sets no maxTimeMS. Normally that is one batch commit. After a process is killed mid-batch it is the time the server takes to expire the transaction, transactionLifetimeLimitSeconds, which is 60 s by default and was 60 on the test server, for up to 500 keys. A live test puts a catalog row whose table id MongoDB refuses as a collection name (it contains $) ahead of a readable table, and checks that the pass returns Ok, counts one stopped table and repairs the second table's document. A second live test reads the profiler entry of the scan and checks that it returned all 20 documents and examined none. --- .../src/data_engine_key_repair.rs | 38 ++++- .../src/data_engine_live_tests.rs | 7 + .../src/data_engine_live_tests_repair_keys.rs | 1 + .../src/data_engine_live_tests_repair_pass.rs | 153 ++++++++++++++++++ 4 files changed, 193 insertions(+), 6 deletions(-) create mode 100644 crates/storage-mongodb/src/data_engine_live_tests_repair_pass.rs diff --git a/crates/storage-mongodb/src/data_engine_key_repair.rs b/crates/storage-mongodb/src/data_engine_key_repair.rs index c5c9e714f..a3cb710cc 100644 --- a/crates/storage-mongodb/src/data_engine_key_repair.rs +++ b/crates/storage-mongodb/src/data_engine_key_repair.rs @@ -33,6 +33,19 @@ //! document's key answer as they did before this repair existed. Progress is //! logged every [`PROGRESS_EVERY`] documents and once at the end. A failure //! on one document is logged and skipped; the summary reports it. +//! +//! A scan that stops on an error, as it does when a table is dropped while the +//! pass runs, is logged, and the pass goes on to the next table. The summary +//! counts it, and the documents that scan had not reached wait for the next +//! start. The scan projects `_id` only, so the `_id` index answers it without +//! reading a document. +//! +//! A client write to a key whose document a repair transaction holds waits for +//! that transaction, and a foreground write sets no `maxTimeMS`. Normally the +//! wait is one batch commit. A process killed mid-batch leaves its transaction +//! open until the server expires it (`transactionLifetimeLimitSeconds`, 60 s +//! by default), and writes to the up to [`REPAIR_BATCH`] keys of that batch +//! wait that long. use std::sync::Arc; @@ -68,6 +81,9 @@ const CONFLICT_WAIT: std::time::Duration = std::time::Duration::from_secs(2); pub struct KeyRepairSummary { /// Tables whose data collection was scanned. pub tables_scanned: u64, + /// Tables among those whose scan stopped on an error. The documents that + /// scan had not reached wait for the next start; each stop is logged. + pub tables_failed: u64, /// Documents rewritten with their key. pub documents_repaired: u64, /// Documents left as they were because they could not be rebuilt or @@ -78,7 +94,7 @@ pub struct KeyRepairSummary { impl KeyRepairSummary { #[must_use] pub fn is_clean(&self) -> bool { - self.documents_repaired == 0 && self.documents_failed == 0 + self.documents_repaired == 0 && self.documents_failed == 0 && self.tables_failed == 0 } } @@ -86,8 +102,9 @@ impl std::fmt::Display for KeyRepairSummary { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!( f, - "{} tables scanned, {} documents stored without their key repaired, {} failed", - self.tables_scanned, self.documents_repaired, self.documents_failed + "{} tables scanned, {} stopped by an error, {} documents stored without their key \ + repaired, {} failed", + self.tables_scanned, self.tables_failed, self.documents_repaired, self.documents_failed ) } } @@ -115,8 +132,16 @@ impl MongoEngine { } }; summary.tables_scanned += 1; - self.repair_table(&table_id, &key_schema, &attribute_definitions, &mut summary) - .await?; + if let Err(e) = self + .repair_table(&table_id, &key_schema, &attribute_definitions, &mut summary) + .await + { + tracing::error!( + "key repair: table {table_id}: the scan stopped ({e}); the documents it \ + had not reached wait for the next start" + ); + summary.tables_failed += 1; + } } Ok(summary) } @@ -134,6 +159,7 @@ impl MongoEngine { .collection::(&data_collection_name(table_id)); let mut cursor = coll .find(doc! { "_id": { "$type": "objectId" } }) + .projection(doc! { "_id": 1 }) .await .map_err(|e| StorageError::Internal(e.to_string()))?; // The rebuilt documents have the key text as their `_id`, outside the @@ -475,7 +501,7 @@ pub(crate) fn spawn_key_repair(engine: Arc) -> tokio::task::JoinHan "startup key repair finished: {repair}, {seconds:.2} s, \ {rate:.0} documents a second" ); - if repair.documents_failed == 0 { + if repair.documents_failed == 0 && repair.tables_failed == 0 { tracing::info!("{line}"); } else { tracing::warn!("{line}"); diff --git a/crates/storage-mongodb/src/data_engine_live_tests.rs b/crates/storage-mongodb/src/data_engine_live_tests.rs index 46d8fa180..e65edeab8 100644 --- a/crates/storage-mongodb/src/data_engine_live_tests.rs +++ b/crates/storage-mongodb/src/data_engine_live_tests.rs @@ -33,6 +33,9 @@ mod repair_keys; #[path = "data_engine_live_tests_txn_create.rs"] mod txn_create; +#[path = "data_engine_live_tests_repair_pass.rs"] +mod repair_pass; + fn uri() -> String { std::env::var("EXTENDDB_TEST_MONGODB_URI").expect( "set EXTENDDB_TEST_MONGODB_URI to a MongoDB replica set, e.g. \ @@ -434,6 +437,7 @@ async fn live_repair_rewrites_documents_stored_without_their_key() { summary, KeyRepairSummary { tables_scanned: 0, + tables_failed: 0, documents_repaired: 1, documents_failed: 0, }, @@ -527,6 +531,7 @@ async fn live_repair_rewrites_a_document_stored_with_no_item_data() { summary, KeyRepairSummary { tables_scanned: 0, + tables_failed: 0, documents_repaired: 1, documents_failed: 0, }, @@ -612,6 +617,7 @@ async fn live_repair_spans_more_than_one_batch() { summary, KeyRepairSummary { tables_scanned: 0, + tables_failed: 0, documents_repaired: COUNT as u64, documents_failed: 0, } @@ -685,6 +691,7 @@ async fn live_repair_reads_the_key_schema_from_the_catalog() { summary, KeyRepairSummary { tables_scanned: 1, + tables_failed: 0, documents_repaired: 1, documents_failed: 0, } diff --git a/crates/storage-mongodb/src/data_engine_live_tests_repair_keys.rs b/crates/storage-mongodb/src/data_engine_live_tests_repair_keys.rs index a0e56672d..039368c37 100644 --- a/crates/storage-mongodb/src/data_engine_live_tests_repair_keys.rs +++ b/crates/storage-mongodb/src/data_engine_live_tests_repair_keys.rs @@ -72,6 +72,7 @@ async fn live_repair_keeps_a_small_numeric_sort_key_writable() { summary, KeyRepairSummary { tables_scanned: 0, + tables_failed: 0, documents_repaired: NUMBERS.len() as u64, documents_failed: 0, } diff --git a/crates/storage-mongodb/src/data_engine_live_tests_repair_pass.rs b/crates/storage-mongodb/src/data_engine_live_tests_repair_pass.rs new file mode 100644 index 000000000..0d5ce8f93 --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_live_tests_repair_pass.rs @@ -0,0 +1,153 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for the pass that finds the documents to repair: how it +//! reads each table and how it goes on after a table it cannot read. +//! +//! The tests read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are +//! `#[ignore]`d, like the tests in the parent module. + +use bson::{Bson, Document, doc}; +use extenddb_core::types::TableKeyInfo; +use futures::TryStreamExt; + +use super::{ + KeyRepairSummary, finish, item_key, s, scratch_engine, store_many_without_key, + store_without_key, stored, table, with, +}; +use crate::data::{data_collection_name, document_to_item}; + +/// A catalog `tables` row for `table_id` with the key schema of `key_info`. +fn catalog_row(table_name: &str, table_id: &str, key_info: &TableKeyInfo) -> Document { + doc! { + "_id": { "account_id": "123456789012", "table_name": table_name }, + "table_id": table_id, + "table_status": "ACTIVE", + "key_schema": bson::to_bson(&key_info.key_schema).expect("bson"), + "attribute_definitions": bson::to_bson(&key_info.attribute_definitions).expect("bson"), + } +} + +/// An integer field of a profiler entry, which MongoDB writes as int32 or +/// int64. +fn count(entry: &Document, field: &str) -> i64 { + match entry.get(field) { + Some(Bson::Int32(n)) => i64::from(*n), + Some(Bson::Int64(n)) => *n, + other => panic!("{field} in {entry:?} is {other:?}, not an integer"), + } +} + +/// A scan that stops on an error does not stop the pass. The first catalog +/// row carries a table id with a `$`, which MongoDB refuses in a collection +/// name, so the scan of its data collection fails. The second row is an +/// ordinary table with one document stored without its key. The pass counts +/// the first table as stopped by an error, repairs the second, and returns Ok. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_goes_on_to_the_next_table_when_a_scan_fails() { + let engine = scratch_engine().await; + let key_info = table(&engine, "pass-good", false); + let key = item_key(false); + let tables = engine.catalog_db.collection::("tables"); + tables + .insert_one(catalog_row("unreadable", "pass-bad$id", &key_info)) + .await + .expect("catalog row of the table that cannot be scanned"); + tables + .insert_one(catalog_row("readable", &key_info.table_id, &key_info)) + .await + .expect("catalog row of the readable table"); + store_without_key(&engine, &key_info, &key, Some("upd")).await; + + let summary = engine + .repair_items_stored_without_key() + .await + .expect("a table that cannot be scanned does not fail the pass"); + assert_eq!( + summary, + KeyRepairSummary { + tables_scanned: 2, + tables_failed: 1, + documents_repaired: 1, + documents_failed: 0, + } + ); + assert!(!summary.is_clean(), "a stopped scan is not a clean pass"); + assert_eq!( + document_to_item(&stored(&engine, &key_info, &key).await).expect("item"), + with(key, "w", s("upd")), + "the table after the failing one was repaired" + ); + finish(engine).await; +} + +/// The scan reads `_id` only. With the projection the `_id` index answers +/// it, and the server's profiler records that the find returned every +/// document it found and examined none; without the projection it examines +/// each one to return its fields. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_scan_is_answered_from_the_id_index() { + const DOCUMENTS: usize = 20; + let engine = scratch_engine().await; + let key_info = table(&engine, "pass-scan", false); + store_many_without_key(&engine, &key_info, DOCUMENTS).await; + + engine + .data_db + .run_command(doc! { "profile": 2 }) + .await + .expect("profiler on"); + let mut summary = KeyRepairSummary::default(); + engine + .repair_table( + &key_info.table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut summary, + ) + .await + .expect("repair"); + engine + .data_db + .run_command(doc! { "profile": 0 }) + .await + .expect("profiler off"); + assert_eq!(summary.documents_repaired, DOCUMENTS as u64); + + let namespace = format!( + "{}.{}", + engine.data_db.name(), + data_collection_name(&key_info.table_id) + ); + let queries: Vec = engine + .data_db + .collection::("system.profile") + .find(doc! { "ns": &namespace, "op": "query" }) + .await + .expect("profile entries") + .try_collect() + .await + .expect("profile entries"); + // The scan filters on the `_id` type; the batch transaction's own read + // filters on a list of ids. + let scans: Vec<&Document> = queries + .iter() + .filter(|entry| { + entry + .get_document("command") + .and_then(|command| command.get_document("filter")) + .and_then(|filter| filter.get_document("_id")) + .is_ok_and(|id| id.contains_key("$type")) + }) + .collect(); + assert_eq!(scans.len(), 1, "one scan of {namespace}: {queries:?}"); + let scan = scans[0]; + assert_eq!(count(scan, "nreturned"), DOCUMENTS as i64, "{scan:?}"); + assert_eq!( + count(scan, "docsExamined"), + 0, + "the scan is answered from the _id index: {scan:?}" + ); + finish(engine).await; +} From 168f80fbeaa763e27371e75f78df0ff051626753 Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 12/14] test(mongodb): cover the task that serve starts for the key repair The live tests called repair_table and repair_items_stored_without_key directly, so spawn_key_repair could stop running the pass and every test would still pass. The new live test registers a catalog table, plants a document in the shape the earlier upsert stored (an ObjectId _id and no key attributes), awaits the JoinHandle that spawn_key_repair returns, and checks that the document is the one item_to_document builds with its _v kept. The call to spawn_key_repair in lib.rs is still not under a test. The MongoDB pytest job starts one server before any test runs and runs the files in parallel, so a pytest cannot plant a document ahead of the startup pass. --- .../src/data_engine_live_tests_repair_pass.rs | 53 +++++++++++++++++-- 1 file changed, 50 insertions(+), 3 deletions(-) diff --git a/crates/storage-mongodb/src/data_engine_live_tests_repair_pass.rs b/crates/storage-mongodb/src/data_engine_live_tests_repair_pass.rs index 0d5ce8f93..3d9cd5d8d 100644 --- a/crates/storage-mongodb/src/data_engine_live_tests_repair_pass.rs +++ b/crates/storage-mongodb/src/data_engine_live_tests_repair_pass.rs @@ -1,20 +1,24 @@ // SPDX-License-Identifier: Apache-2.0 //! Live MongoDB tests for the pass that finds the documents to repair: how it -//! reads each table and how it goes on after a table it cannot read. +//! reads each table, how it goes on after a table it cannot read, and the task +//! that `serve` starts to run it. //! //! The tests read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are //! `#[ignore]`d, like the tests in the parent module. +use std::sync::Arc; + use bson::{Bson, Document, doc}; use extenddb_core::types::TableKeyInfo; use futures::TryStreamExt; use super::{ - KeyRepairSummary, finish, item_key, s, scratch_engine, store_many_without_key, - store_without_key, stored, table, with, + KeyRepairSummary, expected_document, finish, item_key, s, scratch_engine, + store_many_without_key, store_without_key, stored, table, with, }; use crate::data::{data_collection_name, document_to_item}; +use crate::data_engine::key_repair::spawn_key_repair; /// A catalog `tables` row for `table_id` with the key schema of `key_info`. fn catalog_row(table_name: &str, table_id: &str, key_info: &TableKeyInfo) -> Document { @@ -151,3 +155,46 @@ async fn live_repair_scan_is_answered_from_the_id_index() { ); finish(engine).await; } + +/// `spawn_key_repair`, the task `serve` starts (`lib.rs`), runs the whole +/// pass: with a catalog table whose collection holds a document stored +/// without its key, the document is rebuilt by the time the task finishes, +/// and the task does finish. The other tests call `repair_table` and +/// `repair_items_stored_without_key` directly. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_the_spawned_key_repair_rewrites_documents_stored_without_their_key() { + let engine = scratch_engine().await; + let key_info = table(&engine, "pass-spawned", false); + let key = item_key(false); + engine + .catalog_db + .collection::("tables") + .insert_one(catalog_row("spawned", &key_info.table_id, &key_info)) + .await + .expect("catalog row"); + store_without_key(&engine, &key_info, &key, Some("upd")).await; + assert!( + stored(&engine, &key_info, &key) + .await + .get_object_id("_id") + .is_ok(), + "planted in the shape the pre-fix upsert stored" + ); + + let engine = Arc::new(engine); + spawn_key_repair(engine.clone()) + .await + .expect("the repair task finishes"); + + let want = with(key.clone(), "w", s("upd")); + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&want, &key_info, 1), + "the task repaired the document" + ); + let engine = Arc::try_unwrap(engine) + .ok() + .expect("the finished task dropped its engine"); + finish(engine).await; +} From 1ba328e49dbc656c387cede43abac1e742246699 Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 13/14] test(mongodb): cover N and B partition keys and trim the REMOVE case key_fields.rs said item_to_document writes _v, and update_create.rs said the created document is the one PutItem would have written with _v 1. Neither item_to_document nor PutItem writes _v; the update's $inc adds it. Both module comments now say so, and so do the doc comments of two live tests that described the stored document the same way. The live test of the created document asserted, for a REMOVE-only update, that ALL_NEW and the old image of the PutItem that follows both carry the key-only item. The service creates no item for that request and its ALL_NEW answer for one is not measured, so the REMOVE case now asserts the stored document and that the PutItem succeeds, and no returned item. repaired_document decodes a partition key from the pk text through pk_value, and only string partition keys were under test. The new unit test repairs a numeric partition key, a numeric partition key below 1e-6 (whose pk text keeps its plain form) and a binary partition key, and compares the whole rebuilt document with the one item_to_document builds. --- crates/storage-mongodb/src/data/key_fields.rs | 52 +++++++++++++++++-- .../src/data_engine_live_tests.rs | 33 +++++++----- .../src/data_engine_update_create.rs | 5 +- 3 files changed, 70 insertions(+), 20 deletions(-) diff --git a/crates/storage-mongodb/src/data/key_fields.rs b/crates/storage-mongodb/src/data/key_fields.rs index 4d645a5aa..972503c76 100644 --- a/crates/storage-mongodb/src/data/key_fields.rs +++ b/crates/storage-mongodb/src/data/key_fields.rs @@ -4,11 +4,13 @@ //! document from something other than a full item. //! //! [`item_to_document`] writes a base-table document as `{ _id, pk, sk_*, -//! item_data, _v }`, with `_id` the key text and the key attributes inside -//! `item_data`. An UpdateItem that creates an item through a native MongoDB -//! upsert starts from the key filter (`pk` and `sk_*`) and the update -//! operators, so it needs the `_id` and the `item_data` key attributes added -//! with `$setOnInsert`: [`insert_only_key_fields`] builds that document. +//! item_data }`, with `_id` the key text and the key attributes inside +//! `item_data`. It writes no `_v`: the version is added by the UpdateItem +//! write itself, as `$inc` on the native path. An UpdateItem that creates an +//! item through a native MongoDB upsert starts from the key filter (`pk` and +//! `sk_*`) and the update operators, so it needs the `_id` and the +//! `item_data` key attributes added with `$setOnInsert`: +//! [`insert_only_key_fields`] builds that document. //! //! Before that was done, such an upsert stored a document with a generated //! ObjectId `_id` and no key attributes in `item_data` (issue #436), and for @@ -309,6 +311,46 @@ mod tests { assert_eq!(repaired, want); } + #[test] + fn repair_restores_numeric_and_binary_partition_keys() { + // The partition key text is stored as written, so a number below 1e-6 + // keeps its plain form there. + let cases = [ + (ScalarAttributeType::N, AttributeValue::N("42".to_owned())), + ( + ScalarAttributeType::N, + AttributeValue::N("0.0000001".to_owned()), + ), + ( + ScalarAttributeType::B, + AttributeValue::B(vec![0x01, 0xff, 0x10]), + ), + ]; + for (pk_type, pk) in cases { + let schema = vec![key("pk", KeyType::Hash)]; + let defs = vec![def("pk", pk_type)]; + let want_item = item(&[ + ("pk", pk.clone()), + ("w", AttributeValue::S("upd".to_owned())), + ]); + let mut want = item_to_document(&want_item, &schema, &defs).expect("document"); + want.insert("_v", 2_i64); + let bad = doc! { + "_id": bson::oid::ObjectId::new(), + "pk": want.get("pk").cloned().expect("pk text"), + "item_data": { "w": { "S": "upd" } }, + "_v": 2_i64, + }; + let repaired = repaired_document(&bad, &schema, &defs).expect("repair"); + assert_eq!(repaired, want, "{pk:?} partition key"); + assert_eq!( + document_to_item(&repaired).expect("item"), + want_item, + "{pk:?} partition key item" + ); + } + } + /// A number as a request carries it: `{"N": raw}` deserialised, which /// normalises the text to the plain decimal that PutItem stores. fn request_number(raw: &str) -> AttributeValue { diff --git a/crates/storage-mongodb/src/data_engine_live_tests.rs b/crates/storage-mongodb/src/data_engine_live_tests.rs index e65edeab8..997d49b03 100644 --- a/crates/storage-mongodb/src/data_engine_live_tests.rs +++ b/crates/storage-mongodb/src/data_engine_live_tests.rs @@ -156,10 +156,13 @@ fn expected_document(item: &Item, key_info: &TableKeyInfo, v: i64) -> Document { /// An UpdateItem that creates the item, on the native fast path (SET with a /// placeholder, REMOVE) and on the transaction path (numeric ADD), stores the -/// same document PutItem would store: `_id` as the key text, the key -/// attributes in `item_data`, version 1. ALL_NEW returns the item with its -/// key. A PutItem on the key then replaces the item and returns it as the old -/// image. +/// document `item_to_document` builds, plus `_v` 1: `_id` as the key text and +/// the key attributes in `item_data`. For SET and ADD, ALL_NEW returns the +/// item with its key, and a PutItem on the key replaces the item and returns +/// it as the old image. For REMOVE the test pins the stored document and that +/// a PutItem on the key succeeds. The service creates no item for a +/// REMOVE-only update and its ALL_NEW answer for one is not measured, so the +/// REMOVE case asserts neither the ALL_NEW answer nor the old image. #[tokio::test] #[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] async fn live_update_item_that_creates_an_item_stores_its_key_like_put_item() { @@ -212,11 +215,13 @@ async fn live_update_item_that_creates_an_item_stores_its_key_like_put_item() { .await .unwrap_or_else(|e| panic!("{table_id}: update: {e:?}")); assert_eq!(old, None, "{table_id}: no old image on creation"); - assert_eq!( - new, - Some(want.clone()), - "{table_id}: ALL_NEW carries the key" - ); + if *name != "remove" { + assert_eq!( + new, + Some(want.clone()), + "{table_id}: ALL_NEW carries the key" + ); + } assert_eq!( stored(&engine, &key_info, &key).await, expected_document(&want, &key_info, 1), @@ -235,7 +240,9 @@ async fn live_update_item_that_creates_an_item_stores_its_key_like_put_item() { ) .await .unwrap_or_else(|e| panic!("{table_id}: put after update: {e:?}")); - assert_eq!(old, Some(want), "{table_id}: put returns the created item"); + if *name != "remove" { + assert_eq!(old, Some(want), "{table_id}: put returns the created item"); + } let after = stored(&engine, &key_info, &key).await; assert_eq!( document_to_item(&after).expect("item"), @@ -386,9 +393,9 @@ async fn store_without_key( } /// The repair rewrites a document stored without its key into the document -/// PutItem would have written, keeping the attributes and the version. The -/// item then reads with its key and a PutItem on it succeeds. Documents with -/// their key are left alone, and a second run finds nothing. +/// `item_to_document` builds, keeping the attributes and carrying `_v` over. +/// The item then reads with its key and a PutItem on it succeeds. Documents +/// with their key are left alone, and a second run finds nothing. #[tokio::test] #[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] async fn live_repair_rewrites_documents_stored_without_their_key() { diff --git a/crates/storage-mongodb/src/data_engine_update_create.rs b/crates/storage-mongodb/src/data_engine_update_create.rs index f4317e9be..a7f4e7e01 100644 --- a/crates/storage-mongodb/src/data_engine_update_create.rs +++ b/crates/storage-mongodb/src/data_engine_update_create.rs @@ -12,8 +12,9 @@ //! key text and the whole item (`item_to_document`). Every later PutItem on //! the key then fails on the immutable `_id` (issue #436). //! [`MongoEngine::upsert_native_update`] adds a `$setOnInsert` with the `_id` -//! and the key attributes, so the created document is the one PutItem would -//! have written, with `_v` 1 from the update's `$inc`. +//! and the key attributes, so the created document is the one +//! `item_to_document` builds for PutItem, plus `_v` 1 from the update's +//! `$inc`. //! //! Two writers creating the same item at once each find no document and each //! insert; the second insert fails with a duplicate-key error on `_id`. The From 9607fd783103e7148f2e3151c3c65a3b71926723 Mon Sep 17 00:00:00 2001 From: Lee Hannigan <29015211+LeeroyHannigan@users.noreply.github.com> Date: Thu, 29 Oct 2026 22:45:47 +0000 Subject: [PATCH 14/14] test(mongodb): move the creation live tests to their own file data_engine_live_tests.rs had grown to 711 lines. The three tests of the document an UpdateItem writes when it creates an item move to data_engine_live_tests_create.rs, which is 227 lines. The shared helpers and the repair tests stay in the original file, which is now 501 lines. The moved block is byte-identical, and only the imports and the module declaration change. The test names gain the module name create::, so the filter data_engine::live_tests:: that devtools/run-mongodb-tests uses still selects all 16 live tests. --- .../src/data_engine_live_tests.rs | 224 +---------------- .../src/data_engine_live_tests_create.rs | 227 ++++++++++++++++++ 2 files changed, 235 insertions(+), 216 deletions(-) create mode 100644 crates/storage-mongodb/src/data_engine_live_tests_create.rs diff --git a/crates/storage-mongodb/src/data_engine_live_tests.rs b/crates/storage-mongodb/src/data_engine_live_tests.rs index 997d49b03..4ea858a58 100644 --- a/crates/storage-mongodb/src/data_engine_live_tests.rs +++ b/crates/storage-mongodb/src/data_engine_live_tests.rs @@ -1,8 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 -//! Live MongoDB tests for the shape of the document an UpdateItem writes when -//! it creates an item (issue #436), and for the startup repair of documents an -//! earlier build stored without their key. +//! Live MongoDB tests for the startup repair of documents that an earlier build +//! stored without their key (issue #436), and the helpers the live test +//! modules below share. The tests of the document an UpdateItem writes when it +//! creates an item are in `data_engine_live_tests_create.rs`. //! //! They read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are //! `#[ignore]`d so `cargo test --lib` passes without one; run them with @@ -12,18 +13,19 @@ use std::collections::HashMap; use bson::{Document, doc}; -use extenddb_core::expression::{Expr, ExpressionMaps, PathElement, UpdateAction}; +use extenddb_core::expression::{ExpressionMaps, PathElement}; use extenddb_core::types::{ AttributeDefinition, AttributeValue, Item, KeySchemaElement, KeyType, ScalarAttributeType, TableKeyInfo, }; -use extenddb_storage::error::StorageError; - use super::key_repair::KeyRepairSummary; use crate::MongoEngine; use crate::data::{data_collection_name, document_to_item, item_to_document, pk_filter}; +#[path = "data_engine_live_tests_create.rs"] +mod create; + #[path = "data_engine_live_tests_key_names.rs"] mod key_names; @@ -154,216 +156,6 @@ fn expected_document(item: &Item, key_info: &TableKeyInfo, v: i64) -> Document { doc } -/// An UpdateItem that creates the item, on the native fast path (SET with a -/// placeholder, REMOVE) and on the transaction path (numeric ADD), stores the -/// document `item_to_document` builds, plus `_v` 1: `_id` as the key text and -/// the key attributes in `item_data`. For SET and ADD, ALL_NEW returns the -/// item with its key, and a PutItem on the key replaces the item and returns -/// it as the old image. For REMOVE the test pins the stored document and that -/// a PutItem on the key succeeds. The service creates no item for a -/// REMOVE-only update and its ALL_NEW answer for one is not measured, so the -/// REMOVE case asserts neither the ALL_NEW answer nor the old image. -#[tokio::test] -#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] -async fn live_update_item_that_creates_an_item_stores_its_key_like_put_item() { - let engine = scratch_engine().await; - let cases: [(&str, Vec, AttributeValue, Item); 3] = [ - ( - "set", - vec![UpdateAction::Set { - path: path("w"), - value: Expr::Placeholder(":v".to_owned()), - }], - s("upd"), - with(Item::new(), "w", s("upd")), - ), - ( - "remove", - vec![UpdateAction::Remove { path: path("w") }], - s("unused"), - Item::new(), - ), - ( - "add", - vec![UpdateAction::Add { - path: path("c"), - value: Expr::Placeholder(":v".to_owned()), - }], - n("5"), - with(Item::new(), "c", n("5")), - ), - ]; - for range in [false, true] { - for (name, actions, value, extra) in &cases { - let table_id = format!("{name}-{}", if range { "range" } else { "hash" }); - let key_info = table(&engine, &table_id, range); - let key = item_key(range); - let mut want = key.clone(); - want.extend(extra.clone()); - - let (old, new) = engine - .update_item_impl( - &key_info, - &key, - actions, - false, - true, - None, - &maps(value.clone()), - None, - ) - .await - .unwrap_or_else(|e| panic!("{table_id}: update: {e:?}")); - assert_eq!(old, None, "{table_id}: no old image on creation"); - if *name != "remove" { - assert_eq!( - new, - Some(want.clone()), - "{table_id}: ALL_NEW carries the key" - ); - } - assert_eq!( - stored(&engine, &key_info, &key).await, - expected_document(&want, &key_info, 1), - "{table_id}: stored document" - ); - - let put = with(key.clone(), "w", s("put")); - let old = engine - .put_item_impl( - &key_info, - put.clone(), - true, - None, - &ExpressionMaps::default(), - None, - ) - .await - .unwrap_or_else(|e| panic!("{table_id}: put after update: {e:?}")); - if *name != "remove" { - assert_eq!(old, Some(want), "{table_id}: put returns the created item"); - } - let after = stored(&engine, &key_info, &key).await; - assert_eq!( - document_to_item(&after).expect("item"), - put, - "{table_id}: put replaced the item" - ); - } - } - finish(engine).await; -} - -/// Sixty-four UpdateItem calls creating the same item at once, on the fast -/// path, from eight runtime threads: every call succeeds, the stored -/// document has the shape PutItem writes, and `_v` counts every call, so -/// each one applied to the one document after losing the create race. -#[tokio::test(flavor = "multi_thread", worker_threads = 8)] -#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] -async fn live_concurrent_creating_updates_on_one_key_all_apply() { - const WRITERS: i64 = 64; - let engine = std::sync::Arc::new(scratch_engine().await); - for range in [false, true] { - let table_id = format!("race-{}", if range { "range" } else { "hash" }); - let key_info = std::sync::Arc::new(table(&engine, &table_id, range)); - let key = item_key(range); - let mut tasks = Vec::new(); - for i in 0..WRITERS { - let engine = engine.clone(); - let key_info = key_info.clone(); - let key = key.clone(); - tasks.push(tokio::spawn(async move { - let actions = vec![UpdateAction::Set { - path: path("w"), - value: Expr::Placeholder(":v".to_owned()), - }]; - engine - .update_item_impl( - &key_info, - &key, - &actions, - false, - true, - None, - &maps(s(&format!("w{i}"))), - None, - ) - .await - })); - } - for task in tasks { - let (old, new) = task - .await - .expect("task") - .unwrap_or_else(|e| panic!("{table_id}: a creating update failed: {e:?}")); - assert_eq!(old, None, "{table_id}"); - let new = new.expect("ALL_NEW"); - for (name, value) in &key { - assert_eq!( - new.get(name), - Some(value), - "{table_id}: ALL_NEW has the key" - ); - } - } - let doc = stored(&engine, &key_info, &key).await; - let item = document_to_item(&doc).expect("item"); - let w = item.get("w").cloned().expect("w was set"); - let mut want = key.clone(); - want.insert("w".to_owned(), w); - assert_eq!( - doc, - expected_document(&want, &key_info, WRITERS), - "{table_id}" - ); - } - let engine = std::sync::Arc::try_unwrap(engine) - .ok() - .expect("all tasks done"); - finish(engine).await; -} - -/// A creating upsert whose insert always collides on `_id` (another document -/// holds the key text as its `_id` under a different `pk`) exhausts the -/// retry budget and answers `TransactionConflict`, which the server maps to -/// `TransactionConflictException`, not an internal error: the duplicate-key -/// command error `findAndModify` returns is recognised. -#[tokio::test] -#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] -async fn live_creating_update_that_always_loses_on_id_answers_transaction_conflict() { - let engine = scratch_engine().await; - let key_info = table(&engine, "collide-hash", false); - let key = item_key(false); - engine - .data_db - .collection::(&data_collection_name(&key_info.table_id)) - .insert_one(doc! { "_id": "u1", "pk": "other", "item_data": { "pk": { "S": "other" } } }) - .await - .expect("colliding document"); - let actions = vec![UpdateAction::Set { - path: path("w"), - value: Expr::Placeholder(":v".to_owned()), - }]; - let err = engine - .update_item_impl( - &key_info, - &key, - &actions, - false, - true, - None, - &maps(s("upd")), - None, - ) - .await - .expect_err("the insert cannot succeed"); - assert!( - matches!(err, StorageError::TransactionConflict(_)), - "expected TransactionConflict, got {err:?}" - ); - finish(engine).await; -} - /// Store a document the way the fast path did before the fix: an upsert whose /// filter is the key and whose update sets one attribute, or, when `w` is /// `None`, removes one. MongoDB copies the filter's `pk` and `sk_*` fields, diff --git a/crates/storage-mongodb/src/data_engine_live_tests_create.rs b/crates/storage-mongodb/src/data_engine_live_tests_create.rs new file mode 100644 index 000000000..0f76faa39 --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_live_tests_create.rs @@ -0,0 +1,227 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for the document an UpdateItem writes when it creates an +//! item (issue #436), on the native fast path and on the transaction path. +//! +//! The tests read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are +//! `#[ignore]`d, like the tests in the parent module. + +use bson::{Document, doc}; +use extenddb_core::expression::{Expr, ExpressionMaps, UpdateAction}; +use extenddb_core::types::{AttributeValue, Item}; +use extenddb_storage::error::StorageError; + +use super::{ + expected_document, finish, item_key, maps, n, path, s, scratch_engine, stored, table, with, +}; +use crate::data::{data_collection_name, document_to_item}; + +/// An UpdateItem that creates the item, on the native fast path (SET with a +/// placeholder, REMOVE) and on the transaction path (numeric ADD), stores the +/// document `item_to_document` builds, plus `_v` 1: `_id` as the key text and +/// the key attributes in `item_data`. For SET and ADD, ALL_NEW returns the +/// item with its key, and a PutItem on the key replaces the item and returns +/// it as the old image. For REMOVE the test pins the stored document and that +/// a PutItem on the key succeeds. The service creates no item for a +/// REMOVE-only update and its ALL_NEW answer for one is not measured, so the +/// REMOVE case asserts neither the ALL_NEW answer nor the old image. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_update_item_that_creates_an_item_stores_its_key_like_put_item() { + let engine = scratch_engine().await; + let cases: [(&str, Vec, AttributeValue, Item); 3] = [ + ( + "set", + vec![UpdateAction::Set { + path: path("w"), + value: Expr::Placeholder(":v".to_owned()), + }], + s("upd"), + with(Item::new(), "w", s("upd")), + ), + ( + "remove", + vec![UpdateAction::Remove { path: path("w") }], + s("unused"), + Item::new(), + ), + ( + "add", + vec![UpdateAction::Add { + path: path("c"), + value: Expr::Placeholder(":v".to_owned()), + }], + n("5"), + with(Item::new(), "c", n("5")), + ), + ]; + for range in [false, true] { + for (name, actions, value, extra) in &cases { + let table_id = format!("{name}-{}", if range { "range" } else { "hash" }); + let key_info = table(&engine, &table_id, range); + let key = item_key(range); + let mut want = key.clone(); + want.extend(extra.clone()); + + let (old, new) = engine + .update_item_impl( + &key_info, + &key, + actions, + false, + true, + None, + &maps(value.clone()), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: update: {e:?}")); + assert_eq!(old, None, "{table_id}: no old image on creation"); + if *name != "remove" { + assert_eq!( + new, + Some(want.clone()), + "{table_id}: ALL_NEW carries the key" + ); + } + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&want, &key_info, 1), + "{table_id}: stored document" + ); + + let put = with(key.clone(), "w", s("put")); + let old = engine + .put_item_impl( + &key_info, + put.clone(), + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: put after update: {e:?}")); + if *name != "remove" { + assert_eq!(old, Some(want), "{table_id}: put returns the created item"); + } + let after = stored(&engine, &key_info, &key).await; + assert_eq!( + document_to_item(&after).expect("item"), + put, + "{table_id}: put replaced the item" + ); + } + } + finish(engine).await; +} + +/// Sixty-four UpdateItem calls creating the same item at once, on the fast +/// path, from eight runtime threads: every call succeeds, the stored +/// document has the shape PutItem writes, and `_v` counts every call, so +/// each one applied to the one document after losing the create race. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_concurrent_creating_updates_on_one_key_all_apply() { + const WRITERS: i64 = 64; + let engine = std::sync::Arc::new(scratch_engine().await); + for range in [false, true] { + let table_id = format!("race-{}", if range { "range" } else { "hash" }); + let key_info = std::sync::Arc::new(table(&engine, &table_id, range)); + let key = item_key(range); + let mut tasks = Vec::new(); + for i in 0..WRITERS { + let engine = engine.clone(); + let key_info = key_info.clone(); + let key = key.clone(); + tasks.push(tokio::spawn(async move { + let actions = vec![UpdateAction::Set { + path: path("w"), + value: Expr::Placeholder(":v".to_owned()), + }]; + engine + .update_item_impl( + &key_info, + &key, + &actions, + false, + true, + None, + &maps(s(&format!("w{i}"))), + None, + ) + .await + })); + } + for task in tasks { + let (old, new) = task + .await + .expect("task") + .unwrap_or_else(|e| panic!("{table_id}: a creating update failed: {e:?}")); + assert_eq!(old, None, "{table_id}"); + let new = new.expect("ALL_NEW"); + for (name, value) in &key { + assert_eq!( + new.get(name), + Some(value), + "{table_id}: ALL_NEW has the key" + ); + } + } + let doc = stored(&engine, &key_info, &key).await; + let item = document_to_item(&doc).expect("item"); + let w = item.get("w").cloned().expect("w was set"); + let mut want = key.clone(); + want.insert("w".to_owned(), w); + assert_eq!( + doc, + expected_document(&want, &key_info, WRITERS), + "{table_id}" + ); + } + let engine = std::sync::Arc::try_unwrap(engine) + .ok() + .expect("all tasks done"); + finish(engine).await; +} + +/// A creating upsert whose insert always collides on `_id` (another document +/// holds the key text as its `_id` under a different `pk`) exhausts the +/// retry budget and answers `TransactionConflict`, which the server maps to +/// `TransactionConflictException`, not an internal error: the duplicate-key +/// command error `findAndModify` returns is recognised. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_creating_update_that_always_loses_on_id_answers_transaction_conflict() { + let engine = scratch_engine().await; + let key_info = table(&engine, "collide-hash", false); + let key = item_key(false); + engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)) + .insert_one(doc! { "_id": "u1", "pk": "other", "item_data": { "pk": { "S": "other" } } }) + .await + .expect("colliding document"); + let actions = vec![UpdateAction::Set { + path: path("w"), + value: Expr::Placeholder(":v".to_owned()), + }]; + let err = engine + .update_item_impl( + &key_info, + &key, + &actions, + false, + true, + None, + &maps(s("upd")), + None, + ) + .await + .expect_err("the insert cannot succeed"); + assert!( + matches!(err, StorageError::TransactionConflict(_)), + "expected TransactionConflict, got {err:?}" + ); + finish(engine).await; +}