fix bitmap batch parent chain growth with rwlock

Author: roberto-bayardo

storage/src/qmdb/current/batch.rs

@@ -28,7 +28,10 @@ use crate::{
 use ahash::AHasher;
 use commonware_codec::Codec;
 use commonware_cryptography::{Digest, Hasher};
-use commonware_utils::bitmap::{Prunable as BitMap, Readable as BitmapReadable};
+use commonware_utils::{
+    bitmap::{Prunable as BitMap, Readable as BitmapReadable},
+    sync::{RwLock, RwLockReadGuard, RwLockWriteGuard},
+};
 use std::{
     collections::{BTreeSet, HashMap},
     hash::BuildHasherDefault,
@@ -86,11 +89,12 @@ impl<const N: usize> ChunkOverlay<N> {
         chunk[rel / 8] |= 1 << (rel % 8);
     }
 
-    /// Clear a single bit (used for superseded locations).
-    /// Skips locations in pruned chunks — those bits are already inactive.
-    fn clear_bit<B: BitmapReadable<N>>(&mut self, base: &B, loc: u64) {
+    /// Clear a single bit (used for superseded locations). `pruned_chunks` is passed in by the
+    /// caller so the hot loop in `build_chunk_overlay` reads it once rather than per call.
+    /// Skips locations in pruned chunks since those bits are already inactive.
+    fn clear_bit<B: BitmapReadable<N>>(&mut self, base: &B, pruned_chunks: usize, loc: u64) {
         let idx = BitMap::<N>::to_chunk_index(loc);
-        if idx < base.pruned_chunks() {
+        if idx < pruned_chunks {
             return;
         }
         let rel = (loc % Self::CHUNK_BITS) as usize;
@@ -279,11 +283,35 @@ where
     bitmap_parent: BitmapBatch<N>,
 }
 
-/// A speculative batch of operations whose root digest has been computed,
-/// in contrast to [`UnmerkleizedBatch`].
+/// A speculative batch of operations whose root digest has been computed, in contrast to
+/// [`UnmerkleizedBatch`].
+///
+/// Wraps an [`any::batch::MerkleizedBatch`] and adds the bitmap and grafted MMR state needed to
+/// compute the canonical root.
+///
+/// # Stale reads
+///
+/// A `MerkleizedBatch` holds a reference to the DB's committed bitmap via the internal
+/// `BitmapBatch::Base` variant. That bitmap evolves in place as
+/// [`Db::apply_batch`](super::db::Db::apply_batch) commits batches.
+/// Reads through this batch's chain (via its own `get`, by constructing child batches, or by
+/// anything that walks the bitmap chain) are only semantically correct if this batch's chain is
+/// still a valid prefix of the DB's committed state. That is, every
+/// [`apply_batch`](super::db::Db::apply_batch) call since this batch was merkleized has applied an
+/// ancestor of *this* batch.
+///
+/// **The library does not guard against stale reads.** If you apply an unrelated sibling or any
+/// other non-ancestor batch and then read through this one, chunk reads fall through to the
+/// (now-advanced) committed bitmap while your overlays are still applied on top, producing
+/// incoherent state that doesn't correspond to any valid DB view.
 ///
-/// Wraps an [`any::batch::MerkleizedBatch`] and adds the bitmap and grafted MMR state needed
-/// to compute the canonical root.
+/// Rules of thumb:
+/// - Drop any `Arc<MerkleizedBatch>` you no longer intend to apply.
+/// - Extending a batch after `apply_batch` has consumed it (building a child off the just-applied
+///   parent) is safe. The committed bitmap now equals the parent's post-apply state, so child
+///   reads are consistent.
+/// - Extending a batch after a *different* branch has been applied is not safe and will silently
+///   return wrong data.
 pub struct MerkleizedBatch<F: Graftable, D: Digest, U: update::Update + Send + Sync, const N: usize>
 where
     Operation<F, U>: Send + Sync,
@@ -446,13 +474,14 @@ where
 {
     let total_bits = base.len() + batch_len as u64;
     let mut overlay = ChunkOverlay::new(total_bits);
+    let pruned_chunks = base.pruned_chunks();
 
     // 1. CommitFloor (last op) is always active.
     let commit_loc = batch_base + batch_len as u64 - 1;
     overlay.set_bit(base, commit_loc);
 
     // 2. Inactivate previous CommitFloor.
-    overlay.clear_bit(base, batch_base - 1);
+    overlay.clear_bit(base, pruned_chunks, batch_base - 1);
 
     // 3. Set active bits + clear superseded locations from the diff.
     for (key, entry) in diff {
@@ -473,7 +502,7 @@ where
             }
         }
         if let Some(old) = prev_loc {
-            overlay.clear_bit(base, *old);
+            overlay.clear_bit(base, pruned_chunks, *old);
         }
     }
 
@@ -595,7 +624,6 @@ where
     // compute_db_root sees newly completed chunks. Using bitmap_parent alone would miss chunks
     // that transitioned from partial to complete in this batch.
     let bitmap_batch = BitmapBatch::Layer(Arc::new(BitmapBatchLayer {
-        pruned_chunks: bitmap_parent.pruned_chunks(),
         parent: bitmap_parent.clone(),
         overlay: Arc::new(overlay),
     }));
@@ -637,13 +665,86 @@ where
     }))
 }
 
-/// Immutable bitmap state at any point in a batch chain.
+/// The committed bitmap shared between the [`Db`](super::db::Db) and live batches.
 ///
-/// Mirrors the [`crate::merkle::batch::MerkleizedBatch`] pattern.
+/// Wrapped in a [`RwLock`] so that [`Db::apply_batch`](super::db::Db::apply_batch),
+/// [`Db::prune`](super::db::Db::prune), and [`Db::rewind`](super::db::Db::rewind) can mutate
+/// the bitmap in place while live batches concurrently read through it.
+///
+/// # Stale reads
+///
+/// The bitmap behind this lock represents *committed* state. If a caller holds a `MerkleizedBatch`
+/// whose chain is no longer a valid prefix of committed state (e.g., an orphaned sibling branch
+/// after the other branch was applied), reads through that chain's `Base` will silently return
+/// inconsistent data (the chain's overlays mixed with post-divergence committed chunks). The
+/// library does not guard against this; callers must avoid reading through stale batches.
+pub(crate) struct SharedBitmap<const N: usize> {
+    inner: RwLock<BitMap<N>>,
+}
+
+impl<const N: usize> SharedBitmap<N> {
+    pub(crate) const fn new(bitmap: BitMap<N>) -> Self {
+        Self {
+            inner: RwLock::new(bitmap),
+        }
+    }
+
+    /// Acquire a shared read guard over the committed bitmap. Kept private so external callers
+    /// go through [`BitmapReadable`] (which doesn't expose a guard across `.await`).
+    fn read(&self) -> RwLockReadGuard<'_, BitMap<N>> {
+        self.inner.read()
+    }
+
+    /// Acquire an exclusive write guard. By convention only
+    /// [`Db::apply_batch`](super::db::Db::apply_batch), [`Db::prune`](super::db::Db::prune), and
+    /// [`Db::rewind`](super::db::Db::rewind) mutate the shared bitmap.
+    pub(super) fn write(&self) -> RwLockWriteGuard<'_, BitMap<N>> {
+        self.inner.write()
+    }
+}
+
+impl<const N: usize> std::fmt::Debug for SharedBitmap<N> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SharedBitmap")
+            .field("bitmap_len", &BitmapReadable::<N>::len(&*self.read()))
+            .finish()
+    }
+}
+
+/// [`BitmapReadable`] over the DB's committed bitmap. Each call acquires the read lock briefly.
+impl<const N: usize> BitmapReadable<N> for SharedBitmap<N> {
+    fn complete_chunks(&self) -> usize {
+        self.read().complete_chunks()
+    }
+
+    fn get_chunk(&self, idx: usize) -> [u8; N] {
+        *self.read().get_chunk(idx)
+    }
+
+    fn last_chunk(&self) -> ([u8; N], u64) {
+        let guard = self.read();
+        let (chunk, bits) = guard.last_chunk();
+        (*chunk, bits)
+    }
+
+    fn pruned_chunks(&self) -> usize {
+        self.read().pruned_chunks()
+    }
+
+    fn len(&self) -> u64 {
+        BitmapReadable::<N>::len(&*self.read())
+    }
+}
+
+/// A view of the committed bitmap plus zero or more speculative overlay `Layer`s.
+///
+/// The chain terminates in a `Base` that references the shared committed bitmap. No staleness
+/// check is performed. Callers must ensure they only read through batches whose chains are
+/// still valid prefixes of committed state (see [`SharedBitmap`]'s docs).
 #[derive(Clone, Debug)]
 pub(crate) enum BitmapBatch<const N: usize> {
-    /// Committed bitmap (chain terminal).
-    Base(Arc<BitMap<N>>),
+    /// Chain terminal: shared reference to the committed bitmap.
+    Base(Arc<SharedBitmap<N>>),
     /// Speculative layer on top of a parent batch.
     Layer(Arc<BitmapBatchLayer<N>>),
 }
@@ -654,10 +755,6 @@ pub(crate) struct BitmapBatchLayer<const N: usize> {
     pub(crate) parent: BitmapBatch<N>,
     /// Chunk-level overlay: materialized bytes for every chunk that differs from parent.
     pub(crate) overlay: Arc<ChunkOverlay<N>>,
-    /// Pruned-chunk count, copied from `parent` at construction. Invariant across the whole
-    /// layer chain (pruning only happens on the committed base), so caching here lets
-    /// `BitmapBatch::pruned_chunks` return in O(1) instead of walking to the Base.
-    pub(crate) pruned_chunks: usize,
 }
 
 impl<const N: usize> BitmapBatch<N> {
@@ -675,7 +772,7 @@ impl<const N: usize> BitmapReadable<N> for BitmapBatch<N> {
         let mut current = self;
         loop {
             match current {
-                Self::Base(bm) => return *bm.get_chunk(idx),
+                Self::Base(shared) => return shared.get_chunk(idx),
                 Self::Layer(layer) => {
                     if let Some(&chunk) = layer.overlay.get(idx) {
                         return chunk;
@@ -702,96 +799,23 @@ impl<const N: usize> BitmapReadable<N> for BitmapBatch<N> {
     }
 
     fn pruned_chunks(&self) -> usize {
-        match self {
-            Self::Base(bm) => bm.pruned_chunks(),
-            Self::Layer(layer) => layer.pruned_chunks,
+        let mut current = self;
+        loop {
+            match current {
+                Self::Base(shared) => return shared.pruned_chunks(),
+                Self::Layer(layer) => current = &layer.parent,
+            }
         }
     }
 
     fn len(&self) -> u64 {
         match self {
-            Self::Base(bm) => BitmapReadable::<N>::len(bm.as_ref()),
+            Self::Base(shared) => BitmapReadable::<N>::len(shared.as_ref()),
             Self::Layer(layer) => layer.overlay.len,
         }
     }
 }
 
-impl<const N: usize> BitmapBatch<N> {
-    /// Apply a chunk overlay to this bitmap. When `self` is `Base` with sole ownership, writes
-    /// overlay chunks directly into the bitmap. Otherwise creates a new `Layer`.
-    pub(super) fn apply_overlay(&mut self, overlay: Arc<ChunkOverlay<N>>) {
-        // Fast path: write overlay chunks directly into the Base bitmap.
-        if let Self::Base(base) = self {
-            if let Some(bitmap) = Arc::get_mut(base) {
-                // Extend bitmap to the overlay's length.
-                bitmap.extend_to(overlay.len);
-                // Overwrite dirty chunks.
-                for (&idx, chunk_bytes) in &overlay.chunks {
-                    if idx >= bitmap.pruned_chunks() {
-                        bitmap.set_chunk_by_index(idx, chunk_bytes);
-                    }
-                }
-                return;
-            }
-        }
-
-        // Slow path: create a new layer.
-        let pruned_chunks = self.pruned_chunks();
-        let parent = self.clone();
-        *self = Self::Layer(Arc::new(BitmapBatchLayer {
-            parent,
-            overlay,
-            pruned_chunks,
-        }));
-    }
-
-    /// Flatten all layers back to a single `Base(Arc<BitMap<N>>)`.
-    ///
-    /// After flattening, the new `Base` Arc has refcount 1 (assuming no external clones
-    /// are held).
-    pub(super) fn flatten(&mut self) {
-        if matches!(self, Self::Base(_)) {
-            return;
-        }
-
-        // Take ownership of the chain so that Arc refcounts are not
-        // artificially inflated by a clone.
-        let mut owned = std::mem::replace(self, Self::Base(Arc::new(BitMap::default())));
-
-        // Collect overlays from tip to base.
-        let mut overlays = Vec::new();
-        let base = loop {
-            match owned {
-                Self::Base(bm) => break bm,
-                Self::Layer(layer) => match Arc::try_unwrap(layer) {
-                    Ok(inner) => {
-                        overlays.push(inner.overlay);
-                        owned = inner.parent;
-                    }
-                    Err(arc) => {
-                        overlays.push(arc.overlay.clone());
-                        owned = arc.parent.clone();
-                    }
-                },
-            }
-        };
-
-        // Apply overlays from base to tip.
-        let mut bitmap = Arc::try_unwrap(base).unwrap_or_else(|arc| (*arc).clone());
-        for overlay in overlays.into_iter().rev() {
-            // Extend bitmap to the overlay's length.
-            bitmap.extend_to(overlay.len);
-            // Apply dirty chunks.
-            for (&idx, chunk_bytes) in &overlay.chunks {
-                if idx >= bitmap.pruned_chunks() {
-                    bitmap.set_chunk_by_index(idx, chunk_bytes);
-                }
-            }
-        }
-        *self = Self::Base(Arc::new(bitmap));
-    }
-}
-
 impl<F: Graftable, D: Digest, U: update::Update + Send + Sync, const N: usize>
     MerkleizedBatch<F, D, U, N>
 where
@@ -861,7 +885,7 @@ where
         Arc::new(MerkleizedBatch {
             inner: self.any.to_batch(),
             grafted,
-            bitmap: self.status.clone(),
+            bitmap: BitmapBatch::Base(Arc::clone(&self.status)),
             canonical_root: self.root,
         })
     }
@@ -1301,33 +1325,4 @@ mod tests {
         // Empty bitmap, tip = 0: no candidates.
         assert_eq!(scan.next_candidate(Location::new(0), 0), None);
     }
-
-    #[test]
-    fn test_apply_overlay() {
-        // Base: 8 bits all set, sole owner -> fast path.
-        let base = make_bitmap(&[true; 8]);
-        let mut batch = BitmapBatch::Base(Arc::new(base));
-
-        let mut overlay = ChunkOverlay::new(12);
-        let mut c0 = [0u8; N];
-        c0[0] = 0b1111_0111; // bits 0-7 set except bit 3
-        c0[1] = 0b0000_0100; // bit 10 set
-        overlay.chunks.insert(0, c0);
-        batch.apply_overlay(Arc::new(overlay));
-
-        // Fast path keeps Base, extends length, applies chunks.
-        assert!(matches!(batch, BitmapBatch::Base(_)));
-        assert_eq!(batch.len(), 12);
-        assert_eq!(batch.get_chunk(0)[0] & (1 << 3), 0);
-        assert_ne!(batch.get_chunk(0)[1] & (1 << 2), 0);
-
-        // Shared Arc -> slow path creates Layer.
-        let BitmapBatch::Base(ref base_arc) = batch else {
-            panic!("expected Base");
-        };
-        let _shared = Arc::clone(base_arc);
-        let overlay2 = ChunkOverlay::new(12);
-        batch.apply_overlay(Arc::new(overlay2));
-        assert!(matches!(batch, BitmapBatch::Layer(_)));
-    }
 }