fix bitmap batch parent chain growth with rwlock

Author: roberto-bayardo

storage/src/qmdb/current/batch.rs

@@ -27,7 +27,10 @@ use crate::{
 };
 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::{BTreeMap, BTreeSet},
     sync::Arc,
@@ -279,6 +282,29 @@ where
 ///
 /// 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
+/// [`BitmapBatch::Base`]. That bitmap evolves in place as [`Db::apply_batch`](super::db::Db)
+/// 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 — i.e., every
+/// [`apply_batch`](super::db::Db::apply_batch) call since this batch was merkleized 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.
+///
+/// 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,
@@ -590,7 +616,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),
     }));
@@ -632,13 +657,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.
+///
+/// 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. Batches never write;
+/// write access is kept module-private to `current::db`.
+///
+/// # 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.
+    pub(crate) fn read(&self) -> RwLockReadGuard<'_, BitMap<N>> {
+        self.inner.read()
+    }
+
+    /// Acquire an exclusive write guard. 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.
 ///
-/// Mirrors the [`crate::merkle::batch::MerkleizedBatch`] pattern.
+/// 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>>),
 }
@@ -649,10 +747,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> {
@@ -666,7 +760,7 @@ impl<const N: usize> BitmapReadable<N> for BitmapBatch<N> {
 
     fn get_chunk(&self, idx: usize) -> [u8; N] {
         match self {
-            Self::Base(bm) => *bm.get_chunk(idx),
+            Self::Base(shared) => *shared.read().get_chunk(idx),
             Self::Layer(layer) => {
                 // Check overlay first; fall through to parent if unmodified.
                 if let Some(&chunk) = layer.overlay.get(idx) {
@@ -695,95 +789,19 @@ 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,
+            Self::Base(shared) => shared.read().pruned_chunks(),
+            Self::Layer(layer) => layer.parent.pruned_chunks(),
         }
     }
 
     fn len(&self) -> u64 {
         match self {
-            Self::Base(bm) => BitmapReadable::<N>::len(bm.as_ref()),
+            Self::Base(shared) => BitmapReadable::<N>::len(&*shared.read()),
             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
@@ -853,7 +871,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,
         })
     }
@@ -1294,32 +1312,4 @@ mod tests {
         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(_)));
-    }
 }