Concurrency and persistence fix series: deterministic suite, sound queue, CDC-complete vacuum

Cargo.toml

@@ -23,17 +23,16 @@ s3-support = ["dep:rusty-s3", "dep:url", "dep:reqwest", "dep:walkdir", "worktabl
 [dependencies]
 async-trait = "0.1.89"
 convert_case = "0.6.0"
-data_bucket = "=0.3.15"
+data_bucket = "=0.4.0"
 # data_bucket = { git = "https://github.com/pathscale/DataBucket", branch = "page_cdc_correction", version = "0.2.7" }
 # data_bucket = { path = "../DataBucket", version = "0.3.14" }
 derive_more = { version = "2.0.1", features = ["from", "error", "display", "debug", "into"] }
 eyre = "0.6.12"
 fastrand = "2.3.0"
 futures = "0.3.30"
-indexset = { version = "=0.16.0", features = ["concurrent", "cdc", "multimap"] }
+indexset = { package = "WorkTablesIndex", version = "=0.0.1", features = ["concurrent", "cdc", "multimap"] }
 # indexset = { path = "../indexset", version = "0.15.0", features = ["concurrent", "cdc", "multimap"] }
 # indexset = { package = "wt-indexset", version = "=0.12.12", features = ["concurrent", "cdc", "multimap"] }
-lockfree = { version = "0.5.1" }
 log = "0.4.29"
 ordered-float = "5.0.0"
 parking_lot = "0.12.3"

codegen/src/generators/in_memory/locks.rs

@@ -63,6 +63,9 @@ impl InMemoryGenerator {
 
         quote! {
             impl RowLock for #lock_ident {
+                fn new() -> Self {
+                    #lock_ident::new()
+                }
                 #is_locked_fn
                 #lock_fn
                 #with_lock_fn

codegen/src/generators/in_memory/queries/delete.rs

@@ -4,6 +4,7 @@ use convert_case::{Case, Casing};
 use proc_macro2::{Ident, Span, TokenStream};
 use quote::quote;
 
+use crate::common::model::Index;
 use crate::common::model::Operation;
 use crate::common::name_generator::{WorktableNameGenerator, is_float};
 use crate::generators::in_memory::InMemoryGenerator;
@@ -148,7 +149,7 @@ impl InMemoryGenerator {
                     if index.is_unique {
                         Self::gen_unique_delete(type_, &method_ident, index_name)
                     } else {
-                        Self::gen_non_unique_delete(type_, &method_ident, index_name)
+                        Self::gen_non_unique_delete(type_, &method_ident, index)
                     }
                 } else {
                     Self::gen_brute_force_delete_field(&op.by, type_, &method_ident)
@@ -180,7 +181,9 @@ impl InMemoryGenerator {
         }
     }
 
-    fn gen_non_unique_delete(type_: &TokenStream, name: &Ident, index: &Ident) -> TokenStream {
+    fn gen_non_unique_delete(type_: &TokenStream, name: &Ident, index: &Index) -> TokenStream {
+        let by_field = &index.field;
+        let index = &index.name;
         let by = if is_float(type_.to_string().as_str()) {
             quote! {
                 &OrderedFloat(by)
@@ -192,10 +195,46 @@ impl InMemoryGenerator {
         };
         quote! {
             pub async fn #name(&self, by: #type_) -> core::result::Result<(), WorkTableError> {
-                let rows_to_update = self.0.indexes.#index.get(#by).map(|kv| kv.1).collect::<Vec<_>>();
-                for link in rows_to_update {
-                    let row = self.0.data.select_non_ghosted(link.0).map_err(WorkTableError::PagesError)?;
-                    self.delete(row.get_primary_key()).await?;
+                // Snapshot the matching rows as validated primary keys before
+                // deleting anything. Storage links are not stable identities:
+                // a concurrent delete can free a slot and an insert can reuse
+                // it for an unrelated row before this loop runs, so resolving
+                // a stale link later could delete the wrong row. Every
+                // candidate is resolved and checked against the predicate at
+                // snapshot time; a reused slot only stays in the set if the
+                // row now living there genuinely matches. Keys are sorted for
+                // a deterministic delete order (non-unique indexes iterate
+                // equal keys in random-discriminator order) and the per-row
+                // delete takes its own row lock and resolves by primary key,
+                // never through the snapshotted link.
+                let mut pks: Vec<_> = Vec::new();
+                for link in self.0.indexes.#index.get(#by).map(|kv| kv.1.0) {
+                    match self.0.data.select_non_ghosted(link) {
+                        core::result::Result::Ok(r) => {
+                            if r.#by_field == by {
+                                pks.push(r.get_primary_key());
+                            }
+                        }
+                        // The row vanished between the index read and the
+                        // resolve; it is simply not part of the snapshot.
+                        core::result::Result::Err(e) if e.is_row_absent() => {}
+                        // Anything else (corrupt page, invalid link, ...) is a
+                        // real storage error, not an empty snapshot slot.
+                        core::result::Result::Err(e) => {
+                            return core::result::Result::Err(WorkTableError::PagesError(e));
+                        }
+                    }
+                }
+                pks.sort_unstable();
+                pks.dedup();
+                for pk in pks {
+                    match self.delete(pk).await {
+                        core::result::Result::Ok(()) => {}
+                        // Deleted concurrently after the snapshot: the goal
+                        // state for this row is already reached.
+                        core::result::Result::Err(WorkTableError::NotFound) => {}
+                        core::result::Result::Err(e) => return core::result::Result::Err(e),
+                    }
                 }
                 core::result::Result::Ok(())
             }

codegen/src/generators/in_memory/queries/locks.rs

@@ -58,8 +58,16 @@ impl InMemoryGenerator {
 
                 let lock_ident = WorktableNameGenerator::get_update_query_lock_ident(&snake_case_name);
 
-                let columns = &updates.get(name).as_ref().expect("exists").columns;
-                let lock_fn = Self::gen_rows_lock_fn(columns, lock_ident);
+                let op = updates.get(name).expect("exists");
+                // The lock set covers the updated columns AND the predicate
+                // (`by`) column: multi-row updates re-validate the predicate
+                // under this lock, which is only sound if no concurrent
+                // operation can rewrite that column while it is held.
+                let mut columns = op.columns.clone();
+                if !columns.contains(&op.by) {
+                    columns.push(op.by.clone());
+                }
+                let lock_fn = Self::gen_rows_lock_fn(&columns, lock_ident);
 
                 quote! {
                     #lock_fn
@@ -103,31 +111,17 @@ impl InMemoryGenerator {
 
         quote! {
             let lock_id = self.0.lock_manager.next_id();
-            if let Some(lock) = self.0.lock_manager.get(&pk) {
-                let mut lock_guard = lock.write().await;
-                #[allow(clippy::mutable_key_type)]
-                let (locks, op_lock) = lock_guard.lock(lock_id);
-                drop(lock_guard);
-                futures::future::join_all(locks.iter().map(|l| l.wait()).collect::<Vec<_>>()).await;
-
-                op_lock
-            } else {
-                #[allow(clippy::mutable_key_type)]
-                let (lock, op_lock) = #lock_ident::with_lock(lock_id);
-                let lock = std::sync::Arc::new(tokio::sync::RwLock::new(lock));
-                let mut guard = lock.write().await;
-                if let Some(old_lock) = self.0.lock_manager.insert(pk.clone(), lock.clone()) {
-                    let mut old_lock_guard = old_lock.write().await;
-                    #[allow(clippy::mutable_key_type)]
-                    let locks = guard.merge(&mut *old_lock_guard);
-                    drop(old_lock_guard);
-                    drop(guard);
-
-                    futures::future::join_all(locks.iter().map(|l| l.wait()).collect::<Vec<_>>()).await;
-                }
-
-                op_lock
-            }
+            // Same atomic acquire as the per-column path: see LockMap::get_or_insert_with.
+            let lock = self
+                .0
+                .lock_manager
+                .get_or_insert_with(pk.clone(), #lock_ident::new);
+            let mut lock_guard = lock.write().await;
+            #[allow(clippy::mutable_key_type)]
+            let (locks, op_lock) = lock_guard.lock(lock_id);
+            drop(lock_guard);
+            futures::future::join_all(locks.iter().map(|l| l.wait()).collect::<Vec<_>>()).await;
+            op_lock
         }
     }
 
@@ -137,31 +131,21 @@ impl InMemoryGenerator {
 
         quote! {
             let lock_id = self.0.lock_manager.next_id();
-            if let Some(lock) = self.0.lock_manager.get(&pk) {
-                let mut lock_guard = lock.write().await;
-                #[allow(clippy::mutable_key_type)]
-                let (locks, op_lock) = lock_guard.#ident(lock_id);
-                drop(lock_guard);
-                futures::future::join_all(locks.iter().map(|l| l.wait()).collect::<Vec<_>>()).await;
-                op_lock
-            } else {
-                let mut lock = #lock_ident::new();
-                #[allow(clippy::mutable_key_type)]
-                let (_, op_lock) = lock.#ident(lock_id);
-                let lock = std::sync::Arc::new(tokio::sync::RwLock::new(lock));
-                let mut guard = lock.write().await;
-                if let Some(old_lock) = self.0.lock_manager.insert(pk.clone(), lock.clone()) {
-                    let mut old_lock_guard = old_lock.write().await;
-                    #[allow(clippy::mutable_key_type)]
-                    let locks = guard.merge(&mut *old_lock_guard);
-                    drop(old_lock_guard);
-                    drop(guard);
-
-                    futures::future::join_all(locks.iter().map(|l| l.wait()).collect::<Vec<_>>()).await;
-                }
-
-                op_lock
-            }
+            // One atomic acquire, no check-then-act. Splitting this into `get`
+            // then `insert` let two tasks both miss, both build a lock and both
+            // enter the row: the loser merged into the winner's lock, but the
+            // winner had already registered its operation on a lock that was no
+            // longer the map's, so it never waited for the loser.
+            let lock = self
+                .0
+                .lock_manager
+                .get_or_insert_with(pk.clone(), #lock_ident::new);
+            let mut lock_guard = lock.write().await;
+            #[allow(clippy::mutable_key_type)]
+            let (locks, op_lock) = lock_guard.#ident(lock_id);
+            drop(lock_guard);
+            futures::future::join_all(locks.iter().map(|l| l.wait()).collect::<Vec<_>>()).await;
+            op_lock
         }
     }
 }

codegen/src/generators/in_memory/queries/type.rs

@@ -144,7 +144,7 @@ impl InMemoryGenerator {
                 .updates
                 .values()
                 .map(|op| {
-                    let ident = Ident::new(format!("{}By", &op.name).as_str(), Span::mixed_site());
+                    let ident = Ident::new(format!("{}By", op.name).as_str(), Span::mixed_site());
                     let field_type = self
                         .columns
                         .columns_map

codegen/src/generators/in_memory/queries/update.rs

@@ -1,7 +1,7 @@
 use proc_macro2::Literal;
 use std::collections::HashMap;
 
-use crate::common::model::Operation;
+use crate::common::model::{Index, Operation};
 use crate::common::name_generator::{WorktableNameGenerator, is_float};
 use crate::generators::in_memory::InMemoryGenerator;
 use convert_case::{Case, Casing};
@@ -169,7 +169,7 @@ impl InMemoryGenerator {
                         self.gen_non_unique_update(
                             snake_case_name,
                             name,
-                            index_name,
+                            index,
                             idents,
                             indexes_columns.as_ref(),
                             unsized_columns,
@@ -270,6 +270,7 @@ impl InMemoryGenerator {
                     #secondary_events_ident
                 > = Operation::Update(UpdateOperation {
                     id: op_id,
+                    primary_key_events: vec![],
                     secondary_keys_events,
                     bytes: updated_bytes,
                     link,
@@ -449,7 +450,7 @@ impl InMemoryGenerator {
             pub async fn #method_ident<Pk>(&self, row: #query_ident, pk: Pk) -> core::result::Result<(), WorkTableError>
             where #pk_ident: From<Pk>
             {
-                let pk = pk.into();
+                let pk: #pk_ident = pk.into();
                 let op_lock = { #custom_lock };
                 let _guard = LockGuard::new(
                     op_lock,
@@ -489,11 +490,13 @@ 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>>,
     ) -> TokenStream {
+        let by_field = &index.field;
+        let index = &index.name;
         let method_ident = Ident::new(format!("update_{snake_case_name}").as_str(), Span::mixed_site());
 
         let query_ident = Ident::new(format!("{name}Query").as_str(), Span::mixed_site());
@@ -573,19 +576,55 @@ impl InMemoryGenerator {
 
         quote! {
             pub async fn #method_ident(&self, row: #query_ident, by: #by_ident) -> core::result::Result<(), WorkTableError> {
-                let links: Vec<_> = self.0.indexes.#index.get(#by).map(|(_, l)| l.0).collect();
+                // Snapshot the matching rows' primary keys once; the same set
+                // is locked and then processed. Locking one index scan and
+                // processing a fresh second scan would let rows that joined the
+                // range in between be processed without a held lock. The keys
+                // are sorted so concurrent multi-row operations acquire row
+                // locks in one global order (a non-unique index iterates equal
+                // keys in random-discriminator order, which would otherwise
+                // also make the partially-updated subset on a mid-way failure
+                // nondeterministic).
+                let mut pks: Vec<_> = Vec::new();
+                for link in self.0.indexes.#index.get(#by).map(|(_, l)| l.0) {
+                    match self.0.data.select_non_ghosted(link) {
+                        core::result::Result::Ok(r) => pks.push(r.get_primary_key()),
+                        // The row vanished between the index read and the
+                        // resolve; it is simply not part of the snapshot.
+                        core::result::Result::Err(e) if e.is_row_absent() => {}
+                        // Anything else (corrupt page, invalid link, ...) is a
+                        // real storage error, not an empty snapshot slot.
+                        core::result::Result::Err(e) => {
+                            return core::result::Result::Err(WorkTableError::PagesError(e));
+                        }
+                    }
+                }
+                pks.sort_unstable();
+                pks.dedup();
 
                 let mut guards: std::collections::HashMap<_, _> = std::collections::HashMap::new();
-                for link in links.iter() {
-                    let pk = self.0.data.select_non_ghosted(*link)?.get_primary_key().clone();
+                for pk in pks.iter() {
+                    let pk = pk.clone();
                     let op_lock = { #custom_lock };
-                    guards.insert(pk.clone(), LockGuard::new(op_lock, self.0.lock_manager.clone(), pk.clone()));
+                    guards.insert(pk.clone(), LockGuard::new(op_lock, self.0.lock_manager.clone(), pk));
                 }
 
-                let links: Vec<_> = self.0.indexes.#index.get(#by).map(|(_, l)| l.0).collect();
                 let op_id = OperationId::Multi(uuid::Uuid::now_v7());
-                for link in links.into_iter() {
-                    let pk = self.0.data.select_non_ghosted(link)?.get_primary_key().clone();
+                for pk in pks.into_iter() {
+                    // Re-resolve and re-validate under the held lock. The
+                    // query's lock set includes the predicate column, so the
+                    // value read here cannot be rewritten concurrently while
+                    // the lock is held. Rows deleted or updated out of the
+                    // matched range before their lock was acquired are
+                    // skipped; rows that joined the range after the snapshot
+                    // are not touched.
+                    let link: Link = match self.0.primary_index.pk_map.get(&pk) {
+                        Some(v) => v.get().value.into(),
+                        None => continue,
+                    };
+                    if self.0.data.select_non_ghosted(link)?.#by_field != by {
+                        continue;
+                    }
                     let mut bytes = rkyv::to_bytes::<rkyv::rancor::Error>(&row)
                         .map_err(|_| WorkTableError::SerializeError)?;