fix: address PR 48 review findings

Cargo.toml

@@ -58,6 +58,7 @@ prettytable-rs = "^0.10"
 psc-nanoid = { version = "3.1.1", features = ["rkyv", "packed"] }
 rkyv = { version = "0.8.17", features = ["uuid-1"] }
 reqwest = { version = "0.12", optional = true, default-features = false, features = ["rustls-tls-webpki-roots", "charset", "http2"] }
+rustc-hash = "2.1.1"
 rusty-s3 = { version = "0.10.2", optional = true }
 smart-default = "0.7.1"
 tokio = { version = "1", features = ["full"] }

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

@@ -66,8 +66,36 @@ impl InMemoryGenerator {
         // column and therefore always reinserts, correctly.)
         let const_name = name_generator.get_page_inner_size_const_ident();
         let full_row_in_place_eligible = !self.columns.is_sized && self.columns.indexes.is_empty();
-        let size_check = if self.columns.is_sized {
-            quote! {}
+        let update_body = if self.columns.is_sized {
+            quote! {
+                let mut bytes = rkyv::to_bytes::<rkyv::rancor::Error>(&row)
+                    .map_err(|_| WorkTableError::SerializeError)?;
+                let mut archived_row = unsafe {
+                    rkyv::access_unchecked_mut::<<#row_ident as rkyv::Archive>::Archived>(&mut bytes[..])
+                        .unseal_unchecked()
+                };
+
+                let op_id = OperationId::Single(uuid::Uuid::now_v7());
+                #diff_process_insert
+                #persist_op
+
+                unsafe {
+                    self.0
+                        .data
+                        .with_mut_ref(link, move |archived| {
+                            #(#row_updates)*
+                        })
+                        .map_err(WorkTableError::PagesError)?
+                };
+
+                #diff_process_remove
+
+                self.0.update_state.remove(&pk);
+
+                #persist_call
+
+                core::result::Result::Ok(())
+            }
         } else if full_row_in_place_eligible {
             quote! {
                 // No secondary indexes: same-size unsized full-row update may go
@@ -93,29 +121,27 @@ impl InMemoryGenerator {
                     return Err(e);
                 }
                 self.0.update_state.remove(&pk);
-                return core::result::Result::Ok(());
+                core::result::Result::Ok(())
             }
         } else {
             quote! {
-                if true {
-                    drop(_guard);
-                    let op_lock = { #full_row_lock };
-                    let _guard = LockGuard::new_with_mutation(
-                        op_lock,
-                        self.0.lock_manager.clone(),
-                        pk.clone(),
-                    );
-                    let row_old = self.0.data.select_non_ghosted(link)?;
-                    if let Err(e) = self.reinsert(row_old, row).await {
-                        self.0.update_state.remove(&pk);
-
-                        return Err(e);
-                    }
-
+                drop(_guard);
+                let op_lock = { #full_row_lock };
+                let _guard = LockGuard::new_with_mutation(
+                    op_lock,
+                    self.0.lock_manager.clone(),
+                    pk.clone(),
+                );
+                let row_old = self.0.data.select_non_ghosted(link)?;
+                if let Err(e) = self.reinsert(row_old, row).await {
                     self.0.update_state.remove(&pk);
 
-                    return core::result::Result::Ok(());
+                    return Err(e);
                 }
+
+                self.0.update_state.remove(&pk);
+
+                core::result::Result::Ok(())
             }
         };
 
@@ -150,26 +176,7 @@ impl InMemoryGenerator {
                 let row_old = self.0.data.select_non_ghosted(link)?;
                 self.0.update_state.insert(pk.clone(), row_old);
 
-                let mut bytes = rkyv::to_bytes::<rkyv::rancor::Error>(&row).map_err(|_| WorkTableError::SerializeError)?;
-                #size_check
-
-                let mut archived_row = unsafe { rkyv::access_unchecked_mut::<<#row_ident as rkyv::Archive>::Archived>(&mut bytes[..]).unseal_unchecked() };
-
-                let op_id = OperationId::Single(uuid::Uuid::now_v7());
-                #diff_process_insert
-                #persist_op
-
-                unsafe { self.0.data.with_mut_ref(link, move |archived| {
-                    #(#row_updates)*
-                }).map_err(WorkTableError::PagesError)? };
-
-                #diff_process_remove
-
-                self.0.update_state.remove(&pk);
-
-                #persist_call
-
-                core::result::Result::Ok(())
+                #update_body
             }
         }
     }
@@ -367,8 +374,10 @@ impl InMemoryGenerator {
                         return core::result::Result::Ok(());
                     }
 
-                    let row_old_for_reinsert = self.0.select(pk.clone()).expect("should not be deleted by other thread");
-                    if let Err(e) = self.reinsert(row_old_for_reinsert, row_new).await {
+                    // `update_in_place` checks serialization and exact slot
+                    // length before touching page bytes, so its error path
+                    // leaves this locked snapshot authoritative for fallback.
+                    if let Err(e) = self.reinsert(row_old, row_new).await {
                         self.0.update_state.remove(&pk);
                         return Err(e);
                     }

src/in_memory/data.rs

@@ -135,6 +135,10 @@ impl<Row, const DATA_LENGTH: usize> Data<Row, DATA_LENGTH> {
         if length != link.length {
             return Err(ExecutionError::InvalidLink);
         }
+        debug_assert_eq!(
+            length, link.length,
+            "slot length was checked before archived bytes are overwritten"
+        );
 
         let inner_data = unsafe { &mut *self.inner_data.get() };
         inner_data[link.offset as usize..][..link.length as usize].copy_from_slice(bytes.as_slice());

src/in_memory/pages.rs

@@ -645,6 +645,12 @@ where
     /// generated persisted update path deliberately keeps the reinsert path for
     /// this reason.
     ///
+    /// Serialization and the exact-length check finish before any page byte is
+    /// changed. `page_access` excludes low-level archived-page readers during
+    /// the copy, while generated reads continue from the old immutable
+    /// publication until [`Self::publish_wrapped_row`] replaces the complete
+    /// owned row and flags together.
+    ///
     /// # Safety
     /// Same contract as [`Self::update`]: `link` must be valid and no other
     /// mutable references to the row may exist during modification.
@@ -909,6 +915,7 @@ mod tests {
     use crate::in_memory::pages::{DataPages, ExecutionError};
     use crate::in_memory::{DATA_INNER_LENGTH, PagesExecutionError, RowWrapper, StorableRow};
     use crate::prelude::ArchivedRowWrapper;
+    use data_bucket::Link;
 
     #[derive(Archive, Copy, Clone, Deserialize, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
     struct TestRow {
@@ -1089,6 +1096,32 @@ mod tests {
         assert_eq!(pages.select_non_ghosted(link), Ok(TestRow { a: 1, b: 1 }));
     }
 
+    #[test]
+    fn failed_exact_length_update_preserves_page_bytes_and_publication() {
+        let pages = DataPages::<TestRow>::new();
+        let old_row = TestRow { a: 10, b: 20 };
+        let link = pages.insert(old_row).unwrap();
+        unsafe {
+            pages.with_mut_ref(link, |row| row.unghost()).unwrap();
+        }
+        let old_bytes = pages.select_raw(link).unwrap();
+        let wrong_length = Link {
+            length: link.length - 1,
+            ..link
+        };
+
+        let result = unsafe { pages.update_in_place::<DATA_INNER_LENGTH>(TestRow { a: 30, b: 40 }, wrong_length) };
+
+        assert!(matches!(
+            result,
+            Err(ExecutionError::DataPageError(
+                crate::in_memory::DataExecutionError::InvalidLink
+            ))
+        ));
+        assert_eq!(pages.select_raw(link).unwrap(), old_bytes);
+        assert_eq!(pages.select_non_ghosted(link), Ok(old_row));
+    }
+
     #[test]
     fn retired_version_survives_link_reuse_for_in_flight_reader() {
         let pages = DataPages::<TestRow>::new();

src/index/persistent_wti.rs

@@ -10,7 +10,6 @@
 //! that marker and derives the real structural metadata itself.
 
 use std::array;
-use std::collections::hash_map::DefaultHasher;
 use std::fmt::{self, Debug};
 use std::hash::{Hash, Hasher};
 use std::ops::RangeBounds;
@@ -22,17 +21,28 @@ use indexset::cdc::change::{ChangeEvent, Id};
 use indexset::core::node::NodeLike;
 use indexset::core::pair::Pair;
 use parking_lot::{Mutex, MutexGuard};
+use rustc_hash::FxHasher;
 
 use crate::index::UniqueIndex;
 use crate::util::OffsetEqLink;
 use crate::{IndexMap, TableIndexCdc};
 
-// A stripe provides per-key exclusion only: DefaultHasher has no relationship
+// A stripe provides per-key exclusion only: FxHasher has no relationship
 // to key order, so callers must not infer range or cross-key ordering from the
 // selected mutex. Reads never touch this fixed inline table. Logical batches
 // are ordered independently by event id in the persistence worker.
 const MUTATION_STRIPES: usize = 64;
 
+#[inline]
+fn mutation_stripe_index<Q: Hash + ?Sized>(key: &Q) -> usize {
+    // Stripe selection is not a security boundary. FxHasher avoids SipHash's
+    // per-mutation cost and distributes the overwhelmingly common sequential
+    // integer keys across the power-of-two stripe table.
+    let mut hasher = FxHasher::default();
+    key.hash(&mut hasher);
+    hasher.finish() as usize & (MUTATION_STRIPES - 1)
+}
+
 /// A persisted WorkTablesIndex whose foreground mutations emit logical CDC.
 ///
 /// Point reads delegate directly to the native index. There is no runtime
@@ -45,6 +55,9 @@ where
 {
     inner: IndexMap<K, V, Node>,
     next_event_id: AtomicU64,
+    // Fixed inline allocation: 64 parking_lot mutexes per persisted WTI. The
+    // enclosing index's allocation size accounts for these; there is no
+    // per-mutation or per-key mutex allocation.
     mutation_stripes: [Mutex<()>; MUTATION_STRIPES],
 }
 
@@ -96,10 +109,9 @@ where
         &self.inner
     }
 
+    #[inline]
     fn mutation_stripe<Q: Hash + ?Sized>(&self, key: &Q) -> MutexGuard<'_, ()> {
-        let mut hasher = DefaultHasher::new();
-        key.hash(&mut hasher);
-        self.mutation_stripes[hasher.finish() as usize % MUTATION_STRIPES].lock()
+        self.mutation_stripes[mutation_stripe_index(key)].lock()
     }
 
     fn next_event_id(&self) -> Id {
@@ -267,6 +279,8 @@ where
 
 #[cfg(test)]
 mod tests {
+    use std::collections::HashSet;
+
     use data_bucket::page::PageId;
 
     use super::*;
@@ -305,4 +319,12 @@ mod tests {
         assert!(index.insert_checked_cdc(8, link(8)).is_some());
         assert_eq!(index.next_event_id.load(Ordering::Relaxed), 2);
     }
+
+    #[test]
+    fn sequential_integer_keys_use_every_mutation_stripe() {
+        let stripes = (0_u64..4_096)
+            .map(|key| mutation_stripe_index(&key))
+            .collect::<HashSet<_>>();
+        assert_eq!(stripes.len(), MUTATION_STRIPES);
+    }
 }

src/lock/map.rs

@@ -98,6 +98,16 @@ impl Drop for MutationGuard {
     }
 }
 
+/// Registry for per-row async locks and synchronous mutation stripes.
+///
+/// # Sync/async lock boundary
+///
+/// The `parking_lot` map guard is never returned and never crosses an
+/// `.await`. Acquisition clones a tracked `Arc<tokio::sync::RwLock<_>>` before
+/// releasing the map guard. Cleanup may synchronously take the short-lived map
+/// write guard, but only probes the per-row lock with `try_read`; it never waits
+/// on a Tokio lock while holding the map. This one-way boundary prevents a
+/// map-lock/per-row-lock cycle during cancellation and `Drop`.
 #[derive(Debug)]
 pub struct LockMap<LockType, PrimaryKey> {
     map: RwLock<HashMap<PrimaryKey, LockEntry<LockType>>>,
@@ -119,6 +129,12 @@ impl<LockType, PrimaryKey> LockMap<LockType, PrimaryKey>
 where
     PrimaryKey: Hash + Eq + Debug + Clone,
 {
+    /// Inserts a raw lock entry.
+    ///
+    /// A returned or externally retained `Arc` pins cleanup through
+    /// `Arc::strong_count`. Generated operations should prefer
+    /// [`Self::get_or_insert_with`], whose [`LockAcquirer`] makes cancellation
+    /// tracking explicit.
     pub fn insert(
         &self,
         key: PrimaryKey,
@@ -136,6 +152,8 @@ where
             .map(|entry| entry.lock)
     }
 
+    /// Returns an untracked raw lock clone, which keeps the map entry alive
+    /// until that clone is dropped.
     pub fn get(&self, key: &PrimaryKey) -> Option<Arc<tokio::sync::RwLock<LockType>>> {
         self.map.read().get(key).map(|entry| entry.lock.clone())
     }
@@ -278,4 +296,24 @@ mod tests {
         drop(second);
         assert!(!lock_map.map.read().contains_key(&33));
     }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn cancelling_async_waiter_releases_tracking_without_deadlock() {
+        let lock_map: Arc<LockMap<FullRowLock, u64>> = Arc::new(LockMap::default());
+        let owner = lock_map.get_or_insert_with(41, FullRowLock::new);
+        let owner_guard = owner.write().await;
+        let waiter = lock_map.get_or_insert_with(41, FullRowLock::new);
+        let waiting_task = tokio::spawn(async move {
+            let _guard = waiter.write().await;
+        });
+        tokio::task::yield_now().await;
+
+        waiting_task.abort();
+        assert!(waiting_task.await.unwrap_err().is_cancelled());
+        assert!(lock_map.map.read().contains_key(&41));
+
+        drop(owner_guard);
+        drop(owner);
+        assert!(!lock_map.map.read().contains_key(&41));
+    }
 }