fix: revalidate unique-index updates after locking

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
 

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)
 
-

codegen/Cargo.toml

@@ -8,6 +8,7 @@ repository = "https://github.com/pathscale/WorkTable"
 
 [features]
 s3-support = []
+strict-unique-index-revalidation = []
 
 [lib]
 name = "worktable_codegen"

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<Ident>>,
         unsized_fields: Option<Vec<&Ident>>,
@@ -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 || &current.#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());

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<Ident>>,
         unsized_fields: Option<Vec<&Ident>>,
@@ -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 || &current.#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());

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();