diff --git a/crates/datastore/src/locking_tx_datastore/mut_tx.rs b/crates/datastore/src/locking_tx_datastore/mut_tx.rs index 5f9e48c1783..21ce7cf486d 100644 --- a/crates/datastore/src/locking_tx_datastore/mut_tx.rs +++ b/crates/datastore/src/locking_tx_datastore/mut_tx.rs @@ -4072,8 +4072,19 @@ pub(super) fn insert<'a, const GENERATE: bool>( // the temporary table borrow ends before cleanup or mutation. let row = unsafe { tx_table.get_row_ref_from_pinned_unchecked(tx_pinned_row.clone()) }; unsafe { Table::find_same_row_from_ref(commit_table, &row, tx_row_hash) } - }?; - if let (_, Some(commit_ptr)) = existing_row { + }; + let commit_ptr = match existing_row { + Ok((_, commit_ptr)) => commit_ptr, + Err(error) => { + // Lookup can fault after confirmation has indexed and counted + // the new row. Restore tx state before propagating the error. + // SAFETY: Lookup did not delete this row; tx_pinned_row + // still retains its page and blob pins, so cleanup cannot fault. + unsafe { tx_table.delete_unchecked(tx_blob_store, tx_pinned_row) }; + return Err(error.into()); + } + }; + if let Some(commit_ptr) = commit_ptr { // (insert_undelete) // ----------------------------------------------------- // If `row` was already present in the committed state, @@ -4104,14 +4115,16 @@ pub(super) fn insert<'a, const GENERATE: bool>( // tx_pinned_row retains its page and blob pins. unsafe { tx_table.delete_unchecked(tx_blob_store, tx_pinned_row) }; + // Loading the existing row can fault on its blobs. Finish this + // before undeleting it, so an error preserves any prior deletion. + // SAFETY: `find_same_row` found a live row in `commit_table`. + let row_ref = unsafe { commit_table.get_row_ref_unchecked(commit_blob_store, commit_ptr) }?; + // It's possible that `row` appears in the committed state, // but is marked as deleted. // In this case, undelete it, so it remains in the committed state. delete_table.remove(commit_ptr); - // No new row was inserted, but return `committed_ptr`. - // SAFETY: `find_same_row` told us that `ptr` refers to a valid row in `commit_table`. - let row_ref = unsafe { commit_table.get_row_ref_unchecked(commit_blob_store, commit_ptr) }?; return ok(RowRefInsertion::Existed(row_ref)); } @@ -4271,11 +4284,17 @@ impl MutTxId { // If the new row is the same as the old, // skip the update altogether to match the semantics of `Self::insert`. // + // A comparison fault must reach the unconfirmed-row cleanup below. // SAFETY: // 1. `tx_table` is derived from `commit_table` so they have the same layouts. // 2. `old_ptr` was found in an index of `commit_table`, so we know it is valid. // 3. we just inserted `tx_row_ptr` into `tx_table`, so we know it is valid. - if unsafe { Table::eq_row_in_page(commit_table, old_ptr, tx_table, tx_pinned_row.pointer()) }? { + let same_row = + match unsafe { Table::eq_row_in_page(commit_table, old_ptr, tx_table, tx_pinned_row.pointer()) } { + Ok(same_row) => same_row, + Err(error) => throw!(error), + }; + if same_row { // SAFETY: `tx_table.is_row_present(tx_row_ptr)` holds, as noted in 3. unsafe { tx_table.delete_internal_skip_pointer_map(tx_blob_store, tx_pinned_row) }; // SAFETY: `commit_table.is_row_present(old_ptr)` holds, as noted in 2. @@ -4342,6 +4361,8 @@ impl MutTxId { // If we have an identical deleted row in the committed state, // we need to undelete it, just like in `Self::insert`. // The same note (`insert_undelete`) there re. MVCC applies here as well. + // Compare after confirmation has released the old tx row's pins + // to avoid requiring all three rows' data to be resident at once. // // SAFETY: // 1. `tx_table` is derived from `commit_table` so they have the same layouts. diff --git a/crates/table/src/table.rs b/crates/table/src/table.rs index ad5f8418aa5..82d7904660e 100644 --- a/crates/table/src/table.rs +++ b/crates/table/src/table.rs @@ -59,7 +59,7 @@ use spacetimedb_schema::{ table_name::TableName, }; use std::{ - collections::{btree_map, BTreeMap}, + collections::{btree_map, BTreeMap, BTreeSet}, sync::Arc, }; use thiserror::Error; @@ -1110,8 +1110,10 @@ impl Table { // SAFETY: Caller promised that `self.is_row_present(new)` holds. let res = unsafe { self.insert_into_indices::(blobs, new) }; if let Err(e) = res { - // Undo (1). - unsafe { self.insert_into_indices::(blobs, old) } + // Undo (1). Cleanup removed the replacement's index entries, so the + // old unique keys are free. Skip duplicate-row comparison: restoring the + // pinned old row must not try to load another row during rollback. + unsafe { self.insert_into_indices::(blobs, old) } .expect("re-inserting the old row into indices should always work"); return Err(e); } @@ -1168,23 +1170,31 @@ impl Table { } if let Some((index_id, old_ptr)) = violation { // Found unique constraint violation! - let error = if CHECK_SAME_ROW + // Capture comparison errors so they cannot bypass the pinned + // index and physical-row cleanup below. + let same_row = if CHECK_SAME_ROW // If the index was added in this tx, // `old` could be a committed row, // which we want to avoid here. // TODO(centril): not 100% correct, could still be a duplicate, // but this is rather pathological and should be fixed when we restructure. && old_ptr.squashed_offset().is_tx_state() - // SAFETY: - // - The row layouts are the same as it's the same table. - // - We know `old` exists in `self` as we just found it in an index. - // - Caller promised that `new` is valid for `self`. - && unsafe { Self::eq_row_in_page(self, old_ptr, self, new_ref.pointer()) }.map_err(TableError::from)? { - DuplicateError(old_ptr).into() + // SAFETY: + // - The row layouts are the same as it's the same table. + // - We know `old` exists in `self` as we just found it in an index. + // - Caller promised that `new` is valid for `self`. + unsafe { Self::eq_row_in_page(self, old_ptr, self, new_ref.pointer()) } } else { - let index = self.indexes.get(&index_id).unwrap(); - self.build_error_unique(index, index_id, &new_ref).into() + Ok(false) + }; + let error = match same_row { + Ok(true) => DuplicateError(old_ptr).into(), + Ok(false) => { + let index = self.indexes.get(&index_id).unwrap(); + self.build_error_unique(index, index_id, &new_ref).into() + } + Err(e) => TableError::from(e).into(), }; // Delete row from indices. @@ -2351,11 +2361,15 @@ impl Table { pinned_blobs.extend(row.blobs.iter().map(|(_, handle)| handle.clone())); } + let mut comparison_pages = BTreeSet::new(); let mut prepared = self.inner.pages.prepare_commit( self.row_size(), &self.inner.visitor_prog, deletes, inserts.into_iter().map_ok(|insert| { + if let Some(hash) = insert.pointer_map_hash { + comparison_pages.extend(self.pointers_for(hash).iter().map(|ptr| ptr.page_index())); + } let num_granules = match &insert.inner { PreparedInsertInner::Product(row) => self.required_var_len_granules_for_insert(row), PreparedInsertInner::CopyFixed { .. } => 0, @@ -2363,27 +2377,43 @@ impl Table { (insert, num_granules) }), )?; + // Applying an insert may compare it with an untouched row sharing its + // hash. Pin those pages before any mutations, too. Equality reads blob + // hashes from the page, so it does not require the candidate's blobs. + for page_index in comparison_pages { + prepared.pin_page(&self.inner.pages, page_index)?; + } prepared.pin_blobs(pinned_blobs); Ok(prepared) } + /// Prepare a row for insertion into this table during commit. pub fn prepare_insert(&self, row: RowRef<'_>) -> PreparedInsert { assert_eq!(self.row_layout(), row.row_layout()); - if !self.schema.is_event + // Only tiered tables with a pointer map can fault while comparing + // committed candidates. Compute the hash while the source row is pinned; + // hashing its ProductValue later need not produce the same row hash. + let pointer_map_hash = (!self.schema.is_event && !self.inner.pages.is_resident() && self.pointer_map.is_some()) + .then(|| row.row_hash()); + let inner = if !self.schema.is_event && self.inner.row_layout.layout().fixed && let Some(source_page) = &row.page { - return PreparedInsertInner::CopyFixed { + PreparedInsertInner::CopyFixed { source_page: source_page.clone(), source_offset: row.pointer().page_offset(), } - .into(); + } else { + // Resident pages are borrowed, so prepared commits own the row value. + // Unlimited fixed-row resident commits bypass preparation entirely. + PreparedInsertInner::Product(row.to_product_value()) + }; + PreparedInsert { + inner, + pointer_map_hash, } - // Resident pages are borrowed, so prepared commits own the row value. - // Unlimited fixed-row resident commits bypass preparation entirely. - PreparedInsertInner::Product(row.to_product_value()).into() } fn required_var_len_granules_for_insert(&self, row: &ProductValue) -> usize { if self.inner.row_layout.layout().fixed { @@ -3275,6 +3305,8 @@ impl Table { /// normal insert path. pub struct PreparedInsert { inner: PreparedInsertInner, + /// Present when apply may compare this row with evictable pointer-map candidates. + pointer_map_hash: Option, } impl PreparedInsert { @@ -3286,12 +3318,6 @@ impl PreparedInsert { } } -impl From for PreparedInsert { - fn from(inner: PreparedInsertInner) -> Self { - Self { inner } - } -} - enum PreparedInsertInner { Product(ProductValue), CopyFixed { @@ -4686,4 +4712,114 @@ pub(crate) mod test { assert_eq!(table.num_rows(), count_before); assert_eq!(table.scan_all_row_ptrs().unwrap().len() as u64, count_before); } + + fn assert_prepared_commit_protects_pointer_map_candidates(fail_during_preparation: bool, fixed_rows: bool) { + use std::sync::atomic::Ordering; + + let (storage, fail_reads) = TieredStorage::with_failing_reads_for_test(); + let storage = Arc::new(storage); + let resident = Arc::new(ResidentObjects::new(ByteBudget::unlimited())); + let manager = Arc::new(PageManager::new( + resident.clone(), + PagePool::new_for_test(), + storage.clone(), + )); + let blobs = BlobSet::new(resident, storage); + let value_type = if fixed_rows { + AlgebraicType::U64 + } else { + AlgebraicType::String + }; + let schema = TableSchema::from_product_type(ProductType::from([AlgebraicType::U64, value_type])); + let mut table = Table::new( + schema.into(), + SquashedOffset::COMMITTED_STATE, + manager, + PageEvictionPolicy::Evictable, + ); + let value: AlgebraicValue = if fixed_rows { + 42_u64.into() + } else { + "x".repeat(128).into() + }; + let mut first_pointer = None; + let mut deletion = None; + // Leave a full first page and a row on the second page. The commit + // replaces the second-page row, so placement need not read the first. + for id in 0..PAGE_SIZE as u64 { + let (_, row) = table.insert(&blobs, &product![id, value.clone()]).unwrap(); + first_pointer.get_or_insert(row.pointer()); + if row.pointer().page_index() == PageIndex(1) { + deletion = Some((row.to_product_value(), row.into_pinned_row())); + break; + } + } + let first_pointer = first_pointer.unwrap(); + let (deleted_value, deleted_pin) = deletion.expect("fixture must fill the first page"); + let mut source = table.clone_structure(SquashedOffset::TX_STATE); + let incoming = product![u64::MAX, value]; + let (_, row) = source.insert(&blobs, &incoming).unwrap(); + let hash = row.row_hash(); + let insert = table.prepare_insert(row); + // Force the bucket lookup used by a hash collision, without searching + // for a real 64-bit collision. The rows still compare unequal. + assert!(table.pointers_for(hash).is_empty()); + table.pointer_map.as_mut().unwrap().insert(hash, first_pointer); + let pointers_before = table.pointer_map.clone(); + let count_before = table.num_rows(); + + if fail_during_preparation { + // The deletion page stays pinned, isolating the fault to the + // comparison candidate on the otherwise untouched first page. + table.inner.pages.evict().unwrap(); + fail_reads.store(true, Ordering::Relaxed); + } + let result = table.prepare_commit(&blobs, [deleted_pin.pointer()], [Ok(insert)]); + if fail_during_preparation { + assert!(matches!(result, Err(TableError::Page(error)) + if matches!(*error, PageError::Storage(_)))); + fail_reads.store(false, Ordering::Relaxed); + assert_eq!(table.pointer_map, pointers_before); + assert_eq!(table.num_rows(), count_before); + assert_eq!(table.scan_all_row_ptrs().unwrap().len() as u64, count_before); + assert!(table.scan_rows(&blobs).any(|row| row.unwrap() == deleted_value)); + } else { + let prepared = result.unwrap(); + drop(deleted_pin); + // After preparation, eviction and failed storage reads must not + // interrupt apply, even after its deletion has already happened. + table.inner.pages.evict().unwrap(); + fail_reads.store(true, Ordering::Relaxed); + let (deletes, inserts) = prepared.apply(&mut table, &blobs).into_parts(); + fail_reads.store(false, Ordering::Relaxed); + assert_eq!(&*deletes, std::slice::from_ref(&deleted_value)); + assert_eq!(&*inserts, std::slice::from_ref(&incoming)); + assert_eq!(table.num_rows(), count_before); + let rows = table + .scan_rows(&blobs) + .map(|row| row.unwrap().to_product_value()) + .collect::>(); + assert_eq!(rows.len() as u64, count_before); + assert!(rows.contains(&incoming)); + assert!(!rows.contains(&deleted_value)); + } + // Remove the synthetic association; the first row's real hash entry + // has remained intact throughout the test. + assert!(table.pointer_map.as_mut().unwrap().remove(hash, first_pointer)); + } + + #[test] + fn prepared_commit_pointer_map_fault_precedes_mutations() { + assert_prepared_commit_protects_pointer_map_candidates(true, false); + } + + #[test] + fn prepared_commit_retains_pointer_map_candidate_pins() { + assert_prepared_commit_protects_pointer_map_candidates(false, false); + } + + #[test] + fn prepared_fixed_commit_retains_pointer_map_candidate_pins() { + assert_prepared_commit_protects_pointer_map_candidates(false, true); + } } diff --git a/crates/table/src/tiered/page_set.rs b/crates/table/src/tiered/page_set.rs index 57638e3daf4..03ce8f37c5c 100644 --- a/crates/table/src/tiered/page_set.rs +++ b/crates/table/src/tiered/page_set.rs @@ -990,6 +990,21 @@ pub struct PreparedCommit { } impl PreparedCommit { + /// Retain an additional comparison page in the same set as mutation pages. + /// Deduplicating by page index avoids reloading pages already protected by the plan. + /// `pages` must be the page set used to prepare this commit. + pub(crate) fn pin_page(&mut self, pages: &PageSet, index: PageIndex) -> Result<(), PageError> { + if let btree_map::Entry::Vacant(entry) = self.pinned.entry(index) { + let page = pages + .get_page(index)? + .expect("pointer-map candidate belongs to a live page"); + if let PageRef::Tiered(pin) = page { + entry.insert(pin); + } + } + Ok(()) + } + /// Apply a transaction directly when neither source nor destination can fault /// and the destination has no allocation budget or page-index failure ahead. pub fn apply_resident(