diff --git a/Cargo.toml b/Cargo.toml index f6c5bef..7bd1a48 100644 --- a/Cargo.toml +++ b/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" diff --git a/codegen/src/generators/in_memory/locks.rs b/codegen/src/generators/in_memory/locks.rs index 4142277..78b40ae 100644 --- a/codegen/src/generators/in_memory/locks.rs +++ b/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 diff --git a/codegen/src/generators/in_memory/queries/delete.rs b/codegen/src/generators/in_memory/queries/delete.rs index 8014249..9f08eb4 100644 --- a/codegen/src/generators/in_memory/queries/delete.rs +++ b/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::>(); - 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(()) } diff --git a/codegen/src/generators/in_memory/queries/locks.rs b/codegen/src/generators/in_memory/queries/locks.rs index 92f8318..b640008 100644 --- a/codegen/src/generators/in_memory/queries/locks.rs +++ b/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::>()).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::>()).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::>()).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::>()).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::>()).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::>()).await; + op_lock } } } diff --git a/codegen/src/generators/in_memory/queries/type.rs b/codegen/src/generators/in_memory/queries/type.rs index 36f9339..35491ce 100644 --- a/codegen/src/generators/in_memory/queries/type.rs +++ b/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 diff --git a/codegen/src/generators/in_memory/queries/update.rs b/codegen/src/generators/in_memory/queries/update.rs index 8e1c2ad..15503b9 100644 --- a/codegen/src/generators/in_memory/queries/update.rs +++ b/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(&self, row: #query_ident, pk: Pk) -> core::result::Result<(), WorkTableError> where #pk_ident: From { - 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>, unsized_fields: Option>, ) -> 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::(&row) .map_err(|_| WorkTableError::SerializeError)?; diff --git a/codegen/src/generators/in_memory/table/impls.rs b/codegen/src/generators/in_memory/table/impls.rs index e3c26d2..b647c68 100644 --- a/codegen/src/generators/in_memory/table/impls.rs +++ b/codegen/src/generators/in_memory/table/impls.rs @@ -148,21 +148,63 @@ impl InMemoryGenerator { let row_type = name_generator.get_row_type_ident(); quote! { + /// Inserts the row if its primary key is absent, updates it + /// otherwise. + /// + /// Concurrency: **system-wide lock-free, not wait-free per call.** + /// A retry is only taken when a concurrent delete or insert + /// flipped this key's existence between the existence check and + /// the operation, so some operation on this key completes on + /// every iteration -- but under sustained adversarial churn on + /// the same key an individual call can retry indefinitely. There + /// is deliberately no retry limit: upsert is semantically + /// infallible for primary-key conflicts, and a limit would trade + /// theoretical starvation for real spurious errors. Each + /// conflicting round yields to the scheduler before retrying so + /// the interfering task can complete. pub async fn upsert(&self, row: #row_type) -> core::result::Result<(), WorkTableError> { let pk = row.get_primary_key(); - let need_to_update = { - if let Some(link) = self.0.primary_index.pk_map.get(&pk) { - true + loop { + let need_to_update = self.0.primary_index.pk_map.get(&pk).is_some(); + if need_to_update { + match self.update(row.clone()).await { + core::result::Result::Ok(_) => return core::result::Result::Ok(()), + // Row was deleted concurrently between the check and the + // update; retry as an insert. + core::result::Result::Err(WorkTableError::NotFound) => { + tokio::task::yield_now().await; + continue; + } + // Row is mid-flight: a concurrent insert publishes + // the primary-key entry before unghosting the row + // data (and insert takes no row lock), and a + // concurrent delete ghosts data it is about to + // unindex. Both are transient; retry. + core::result::Result::Err(WorkTableError::PagesError(e)) if e.is_row_absent() => { + tokio::task::yield_now().await; + continue; + } + core::result::Result::Err(e) => return core::result::Result::Err(e), + } } else { - false + match self.insert(row.clone()) { + core::result::Result::Ok(_) => return core::result::Result::Ok(()), + // Row was inserted concurrently between the check and the + // insert; retry as an update. Secondary-index conflicts are + // real errors and are propagated. Progress is lock-free, + // not wait-free: a retry is only taken when a concurrent + // delete/insert flipped this key's existence between the + // check and the operation, so the system as a whole makes + // progress on every retry, but this call can in principle + // retry unboundedly under sustained same-key churn. + core::result::Result::Err(WorkTableError::PrimaryAlreadyExists) => { + tokio::task::yield_now().await; + continue; + } + core::result::Result::Err(e) => return core::result::Result::Err(e), + } } - }; - if need_to_update { - self.update(row).await?; - } else { - self.insert(row)?; } - core::result::Result::Ok(()) } } } diff --git a/codegen/src/generators/persist/locks.rs b/codegen/src/generators/persist/locks.rs index 90547a2..92a88a0 100644 --- a/codegen/src/generators/persist/locks.rs +++ b/codegen/src/generators/persist/locks.rs @@ -63,6 +63,9 @@ impl PersistGenerator { quote! { impl RowLock for #lock_ident { + fn new() -> Self { + #lock_ident::new() + } #is_locked_fn #lock_fn #with_lock_fn diff --git a/codegen/src/generators/persist/queries/delete.rs b/codegen/src/generators/persist/queries/delete.rs index 0d8fb9d..1482d3d 100644 --- a/codegen/src/generators/persist/queries/delete.rs +++ b/codegen/src/generators/persist/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::persist::PersistGenerator; @@ -141,7 +142,7 @@ impl PersistGenerator { 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) @@ -173,7 +174,9 @@ impl PersistGenerator { } } - 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) @@ -185,10 +188,46 @@ impl PersistGenerator { }; 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::>(); - 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(()) } diff --git a/codegen/src/generators/persist/queries/locks.rs b/codegen/src/generators/persist/queries/locks.rs index e2cc9df..2b8c033 100644 --- a/codegen/src/generators/persist/queries/locks.rs +++ b/codegen/src/generators/persist/queries/locks.rs @@ -58,8 +58,16 @@ impl PersistGenerator { 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 PersistGenerator { 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::>()).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::>()).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::>()).await; + op_lock } } @@ -137,31 +131,21 @@ impl PersistGenerator { 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::>()).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::>()).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::>()).await; + op_lock } } } diff --git a/codegen/src/generators/persist/queries/type.rs b/codegen/src/generators/persist/queries/type.rs index 8a3d2bb..ec10a1f 100644 --- a/codegen/src/generators/persist/queries/type.rs +++ b/codegen/src/generators/persist/queries/type.rs @@ -144,7 +144,7 @@ impl PersistGenerator { .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 diff --git a/codegen/src/generators/persist/queries/update.rs b/codegen/src/generators/persist/queries/update.rs index ce8e6f7..99149ff 100644 --- a/codegen/src/generators/persist/queries/update.rs +++ b/codegen/src/generators/persist/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::persist::PersistGenerator; use convert_case::{Case, Casing}; @@ -169,7 +169,7 @@ impl PersistGenerator { self.gen_non_unique_update( snake_case_name, name, - index_name, + index, idents, indexes_columns.as_ref(), unsized_columns, @@ -265,6 +265,7 @@ impl PersistGenerator { #secondary_events_ident > = Operation::Update(UpdateOperation { id: op_id, + primary_key_events: vec![], secondary_keys_events, bytes: updated_bytes, link, @@ -406,7 +407,7 @@ impl PersistGenerator { pub async fn #method_ident(&self, row: #query_ident, pk: Pk) -> core::result::Result<(), WorkTableError> where #pk_ident: From { - let pk = pk.into(); + let pk: #pk_ident = pk.into(); let op_lock = { #custom_lock }; let _guard = LockGuard::new( op_lock, @@ -446,11 +447,13 @@ impl PersistGenerator { &self, snake_case_name: String, name: &Ident, - index: &Ident, + index: &Index, idents: &[Ident], idx_idents: Option<&Vec>, unsized_fields: Option>, ) -> 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()); @@ -530,19 +533,55 @@ impl PersistGenerator { 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::(&row) .map_err(|_| WorkTableError::SerializeError)?; diff --git a/codegen/src/generators/persist/table/impls.rs b/codegen/src/generators/persist/table/impls.rs index 1e7ca35..eafdc3a 100644 --- a/codegen/src/generators/persist/table/impls.rs +++ b/codegen/src/generators/persist/table/impls.rs @@ -244,21 +244,63 @@ impl PersistGenerator { let row_type = name_generator.get_row_type_ident(); quote! { + /// Inserts the row if its primary key is absent, updates it + /// otherwise. + /// + /// Concurrency: **system-wide lock-free, not wait-free per call.** + /// A retry is only taken when a concurrent delete or insert + /// flipped this key's existence between the existence check and + /// the operation, so some operation on this key completes on + /// every iteration -- but under sustained adversarial churn on + /// the same key an individual call can retry indefinitely. There + /// is deliberately no retry limit: upsert is semantically + /// infallible for primary-key conflicts, and a limit would trade + /// theoretical starvation for real spurious errors. Each + /// conflicting round yields to the scheduler before retrying so + /// the interfering task can complete. pub async fn upsert(&self, row: #row_type) -> core::result::Result<(), WorkTableError> { let pk = row.get_primary_key(); - let need_to_update = { - if let Some(link) = self.0.primary_index.pk_map.get(&pk) { - true + loop { + let need_to_update = self.0.primary_index.pk_map.get(&pk).is_some(); + if need_to_update { + match self.update(row.clone()).await { + core::result::Result::Ok(_) => return core::result::Result::Ok(()), + // Row was deleted concurrently between the check and the + // update; retry as an insert. + core::result::Result::Err(WorkTableError::NotFound) => { + tokio::task::yield_now().await; + continue; + } + // Row is mid-flight: a concurrent insert publishes + // the primary-key entry before unghosting the row + // data (and insert takes no row lock), and a + // concurrent delete ghosts data it is about to + // unindex. Both are transient; retry. + core::result::Result::Err(WorkTableError::PagesError(e)) if e.is_row_absent() => { + tokio::task::yield_now().await; + continue; + } + core::result::Result::Err(e) => return core::result::Result::Err(e), + } } else { - false + match self.insert(row.clone()) { + core::result::Result::Ok(_) => return core::result::Result::Ok(()), + // Row was inserted concurrently between the check and the + // insert; retry as an update. Secondary-index conflicts are + // real errors and are propagated. Progress is lock-free, + // not wait-free: a retry is only taken when a concurrent + // delete/insert flipped this key's existence between the + // check and the operation, so the system as a whole makes + // progress on every retry, but this call can in principle + // retry unboundedly under sustained same-key churn. + core::result::Result::Err(WorkTableError::PrimaryAlreadyExists) => { + tokio::task::yield_now().await; + continue; + } + core::result::Result::Err(e) => return core::result::Result::Err(e), + } } - }; - if need_to_update { - self.update(row).await?; - } else { - self.insert(row)?; } - core::result::Result::Ok(()) } } } @@ -400,7 +442,7 @@ impl PersistGenerator { std::sync::Arc::clone(&self.0.lock_manager), std::sync::Arc::clone(&self.0.primary_index), std::sync::Arc::clone(&self.0.indexes), - )) + ).with_persistence(self.1.vacuum_sink())) } } } diff --git a/codegen/src/generators/read_only/locks.rs b/codegen/src/generators/read_only/locks.rs index 6e3a46a..280afd2 100644 --- a/codegen/src/generators/read_only/locks.rs +++ b/codegen/src/generators/read_only/locks.rs @@ -63,6 +63,9 @@ impl ReadOnlyGenerator { quote! { impl RowLock for #lock_ident { + fn new() -> Self { + #lock_ident::new() + } #is_locked_fn #lock_fn #with_lock_fn diff --git a/src/in_memory/pages.rs b/src/in_memory/pages.rs index 1825582..08a41fb 100644 --- a/src/in_memory/pages.rs +++ b/src/in_memory/pages.rs @@ -448,6 +448,17 @@ impl ExecutionError { pub fn is_vacuumed(&self) -> bool { matches!(self, Self::Vacuumed) } + + /// True when the error means "no row lives at this link (any more)" — + /// the row was deleted, ghosted, vacuumed away, or its page is gone. + /// Snapshot-building code skips such candidates; every other variant is + /// a real storage error and must propagate. + pub fn is_row_absent(&self) -> bool { + matches!( + self, + Self::Ghosted | Self::Deleted | Self::Vacuumed | Self::PageNotFound(_) + ) + } } #[cfg(test)] diff --git a/src/lib.rs b/src/lib.rs index 070cd2f..aead1ab 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -46,7 +46,7 @@ pub mod prelude { AvailableIndex, Difference, IndexError, IndexMap, IndexMultiMap, MultiPairRecreate, PrimaryIndex, TableIndex, TableIndexCdc, TableRow, TableSecondaryIndex, TableSecondaryIndexCdc, TableSecondaryIndexEventsOps, TableSecondaryIndexInfo, UnsizedNode, WorkTable, WorkTableError, vacuum::EmptyDataVacuum, - vacuum::WorkTableVacuum, + vacuum::VacuumPersistence, vacuum::WorkTableVacuum, }; pub use data_bucket::{ DATA_VERSION, DataPage, GENERAL_HEADER_SIZE, GeneralHeader, GeneralPage, INNER_PAGE_SIZE, IndexPage, Interval, diff --git a/src/lock/map.rs b/src/lock/map.rs index 7be0bf7..4e4b538 100644 --- a/src/lock/map.rs +++ b/src/lock/map.rs @@ -39,6 +39,37 @@ where self.map.read().get(key).cloned() } + /// Returns the lock for `key`, inserting one built by `f` if absent. + /// + /// The check and the insert happen under a single write guard. Doing them + /// as separate `get` then `insert` calls is a check-then-act race: two + /// tasks can both observe no entry, both build a lock, and both believe + /// they hold the row. The loser's `insert` returns the winner's lock and + /// can merge into it, but the *winner* already registered its operation on + /// a lock that is no longer in the map, so it never waits for the loser and + /// both proceed into the row at once. + pub fn get_or_insert_with(&self, key: PrimaryKey, f: F) -> Arc> + where + F: FnOnce() -> LockType, + { + // Fast path: the row is usually already locked by someone, and a read + // guard keeps unrelated rows' acquisitions concurrent. The clone happens + // under the guard, so `remove_with_lock_check` (which needs the write + // lock) either runs before we looked or sees our extra strong reference + // and keeps the entry. + if let Some(lock) = self.map.read().get(&key) { + return lock.clone(); + } + let mut map = self.map.write(); + // Re-check: another task can insert between the read and write guards. + if let Some(lock) = map.get(&key) { + return lock.clone(); + } + let lock = Arc::new(tokio::sync::RwLock::new(f())); + map.insert(key, lock.clone()); + lock + } + pub fn remove(&mut self, key: &PrimaryKey) { self.map.write().remove(key); } @@ -51,6 +82,20 @@ where if let Some(lock) = set.get(key).cloned() && let Ok(guard) = lock.try_read() && !guard.is_locked() + // Two strong references means this map entry and our own `lock` + // clone above, and nothing else. Any higher count is a task that + // has already taken this Arc out of `get_or_insert_with` and is + // about to register on it; removing the entry now would let the + // next caller build a *second* lock for the same row, and the two + // would not serialise against each other. + // + // Known trade-off: if that other task is cancelled between taking + // the Arc and registering its operation, nothing re-triggers this + // cleanup and the (unlocked, unused) entry stays in the map until + // the next operation on the same key drops its guard. That leaks at + // most one empty lock per abandoned key and never affects mutual + // exclusion. + && Arc::strong_count(&lock) == 2 { set.remove(key); } diff --git a/src/lock/mod.rs b/src/lock/mod.rs index 578fe00..90432ee 100644 --- a/src/lock/mod.rs +++ b/src/lock/mod.rs @@ -106,6 +106,19 @@ impl Lock { } } + /// A lock born in the released state (`is_locked() == false`, waiting on + /// it returns immediately). Used for placeholder state, e.g. a fresh + /// [`FullRowLock`](crate::lock::FullRowLock) that no operation holds yet. + /// "Released" refers to the acquisition flag, distinguishing it from + /// [`Lock::new`], which starts held by the creating operation. + pub fn new_released(id: u16) -> Self { + Self { + id, + locked: Arc::new(AtomicBool::new(false)), + wakers: Mutex::new(vec![]), + } + } + pub fn id(&self) -> u16 { self.id } diff --git a/src/lock/row_lock.rs b/src/lock/row_lock.rs index e874373..9e8aff4 100644 --- a/src/lock/row_lock.rs +++ b/src/lock/row_lock.rs @@ -6,6 +6,10 @@ use std::sync::Arc; use crate::lock::{Lock, LockGuard, LockMap, LockWait}; pub trait RowLock { + /// Creates a new [`RowLock`] with no columns locked. + fn new() -> Self + where + Self: Sized; /// Checks if any column of this row is locked. fn is_locked(&self) -> bool; /// Creates new [`RowLock`] with all columns locked. @@ -52,6 +56,19 @@ impl FullRowLock { #[allow(clippy::mutable_key_type)] impl RowLock for FullRowLock { + fn new() -> Self + where + Self: Sized, + { + // Placeholder: no operation holds this row yet, so the initial lock is + // born released and any wait on it completes immediately. The id is + // never observed because `lock()` replaces the placeholder before + // handing anything out. + FullRowLock { + l: Arc::new(Lock::new_released(0)), + } + } + fn is_locked(&self) -> bool { self.l.is_locked() } diff --git a/src/persistence/engine.rs b/src/persistence/engine.rs index 9237d1d..52bfefd 100644 --- a/src/persistence/engine.rs +++ b/src/persistence/engine.rs @@ -143,6 +143,9 @@ where } Operation::Update(update) => { self.data.save_data(update.link, update.bytes.as_ref()).await?; + for event in update.primary_key_events { + self.primary_index.process_change_event(event).await?; + } self.secondary_indexes .process_change_events(update.secondary_keys_events) .await @@ -170,16 +173,53 @@ where let (pk_evs, secondary_evs) = batch_op.get_indexes_evs()?; { + let data = &mut self.data; + let primary_index = &mut self.primary_index; + let secondary_indexes = &mut self.secondary_indexes; let mut futs = FuturesUnordered::new(); - futs.push(Either::Left(Either::Right(self.data.save_batch_data(batch_data_op)))); - futs.push(Either::Left(Either::Left( - self.primary_index.process_change_event_batch(pk_evs), - ))); - futs.push(Either::Right( - self.secondary_indexes.process_change_event_batch(secondary_evs), - )); - - while (futs.next().await).is_some() {} + futs.push(Either::Left(Either::Right(async move { + data.save_batch_data(batch_data_op) + .await + .map_err(|e| e.wrap_err("batch data write")) + }))); + futs.push(Either::Left(Either::Left(async move { + primary_index + .process_change_event_batch(pk_evs) + .await + .map_err(|e| e.wrap_err("primary index batch apply")) + }))); + futs.push(Either::Right(async move { + secondary_indexes + .process_change_event_batch(secondary_evs) + .await + .map_err(|e| e.wrap_err("secondary index batch apply")) + })); + + // Drain every future before surfacing errors: `?` on the first + // failure would drop the FuturesUnordered and cancel the remaining + // sub-operations at arbitrary await points, leaving e.g. a data + // page half-written while its index events were abandoned. These + // futures are not cancellation-safe, so let all started work run + // to completion. Every failed component is reported (each error is + // wrapped with its component name above); a mixed outcome means + // the durable state may already be inconsistent, and knowing WHICH + // parts failed is what an operator needs. + let mut errors: Vec = Vec::new(); + while let Some(res) = futs.next().await { + if let Err(e) = res { + errors.push(e); + } + } + if errors.len() == 1 { + return Err(errors.pop().expect("len checked")); + } + if !errors.is_empty() { + let summary = errors.iter().map(|e| format!("{e:#}")).collect::>().join("; "); + return Err(eyre::eyre!( + "batch apply failed in {} sub-operations: {summary}", + errors.len() + )); + } } if let Some(pk_gen_state_update) = batch_op.get_pk_gen_state()? { diff --git a/src/persistence/operation/batch.rs b/src/persistence/operation/batch.rs index d36bf25..914d065 100644 --- a/src/persistence/operation/batch.rs +++ b/src/persistence/operation/batch.rs @@ -246,22 +246,24 @@ where && !id.is_next_for(last_ids.primary_id) && last_ids.primary_id != IndexChangeEventId::default() { - let mut possibly_valid = false; - if id.inner().overflowing_sub(last_ids.primary_id.inner()).0 == 2 { - // TODO: for split sometimes this happens - let ev = prepared_evs.primary_evs.first().unwrap(); - if let ChangeEvent::SplitNode { .. } = ev { - possibly_valid = true - } - if attempts > 8 { - possibly_valid = true - } - } - - if !possibly_valid { - self.ops.extend(ops_to_remove); - return Ok(None); + // Change events are positional (InsertAt/RemoveAt carry node + // indices), so a stream with a missing id must never be applied: + // the disk index would apply later events against node state the + // missing event was supposed to produce. Always defer. A gap is + // transient (the op carrying the missing event has not been + // batched yet) unless an event was discarded after its id was + // assigned — only non-CDC index mutations do that — so a gap + // that persists is a bug upstream of the analyzer; report it + // loudly instead of force-applying and corrupting the file. + if attempts > 8 { + tracing::error!( + "persistence stalled on primary index event gap: last applied {:?}, next available {:?} (attempt {attempts}); an event id was likely consumed without its event being queued", + last_ids.primary_id, + id, + ); } + self.ops.extend(ops_to_remove); + return Ok(None); } let secondary_first = prepared_evs.secondary_evs.first_evs(); for (index, id) in secondary_first { @@ -272,19 +274,16 @@ where && !id.is_next_for(*last) && *last != IndexChangeEventId::default() { - let mut possibly_valid = false; - if id.inner().overflowing_sub(last.inner()).0 == 2 { - // TODO: for split sometimes this happens - possibly_valid = prepared_evs.secondary_evs.is_first_ev_is_split(index); - if attempts > 8 { - possibly_valid = true - } - } - - if !possibly_valid { - self.ops.extend(ops_to_remove); - return Ok(None); + // Same rule as the primary index above: never apply a gapped + // stream, defer until the missing event arrives, and report + // a persistent gap as the bug it is. + if attempts > 8 { + tracing::error!( + "persistence stalled on secondary index {index:?} event gap: last applied {last:?}, next available {id:?} (attempt {attempts}); an event id was likely consumed without its event being queued", + ); } + self.ops.extend(ops_to_remove); + return Ok(None); } } } diff --git a/src/persistence/operation/operation.rs b/src/persistence/operation/operation.rs index f45e65d..36af829 100644 --- a/src/persistence/operation/operation.rs +++ b/src/persistence/operation/operation.rs @@ -11,7 +11,7 @@ use crate::persistence::{OperationId, OperationType}; #[derive(Clone, Debug)] pub enum Operation { Insert(InsertOperation), - Update(UpdateOperation), + Update(UpdateOperation), Delete(DeleteOperation), Acknowledge(AcknowledgeOperation), } @@ -72,7 +72,7 @@ impl Operation Option<&Vec>>> { match &self { Operation::Insert(insert) => Some(&insert.primary_key_events), - Operation::Update(_) => None, + Operation::Update(update) => Some(&update.primary_key_events), Operation::Delete(delete) => Some(&delete.primary_key_events), Operation::Acknowledge(ack) => Some(&ack.primary_key_events), } @@ -120,8 +120,9 @@ pub struct InsertOperation { } #[derive(Clone, Debug)] -pub struct UpdateOperation { +pub struct UpdateOperation { pub id: OperationId, + pub primary_key_events: Vec>>, pub secondary_keys_events: SecondaryKeys, pub bytes: Vec, pub link: Link, diff --git a/src/persistence/operation/util.rs b/src/persistence/operation/util.rs index 6d1de30..405a9ce 100644 --- a/src/persistence/operation/util.rs +++ b/src/persistence/operation/util.rs @@ -3,8 +3,6 @@ use indexset::cdc::change::{self, ChangeEvent}; use indexset::core::pair::Pair; use std::fmt::Debug; -pub const MAX_CHECK_DEPTH: usize = 30; - pub fn validate_events(evs: &mut Vec>>) -> Vec>> where T: Debug, @@ -36,7 +34,12 @@ fn validate_events_iteration(evs: &[ChangeEvent>]) -> (Vec(evs: &[ChangeEvent>]) -> (Vec ChangeEvent> { + ChangeEvent::InsertAt { + event_id: id.into(), + max_value: Pair { + key: id, + value: Link::default(), + }, + value: Pair { + key: id, + value: Link::default(), + }, + index: 0, + } + } + + #[test] + fn detects_gap_behind_long_contiguous_tail() { + // Interior gap (100..=139, then 1000..=1049) whose tail is longer than + // any bounded scan window: everything after the gap must be deferred. + let mut evs: Vec<_> = (100..140).map(insert_at).collect(); + evs.extend((1000..1050).map(insert_at)); + + let removed = validate_events(&mut evs); + + assert_eq!(evs.len(), 40); + assert!(evs.iter().all(|ev| ev.id() < 140.into())); + assert_eq!(removed.len(), 50); + assert!(removed.iter().all(|ev| ev.id() >= 1000.into())); + } + + #[test] + fn keeps_gapless_stream_untouched() { + let mut evs: Vec<_> = (100..200).map(insert_at).collect(); + let removed = validate_events(&mut evs); + assert!(removed.is_empty()); + assert_eq!(evs.len(), 100); + } +} diff --git a/src/persistence/space/data.rs b/src/persistence/space/data.rs index 9bb72b4..4685a97 100644 --- a/src/persistence/space/data.rs +++ b/src/persistence/space/data.rs @@ -131,8 +131,16 @@ where .cloned() .collect::>(); - if let Some(max) = ids_to_create.last() { - self.last_page_id = *max; + // `page_ids` iterates a HashMap, so `ids_to_create` is unordered: + // taking `.last()` here picked an arbitrary created page, and a batch + // creating several pages could leave `last_page_id` below a page that + // now exists. The next batch touching that page would see it as "new" + // and re-create it zero-filled, wiping the rows persisted before. + if let Some(max) = ids_to_create.iter().max() { + // High-water mark: every id in `ids_to_create` is > last_page_id by + // construction, but state the monotonic invariant directly so a + // future refactor of the filter above cannot regress it. + self.last_page_id = self.last_page_id.max(*max); } let created_pages = ids_to_create .into_iter() diff --git a/src/persistence/task.rs b/src/persistence/task.rs index e6fcc91..7fefe41 100644 --- a/src/persistence/task.rs +++ b/src/persistence/task.rs @@ -1,12 +1,13 @@ -use std::collections::{HashMap, HashSet}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::fmt::Debug; use std::hash::Hash; use std::marker::PhantomData; use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicU16, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::time::Duration; use data_bucket::page::PageId; +use parking_lot::Mutex as ParkingMutex; use tokio::sync::Notify; use worktable_codegen::worktable; @@ -14,6 +15,7 @@ use crate::persistence::PersistenceEngine; use crate::persistence::operation::{BatchInnerRow, BatchInnerWorkTable, BatchOperation, OperationId, PosByOpIdQuery}; use crate::prelude::*; use crate::util::OptimizedVec; +use crate::vacuum::VacuumPersistence; worktable! ( name: QueueInner, @@ -271,32 +273,51 @@ where #[derive(Debug)] pub struct Queue { - queue: lockfree::queue::Queue>, + // Not `lockfree::queue::Queue`: its `Removable::empty` materializes the + // element type via `mem::uninitialized`, which aborts at runtime for + // `Operation` layouts that reject uninit bytes. The queue has a single + // consumer (the engine task), so a mutexed deque is uncontended here. + queue: ParkingMutex>>, notify: Notify, - len: Arc, + // usize, not u16: the queue is unbounded and a 16-bit counter wraps at + // 65_536 queued operations, making the wait triggers see an "empty" + // queue that still holds work. + len: Arc, } impl Queue { pub fn new() -> Self { Self { - queue: lockfree::queue::Queue::new(), + queue: ParkingMutex::new(VecDeque::new()), notify: Notify::new(), - len: Arc::new(AtomicU16::new(0)), + len: Arc::new(AtomicUsize::new(0)), } } pub fn push(&self, value: Operation) { self.len.fetch_add(1, Ordering::Release); - self.queue.push(value); + self.queue.lock().push_back(value); self.notify.notify_one(); } - pub async fn pop(&self) -> Operation { + /// Pops the next operation, marking `in_progress` `true` before the queue + /// length is decremented. The store must precede `len.fetch_sub` so that a + /// waiter that reads `len() == 0` (`Acquire`) is guaranteed to observe + /// `in_progress == true` for the popped-but-unprocessed operation; + /// otherwise `wait_for_ops` can return while that operation is in flight. + pub async fn pop_marking_in_progress( + &self, + in_progress: &AtomicBool, + ) -> Operation { loop { // Drain values - if let Some(value) = self.queue.pop() { - self.len.fetch_sub(1, Ordering::Release); - return value; + { + let mut queue = self.queue.lock(); + if let Some(value) = queue.pop_front() { + in_progress.store(true, Ordering::Release); + self.len.fetch_sub(1, Ordering::Release); + return value; + } } // Wait for values to be available @@ -305,7 +326,7 @@ impl Queue Option> { - if let Some(v) = self.queue.pop() { + if let Some(v) = self.queue.lock().pop_front() { self.len.fetch_sub(1, Ordering::Release); Some(v) } else { @@ -314,20 +335,40 @@ impl Queue impl Iterator> { - let iter_count = self.len.clone(); - self.queue.pop_iter().inspect(move |_| { - iter_count.fetch_sub(1, Ordering::Release); - }) + std::iter::from_fn(|| self.immediate_pop()) } pub fn len(&self) -> usize { - self.len.load(Ordering::Acquire) as usize + self.len.load(Ordering::Acquire) + } +} + +impl VacuumPersistence + for Queue +where + PrimaryKeyGenState: Send, + PrimaryKey: Send, + SecondaryKeys: Send, +{ + fn apply_move( + &self, + bytes: Vec, + new_link: Link, + primary_key_events: Vec>>, + secondary_keys_events: SecondaryKeys, + ) { + self.push(Operation::Update(UpdateOperation { + id: OperationId::Single(uuid::Uuid::now_v7()), + primary_key_events, + secondary_keys_events, + bytes, + link: new_link, + })); } } #[derive(Debug)] pub struct PersistenceTask { - #[allow(dead_code)] engine_task_handle: tokio::task::AbortHandle, queue: Arc>, analyzer_inner_wt: Arc, @@ -336,6 +377,35 @@ pub struct PersistenceTask, } +impl Drop + for PersistenceTask +{ + /// Aborts the engine task so it cannot outlive the table it persists. + /// Without this the detached task keeps running on the runtime after the + /// table is dropped, and a re-opened table can read the same files while + /// the old engine is still writing them. + /// + /// The abort only happens when the engine is provably idle (queue and + /// analyzer empty, no operation in flight) — that is the normal state + /// after `wait_for_ops`, and an idle task is parked at the queue pop, an + /// await point where cancellation is clean. Aborting a *busy* engine + /// would cancel persistence futures that are not cancellation-safe (a + /// data page could be left half-written while its index events are + /// abandoned), so a busy engine is left running and reported instead: + /// callers must drain with `wait_for_ops` before dropping. A proper + /// `close()` lifecycle (drain, join, surface terminal errors) is the + /// long-term replacement for this heuristic. + fn drop(&mut self) { + if self.check_wait_triggers() { + self.engine_task_handle.abort(); + } else { + tracing::error!( + "PersistenceTask dropped with work in flight; the engine task keeps running detached. Call wait_for_ops() before dropping to guarantee a clean shutdown." + ); + } + } +} + impl PersistenceTask { @@ -343,6 +413,17 @@ impl self.queue.push(op); } + /// Returns a sink that lets vacuum queue persistence operations for row + /// moves into this task's operation queue. + pub fn vacuum_sink(&self) -> Arc> + where + PrimaryKeyGenState: Send + Sync + 'static, + PrimaryKey: Send + Sync + 'static, + SecondaryKeys: Send + Sync + 'static, + { + self.queue.clone() + } + pub fn run_engine(mut engine: E) -> Self where E: PersistenceEngine + Send + 'static, @@ -366,11 +447,12 @@ impl let op = if let Some(next_op) = engine_queue.immediate_pop() { Some(next_op) } else if analyzer.len() == 0 { - engine_progress_notify.notify_waiters(); task_analyzer_in_progress.store(false, Ordering::Release); - let res = Some(engine_queue.pop().await); - task_analyzer_in_progress.store(true, Ordering::Release); - res + engine_progress_notify.notify_waiters(); + // The pop sets the flag back to `true` atomically with the + // dequeue, so waiters never observe an empty queue with an + // idle analyzer while an operation is in flight. + Some(engine_queue.pop_marking_in_progress(&task_analyzer_in_progress).await) } else { None }; diff --git a/src/table/mod.rs b/src/table/mod.rs index 8385ea5..cb54aaa 100644 --- a/src/table/mod.rs +++ b/src/table/mod.rs @@ -179,7 +179,7 @@ where let link = self.data.insert(row.clone()).map_err(WorkTableError::PagesError)?; if self.primary_index.insert_checked(pk.clone(), link).is_none() { self.data.delete(link).map_err(WorkTableError::PagesError)?; - return Err(WorkTableError::AlreadyExists("Primary".to_string())); + return Err(WorkTableError::PrimaryAlreadyExists); }; if let Err(e) = self.indexes.save_row(row.clone(), link) { return match e { @@ -237,7 +237,7 @@ where if let Err(e) = self.data.delete(link) { return (None, Err(WorkTableError::PagesError(e))); } - return (None, Err(WorkTableError::AlreadyExists("Primary".to_string()))); + return (None, Err(WorkTableError::PrimaryAlreadyExists)); }; let primary_key_events = convert_change_events(primary_key_events); @@ -516,6 +516,8 @@ pub enum WorkTableError { NotFound, #[display("Value already exists for `{}` index", _0)] AlreadyExists(#[error(not(source))] String), + #[display("Row with this primary key already exists")] + PrimaryAlreadyExists, SerializeError, SecondaryIndexError, PrimaryUpdateTry, diff --git a/src/table/vacuum/mod.rs b/src/table/vacuum/mod.rs index 5396c2e..21ab55e 100644 --- a/src/table/vacuum/mod.rs +++ b/src/table/vacuum/mod.rs @@ -1,5 +1,9 @@ use async_trait::async_trait; +use data_bucket::Link; +use indexset::cdc::change::ChangeEvent; +use indexset::core::pair::Pair; + use crate::vacuum::fragmentation_info::FragmentationInfo; mod fragmentation_info; @@ -10,6 +14,31 @@ mod vacuum; pub use manager::{VacuumManager, VacuumManagerConfig}; pub use vacuum::EmptyDataVacuum; +/// Sink for persisting vacuum row moves. +/// +/// Vacuum relocates rows between data pages, which changes the [`Link`] stored +/// in the primary and secondary indexes. On persisted tables those index +/// mutations must go through the CDC event stream, and the moved row bytes must +/// be written at the new link — otherwise the on-disk state goes stale and the +/// event-id sequence gets a permanent gap that stalls persistence. Implementors +/// receive everything needed to queue a proper persistence operation for one +/// moved row. +/// Not intended for downstream implementation: this is macro-support API for +/// the generated persisted-table `vacuum()`, and it leaks low-level CDC event +/// types. Hidden from docs; semver stability is not promised for it. +#[doc(hidden)] +pub trait VacuumPersistence: Send + Sync { + /// Queue a persistence operation for a row moved to `new_link`, carrying + /// the row bytes and the CDC events produced by the index updates. + fn apply_move( + &self, + bytes: Vec, + new_link: Link, + primary_key_events: Vec>>, + secondary_keys_events: SecondaryEvents, + ); +} + /// Trait for unifying different [`WorkTable`] related [`EmptyDataVacuum`]'s. /// /// [`WorkTable`]: crate::prelude::WorkTable diff --git a/src/table/vacuum/vacuum.rs b/src/table/vacuum/vacuum.rs index f5b87e4..1cb3d8b 100644 --- a/src/table/vacuum/vacuum.rs +++ b/src/table/vacuum/vacuum.rs @@ -18,15 +18,18 @@ use rkyv::{Archive, Deserialize, Serialize}; use crate::in_memory::{ArchivedRowWrapper, DataPages, RowWrapper, StorableRow}; use crate::lock::{Lock, LockMap, RowLock}; use crate::prelude::{OffsetEqLink, TablePrimaryKey}; +use crate::vacuum::VacuumPersistence; use crate::vacuum::VacuumStats; use crate::vacuum::WorkTableVacuum; use crate::vacuum::fragmentation_info::FragmentationInfo; -use crate::{AvailableIndex, PrimaryIndex, TableIndex, TableRow, TableSecondaryIndex, TableSecondaryIndexCdc}; +use crate::{ + AvailableIndex, PrimaryIndex, TableIndex, TableIndexCdc, TableRow, TableSecondaryIndex, TableSecondaryIndexCdc, +}; use async_trait::async_trait; use ordered_float::OrderedFloat; use rkyv::api::high::HighDeserializer; -#[derive(Debug)] +#[derive(derive_more::Debug)] pub struct EmptyDataVacuum< Row, PrimaryKey, @@ -51,6 +54,11 @@ pub struct EmptyDataVacuum< primary_index: Arc>, secondary_indexes: Arc, + /// Persistence sink for row moves. `None` for in-memory tables; persisted + /// tables must set it so index updates go through CDC and reach disk. + #[debug(ignore)] + persistence: Option>>, + phantom_data: PhantomData<(SecondaryEvents, AvailableTypes, AvailableIndexes)>, } @@ -108,10 +116,19 @@ where lock_manager, primary_index, secondary_indexes, + persistence: None, phantom_data: PhantomData, } } + /// Attaches a persistence sink. Index updates for moved rows then use the + /// CDC mutation variants and their events are queued as persistence + /// operations. Required for persisted tables. + pub fn with_persistence(mut self, sink: Arc>) -> Self { + self.persistence = Some(sink); + self + } + async fn defragment(&self) -> VacuumStats { let now = Instant::now(); @@ -262,7 +279,7 @@ where let new_link = to_page .save_raw_row(&raw_data) .expect("page is not full as checked on links collection"); - self.update_index_after_move(pk.clone(), from_link.0, new_link); + self.update_index_after_move(pk.clone(), from_link.0, new_link, raw_data); lock.unlock(); self.lock_manager.remove_with_lock_check(&pk); @@ -273,43 +290,39 @@ where async fn full_row_lock(&self, pk: &PrimaryKey) -> Arc { let lock_id = self.lock_manager.next_id(); - if let Some(lock) = self.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::>()).await; - - op_lock - } else { - #[allow(clippy::mutable_key_type)] - let (lock, op_lock) = LockType::with_lock(lock_id); - let lock = Arc::new(tokio::sync::RwLock::new(lock)); - let mut guard = lock.write().await; - if let Some(old_lock) = self.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::>()).await; - } - - op_lock - } + // One atomic acquire, no check-then-act: see LockMap::get_or_insert_with. + let lock = self.lock_manager.get_or_insert_with(pk.clone(), LockType::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::>()).await; + + op_lock } - fn update_index_after_move(&self, pk: PrimaryKey, old_link: Link, new_link: Link) { + fn update_index_after_move(&self, pk: PrimaryKey, old_link: Link, new_link: Link, raw_data: Vec) { let row = self .data_pages .select(new_link) .expect("should exist as link was moved correctly"); - self.secondary_indexes - .reinsert_row(row.clone(), old_link, row, new_link) - .expect("should be ok as index were no violated"); - self.primary_index.insert(pk.clone(), new_link); + if let Some(persistence) = &self.persistence { + // Persisted table: mutate indexes through the CDC variants and queue + // the events with the moved bytes, so the on-disk state follows the + // move and the event-id stream stays gapless. + let (secondary_keys_events, res) = + self.secondary_indexes + .reinsert_row_cdc(row.clone(), old_link, row, new_link); + res.expect("should be ok as index were no violated"); + let (_, primary_key_events) = self.primary_index.insert_cdc(pk.clone(), new_link); + persistence.apply_move(raw_data, new_link, primary_key_events, secondary_keys_events); + } else { + self.secondary_indexes + .reinsert_row(row.clone(), old_link, row, new_link) + .expect("should be ok as index were no violated"); + self.primary_index.insert(pk.clone(), new_link); + } } } diff --git a/tests/data/expected/space_index/indexset/process_insert_at_big_amount.wt.idx b/tests/data/expected/space_index/indexset/process_insert_at_big_amount.wt.idx index 681ca43..bad6e69 100644 Binary files a/tests/data/expected/space_index/indexset/process_insert_at_big_amount.wt.idx and b/tests/data/expected/space_index/indexset/process_insert_at_big_amount.wt.idx differ diff --git a/tests/persistence/bulk_load_stall.rs b/tests/persistence/bulk_load_stall.rs new file mode 100644 index 0000000..cf55b02 --- /dev/null +++ b/tests/persistence/bulk_load_stall.rs @@ -0,0 +1,107 @@ +use std::collections::HashMap; +use std::time::Duration; + +use tokio::time::timeout; + +use crate::remove_dir_if_exists; +use worktable::prelude::PersistedWorkTable; +use worktable::prelude::*; +use worktable::worktable; + +worktable!( + name: BulkLoadStall, + persist: true, + columns: { + id: u64 primary_key autoincrement, + test: i64, + another: u64, + exchange: String, + }, + indexes: { + test_idx: test unique, + another_idx: another, + exchange_idx: exchange, + }, +); + +/// Regression test for an index-space batching bug: an unthrottled bulk +/// insert+delete used to panic the persistence engine task with "page should +/// be available in table of contents" (process_change_event_batch), after +/// which wait_for_ops hung forever; the timeout below turns that hang into a +/// failure. +/// +/// Page-grouped batch collection puts the delete events (ids far ahead of the +/// batched inserts) into the same batch as the first data page's inserts, so +/// the prepared event stream has an interior id gap. `validate_events` used to +/// scan only 30 events back from the end; the 50 contiguous delete events hid +/// the gap, the batch was applied with the hole, and the on-disk index lost +/// track of node max transitions carried by the missing events. The scan is +/// now unbounded, so the tail after the gap is deferred until the missing +/// events arrive. +/// +/// The same unthrottled load also used to hit a second bug: a lagging batch +/// creating several data pages at once picked `last_page_id` from unordered +/// HashMap keys in `save_batch_data`, so a later batch could re-create an +/// existing page zero-filled and this test's reload phase read back zeroed +/// rows. Both fixes are needed for this test to be stable. +#[test] +fn test_bulk_insert_delete_persistence() { + let config = DiskConfig::new_with_table_name( + "tests/data/bulk_load_stall/persisted", + BulkLoadStallWorkTable::name_snake_case(), + BulkLoadStallWorkTable::version(), + ); + + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_io() + .enable_time() + .build() + .unwrap(); + + runtime.block_on(async { + remove_dir_if_exists("tests/data/bulk_load_stall/persisted".to_string()).await; + + let mut rows = HashMap::new(); + { + let engine = BulkLoadStallPersistenceEngine::new(config.clone()).await.unwrap(); + let table = BulkLoadStallWorkTable::load(engine).await.unwrap(); + + for i in 0..1000i64 { + let row = BulkLoadStallRow { + id: table.get_next_pk().into(), + test: i, + another: i as u64, + exchange: format!("test{i}"), + }; + let id = row.id; + table.insert(row.clone()).unwrap(); + rows.insert(id, row); + } + + let mut ids: Vec<_> = rows.keys().cloned().collect(); + ids.sort_unstable(); + let deleted: Vec = ids.into_iter().take(50).collect(); + for id in &deleted { + table.delete(*id).await.unwrap(); + } + + timeout(Duration::from_secs(30), table.wait_for_ops()) + .await + .expect("persistence stalled on bulk insert+delete"); + + for id in &deleted { + rows.remove(id); + } + } + { + let engine = BulkLoadStallPersistenceEngine::new(config.clone()).await.unwrap(); + let table = BulkLoadStallWorkTable::load(engine).await.unwrap(); + + assert_eq!(table.select_all().execute().unwrap().len(), rows.len()); + for (id, expected) in &rows { + assert_eq!(table.select(*id).as_ref(), Some(expected)); + } + } + }) +} diff --git a/tests/persistence/mod.rs b/tests/persistence/mod.rs index dc766d1..bdcfa80 100644 --- a/tests/persistence/mod.rs +++ b/tests/persistence/mod.rs @@ -2,6 +2,7 @@ use worktable::prelude::PersistedWorkTable; use worktable::prelude::*; use worktable::worktable; +mod bulk_load_stall; mod concurrent; mod failure; mod index_page; @@ -9,6 +10,7 @@ mod read; mod space_index; mod sync; mod toc; +mod vacuum; #[cfg(feature = "s3-support")] mod s3; diff --git a/tests/persistence/vacuum.rs b/tests/persistence/vacuum.rs new file mode 100644 index 0000000..1e510e3 --- /dev/null +++ b/tests/persistence/vacuum.rs @@ -0,0 +1,136 @@ +use std::collections::HashMap; +use std::time::Duration; + +use tokio::time::timeout; + +use crate::remove_dir_if_exists; +use worktable::prelude::PersistedWorkTable; +use worktable::prelude::*; +use worktable::worktable; + +worktable!( + name: VacuumPersist, + persist: true, + columns: { + id: u64 primary_key autoincrement, + test: i64, + another: u64, + exchange: String, + }, + indexes: { + test_idx: test unique, + another_idx: another, + exchange_idx: exchange, + }, +); + +/// Vacuum on a persisted table must push the row moves through the CDC +/// persistence stream: the moved links have to reach the on-disk indexes, and +/// the event-id sequence must stay gapless so persistence does not stall. +#[test] +fn test_vacuum_on_persisted_table_survives_reload() { + let config = DiskConfig::new_with_table_name( + "tests/data/vacuum/persisted", + VacuumPersistWorkTable::name_snake_case(), + VacuumPersistWorkTable::version(), + ); + + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_io() + .enable_time() + .build() + .unwrap(); + + runtime.block_on(async { + remove_dir_if_exists("tests/data/vacuum/persisted".to_string()).await; + + let mut rows = HashMap::new(); + let deleted: Vec; + { + let engine = VacuumPersistPersistenceEngine::new(config.clone()).await.unwrap(); + let table = VacuumPersistWorkTable::load(engine).await.unwrap(); + + // row is ~40 bytes so ~409 rows per page; use multiple pages so + // defragment has non-current pages to move rows from. + for i in 0..1000i64 { + let row = VacuumPersistRow { + id: table.get_next_pk().into(), + test: i, + another: i as u64, + exchange: format!("test{i}"), + }; + let id = row.id; + table.insert(row.clone()).unwrap(); + rows.insert(id, row); + } + + let mut ids: Vec<_> = rows.keys().cloned().collect(); + ids.sort_unstable(); + deleted = ids.into_iter().take(50).collect(); + for id in &deleted { + table.delete(*id).await.unwrap(); + } + + timeout(Duration::from_secs(30), table.wait_for_ops()) + .await + .expect("persistence should catch up before vacuum"); + + let vacuum = table.vacuum(); + let stats = vacuum.vacuum().await.unwrap(); + assert!(stats.pages_freed > 0, "vacuum should have moved rows off a page"); + + // Insert after vacuum: these operations carry event ids issued + // after the moves, so if vacuum consumed ids without queueing the + // events, the batch validator defers on the gap forever. + for i in 1000..1100i64 { + let row = VacuumPersistRow { + id: table.get_next_pk().into(), + test: i, + another: i as u64, + exchange: format!("test{i}"), + }; + let id = row.id; + table.insert(row.clone()).unwrap(); + rows.insert(id, row); + if i % 50 == 49 { + timeout(Duration::from_secs(30), table.wait_for_ops()) + .await + .expect("persistence stalled after vacuum on persisted table"); + } + } + + // Without CDC-aware vacuum this stalls forever: the moved links + // never reach the persistence stream while their event ids are + // consumed, leaving a permanent gap the batch validator defers on. + timeout(Duration::from_secs(30), table.wait_for_ops()) + .await + .expect("persistence stalled after vacuum on persisted table"); + + for id in &deleted { + rows.remove(id); + } + for (id, expected) in &rows { + assert_eq!(table.select(*id).as_ref(), Some(expected)); + } + } + { + let engine = VacuumPersistPersistenceEngine::new(config.clone()).await.unwrap(); + let table = VacuumPersistWorkTable::load(engine).await.unwrap(); + + assert_eq!(table.select_all().execute().unwrap().len(), rows.len()); + for (id, expected) in &rows { + assert_eq!(table.select(*id).as_ref(), Some(expected)); + // Secondary indexes must follow the moved links too. + assert_eq!(table.select_by_test(expected.test).as_ref(), Some(expected)); + assert_eq!( + table.select_by_exchange(expected.exchange.clone()).execute().unwrap(), + vec![expected.clone()] + ); + } + for id in &deleted { + assert_eq!(table.select(*id), None); + } + } + }) +} diff --git a/tests/worktable/mod.rs b/tests/worktable/mod.rs index 9122137..b4db304 100644 --- a/tests/worktable/mod.rs +++ b/tests/worktable/mod.rs @@ -12,6 +12,7 @@ mod nid; mod option; mod tuple_primary_key; mod unsized_; +mod upsert; mod uuid; mod vacuum; mod with_enum; diff --git a/tests/worktable/upsert.rs b/tests/worktable/upsert.rs new file mode 100644 index 0000000..47b87cb --- /dev/null +++ b/tests/worktable/upsert.rs @@ -0,0 +1,101 @@ +use std::sync::Arc; +use std::time::Duration; + +use tokio::time::timeout; +use worktable::prelude::*; +use worktable::worktable; + +worktable!( + name: UpsertChurn, + columns: { + id: u64 primary_key, + val: u64, + }, +); + +/// Upsert is system-wide lock-free but not wait-free per call: a retry is +/// taken exactly when a concurrent delete/insert flips the key's existence +/// between the check and the operation. This test drives sustained +/// adversarial churn on ONE key while several tasks upsert it and asserts +/// that every upsert completes without surfacing a spurious conflict error, +/// under a timeout that turns pathological starvation into a failure +/// instead of a hang. +/// Moderate tier: 40/40 stable when run solo, but under full-suite parallel +/// load the pre-existing #169 stall still fires (~1/30 suite runs), so both +/// tiers stay ignored until that lands. Run with `-- --ignored` to validate +/// the upsert retry behavior. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "exposes the pre-existing #169 lock-free-insert stall under suite load"] +async fn upsert_completes_under_same_key_churn() { + churn_run(100, 200).await; +} + +/// Intense variant: reliably exposes pre-existing engine races under extreme +/// same-key churn that are unrelated to the upsert retry loop (raw `insert` +/// takes no row lock and publishes the pk entry before unghosting the data; +/// under saturation a churn round can stall past the timeout). Tracked in the +/// lock-free-insert issue; un-ignore when that lands. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "exposes pre-existing lock-free-insert races under extreme same-key churn"] +async fn upsert_completes_under_extreme_same_key_churn() { + churn_run(5_000, 2_000).await; +} + +async fn churn_run(churn_flips: u64, upserts_per_task: u64) { + #[allow(non_snake_case)] + let CHURN_FLIPS = churn_flips; + #[allow(non_snake_case)] + let UPSERTS_PER_TASK = upserts_per_task; + let table = Arc::new(UpsertChurnWorkTable::default()); + const KEY: u64 = 7; + + let churn = { + let table = table.clone(); + tokio::spawn(async move { + for i in 0..CHURN_FLIPS { + // Flip the key's existence as fast as possible through the + // locked operations. (Raw `insert` is deliberately not used + // here: it takes no row lock and publishes the pk entry + // before unghosting the data, which trips unrelated + // pre-existing races tracked separately in the issue on + // lock-free insert vs locked mutations.) + table + .upsert(UpsertChurnRow { id: KEY, val: i }) + .await + .expect("churn upsert must not fail"); + let _ = table.delete(KEY).await; + } + }) + }; + + let mut upserters = Vec::new(); + for w in 0..4u64 { + let table = table.clone(); + upserters.push(tokio::spawn(async move { + for i in 0..UPSERTS_PER_TASK { + table + .upsert(UpsertChurnRow { + id: KEY, + val: w * 10_000 + i, + }) + .await + .expect("upsert must never surface a primary-key conflict"); + } + })); + } + + timeout(Duration::from_secs(60), churn) + .await + .expect("churn task starved") + .unwrap(); + for handle in upserters { + timeout(Duration::from_secs(60), handle) + .await + .expect("upserter starved") + .unwrap(); + } + + // Quiesced: a final upsert must land and be visible. + table.upsert(UpsertChurnRow { id: KEY, val: 424_242 }).await.unwrap(); + assert_eq!(table.select(KEY).map(|r| r.val), Some(424_242)); +}