Use pinned rows to make datastore mutations failure-safe

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.

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::<true>(blobs, new) };
         if let Err(e) = res {
-            // Undo (1).
-            unsafe { self.insert_into_indices::<true>(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::<false>(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,39 +2361,59 @@ 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,
                 };
                 (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<RowHash>,
 }
 
 impl PreparedInsert {
@@ -3286,12 +3318,6 @@ impl PreparedInsert {
     }
 }
 
-impl From<PreparedInsertInner> 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::<Vec<_>>();
+            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);
+    }
 }

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(