diff --git a/Cargo.toml b/Cargo.toml index 93bf735..d8a1d9a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ categories = ["database-implementations", "data-structures", "caching"] [features] perf_measurements = ["dep:performance_measurement", "dep:performance_measurement_codegen"] s3-support = ["dep:rusty-s3", "dep:url", "dep:reqwest", "dep:walkdir", "worktable_codegen/s3-support"] +strict-unique-index-revalidation = ["worktable_codegen/strict-unique-index-revalidation"] # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html diff --git a/README.md b/README.md index d2d09f3..80159a2 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,13 @@ S3 support layers *on top of* the disk engine rather than replacing it. worktable = { version = "0.9", features = ["s3-support"] } # S3 sync, optional ``` +Generated updates through a unique secondary index can opt into post-lock +primary-key and predicate revalidation with the +`strict-unique-index-revalidation` feature. It prevents a changed or stale +unique-index entry from redirecting an update to another row while the update +holds the original row's lock. The feature is off by default so existing +latency-sensitive builds retain their current generated update path. + ## Relationship to `data_bucket` WorkTable is built on [`data_bucket`](https://crates.io/crates/data_bucket), which @@ -396,4 +403,3 @@ enum WorkTableError Check out - [Examples](./examples) - diff --git a/codegen/Cargo.toml b/codegen/Cargo.toml index 05ace59..bd8920a 100644 --- a/codegen/Cargo.toml +++ b/codegen/Cargo.toml @@ -8,6 +8,7 @@ repository = "https://github.com/pathscale/WorkTable" [features] s3-support = [] +strict-unique-index-revalidation = [] [lib] name = "worktable_codegen" diff --git a/codegen/src/generators/in_memory/queries/update.rs b/codegen/src/generators/in_memory/queries/update.rs index 15503b9..e9d5927 100644 --- a/codegen/src/generators/in_memory/queries/update.rs +++ b/codegen/src/generators/in_memory/queries/update.rs @@ -154,13 +154,11 @@ impl InMemoryGenerator { let idents = &op.columns; if let Some(index) = index { - let index_name = &index.name; - if index.is_unique { self.gen_unique_update( snake_case_name, name, - index_name, + index, idents, indexes_columns.as_ref(), unsized_columns, @@ -658,7 +656,7 @@ impl InMemoryGenerator { &self, snake_case_name: String, name: &Ident, - index: &Ident, + index: &Index, idents: &[Ident], idx_idents: Option<&Vec>, unsized_fields: Option>, @@ -668,6 +666,8 @@ impl InMemoryGenerator { let query_ident = Ident::new(format!("{name}Query").as_str(), Span::mixed_site()); let by_ident = Ident::new(format!("{name}By").as_str(), Span::mixed_site()); let lock_ident = WorktableNameGenerator::get_update_query_lock_ident(&snake_case_name); + let by_field = &index.field; + let index = &index.name; let row_updates = idents .iter() @@ -692,6 +692,36 @@ impl InMemoryGenerator { } }; let custom_lock = self.gen_custom_lock_for_update(lock_ident); + let target_validation = if cfg!(feature = "strict-unique-index-revalidation") { + quote! { + match self.0.data.select_non_vacuumed(link) { + core::result::Result::Ok(current) => { + // The unique-index entry may have changed while we + // waited for the original row's lock. Never mutate + // the newly indexed row while holding another row's + // lock, and reject stale index entries whose row no + // longer satisfies the generated predicate. + if current.get_primary_key() != pk || ¤t.#by_field != &by { + return core::result::Result::Err(WorkTableError::NotFound); + } + break link; + } + core::result::Result::Err(e) if e.is_vacuumed() => continue, + core::result::Result::Err(e) => return core::result::Result::Err(e.into()), + } + } + } else { + quote! { + if let Err(e) = self.0.data.select_non_vacuumed(link) { + if e.is_vacuumed() { + continue; + } + return Err(e.into()); + } else { + break link; + } + } + }; quote! { pub async fn #method_ident(&self, row: #query_ident, by: #by_ident) -> core::result::Result<(), WorkTableError> { @@ -724,14 +754,7 @@ impl InMemoryGenerator { .map(|v| v.get().value.into()) .ok_or(WorkTableError::NotFound)?; - if let Err(e) = self.0.data.select_non_vacuumed(link) { - if e.is_vacuumed() { - continue; - } - return Err(e.into()); - } else { - break link; - } + #target_validation }; let op_id = OperationId::Single(uuid::Uuid::now_v7()); diff --git a/codegen/src/generators/persist/queries/update.rs b/codegen/src/generators/persist/queries/update.rs index 99149ff..59442bb 100644 --- a/codegen/src/generators/persist/queries/update.rs +++ b/codegen/src/generators/persist/queries/update.rs @@ -154,13 +154,11 @@ impl PersistGenerator { let idents = &op.columns; if let Some(index) = index { - let index_name = &index.name; - if index.is_unique { self.gen_unique_update( snake_case_name, name, - index_name, + index, idents, indexes_columns.as_ref(), unsized_columns, @@ -615,7 +613,7 @@ impl PersistGenerator { &self, snake_case_name: String, name: &Ident, - index: &Ident, + index: &Index, idents: &[Ident], idx_idents: Option<&Vec>, unsized_fields: Option>, @@ -625,6 +623,8 @@ impl PersistGenerator { let query_ident = Ident::new(format!("{name}Query").as_str(), Span::mixed_site()); let by_ident = Ident::new(format!("{name}By").as_str(), Span::mixed_site()); let lock_ident = WorktableNameGenerator::get_update_query_lock_ident(&snake_case_name); + let by_field = &index.field; + let index = &index.name; let row_updates = idents .iter() @@ -649,6 +649,36 @@ impl PersistGenerator { } }; let custom_lock = self.gen_custom_lock_for_update(lock_ident); + let target_validation = if cfg!(feature = "strict-unique-index-revalidation") { + quote! { + match self.0.data.select_non_vacuumed(link) { + core::result::Result::Ok(current) => { + // The unique-index entry may have changed while we + // waited for the original row's lock. Never mutate + // the newly indexed row while holding another row's + // lock, and reject stale index entries whose row no + // longer satisfies the generated predicate. + if current.get_primary_key() != pk || ¤t.#by_field != &by { + return core::result::Result::Err(WorkTableError::NotFound); + } + break link; + } + core::result::Result::Err(e) if e.is_vacuumed() => continue, + core::result::Result::Err(e) => return core::result::Result::Err(e.into()), + } + } + } else { + quote! { + if let Err(e) = self.0.data.select_non_vacuumed(link) { + if e.is_vacuumed() { + continue; + } + return Err(e.into()); + } else { + break link; + } + } + }; quote! { pub async fn #method_ident(&self, row: #query_ident, by: #by_ident) -> core::result::Result<(), WorkTableError> { @@ -681,14 +711,7 @@ impl PersistGenerator { .map(|v| v.get().value.into()) .ok_or(WorkTableError::NotFound)?; - if let Err(e) = self.0.data.select_non_vacuumed(link) { - if e.is_vacuumed() { - continue; - } - return Err(e.into()); - } else { - break link; - } + #target_validation }; let op_id = OperationId::Single(uuid::Uuid::now_v7()); diff --git a/tests/worktable/index/update_query.rs b/tests/worktable/index/update_query.rs index 20d0c56..d681bb9 100644 --- a/tests/worktable/index/update_query.rs +++ b/tests/worktable/index/update_query.rs @@ -2,6 +2,8 @@ use crate::worktable::index::{ Test3NonUniqueRow, Test3NonUniqueWorkTable, Test3UniqueRow, Test3UniqueWorkTable, TwoAttrByThirdQuery, UniqueTwoAttrByThirdQuery, }; +#[cfg(feature = "strict-unique-index-revalidation")] +use worktable::WorkTableError; use worktable::prelude::SelectQueryExecutor; #[tokio::test] @@ -49,6 +51,47 @@ async fn update_two_via_query_unique_indexes() { assert_eq!(updated, Some(new_row)) } +#[tokio::test] +#[cfg(feature = "strict-unique-index-revalidation")] +async fn unique_update_rejects_stale_index_link() { + let table = Test3UniqueWorkTable::default(); + let row1 = Test3UniqueRow { + val: 1, + attr1: "row-1".to_string(), + attr2: 10, + attr3: 100, + id: 0, + }; + let row2 = Test3UniqueRow { + val: 2, + attr1: "row-2".to_string(), + attr2: 20, + attr3: 200, + id: 1, + }; + let pk1 = table.insert(row1.clone()).unwrap(); + let pk2 = table.insert(row2.clone()).unwrap(); + + // Model the stale-link state that a concurrent unique-index rewrite can + // expose: the requested key resolves to a row that does not own that key. + let row2_link = table.0.primary_index.pk_map.get(&pk2).unwrap().get().value; + table.0.indexes.idx3.insert(row1.attr3, (*row2_link).into()); + + let result = table + .update_unique_two_attr_by_third( + UniqueTwoAttrByThirdQuery { + attr1: "must-not-apply".to_string(), + attr2: 30, + }, + row1.attr3, + ) + .await; + + assert!(matches!(result, Err(WorkTableError::NotFound))); + assert_eq!(table.select(pk1), Some(row1)); + assert_eq!(table.select(pk2), Some(row2)); +} + #[tokio::test] async fn update_with_reinsert_and_secondary_unique_violation() { let test_table = Test3UniqueWorkTable::default();