tests for old_locs returned from commit

Author: roberto-bayardo

storage/src/adb/any.rs

@@ -33,7 +33,7 @@ use tracing::{debug, warn};
 /// Configuration for an `Any` authenticated db.
 #[derive(Clone)]
 pub struct Config {
-    /// The name of the `Storage` partition used for the MMR's backing journal.
+    /// The name of the `Storage` ˘ used for the MMR's backing journal.
     pub mmr_journal_partition: String,
 
     /// The items per blob configuration value used by the MMR journal.
@@ -69,8 +69,9 @@ pub struct Any<E: RStorage + Clock + Metrics, K: Array, V: Array, H: CHasher> {
     /// are over keys that have been updated by some operation at or after this point).
     inactivity_floor_loc: u64,
 
-    /// A map from each key to the location in the log containing its most recent update. Only
-    /// contains the keys that currently have a value (that is, deleted keys are not in the map).
+    /// A snapshot of all currently active operations in the form of a map from each key to the
+    /// location in the log containing its most recent update. Only contains the keys that currently
+    /// have a value (that is, deleted keys are not in the map).
     snapshot: Index<EightCap, u64>,
 
     /// The number of operations that are pending commit.
@@ -79,7 +80,9 @@ pub struct Any<E: RStorage + Clock + Metrics, K: Array, V: Array, H: CHasher> {
 
 impl<E: RStorage + Clock + Metrics, K: Array, V: Array, H: CHasher> Any<E, K, V, H> {
     /// Return an MMR initialized from `cfg`. Any uncommitted operations in the log will be
-    /// discarded and the state of the db will be as of the last committed operation.
+    /// discarded and the state of the db will be as of the last committed operation. If a bitmap is
+    /// provided, then a bit is appended for each operation in the operation log, with its value
+    /// appropriately set to its activity status.
     pub async fn init<const N: usize>(
         context: E,
         hasher: &mut H,
@@ -167,6 +170,9 @@ impl<E: RStorage + Clock + Metrics, K: Array, V: Array, H: CHasher> Any<E, K, V,
                     while let Some(loc) = loc_iter.next() {
                         let op = log.read(*loc).await?;
                         if op.to_key() == key {
+                            if let Some(ref mut bitmap_ref) = bitmap {
+                                bitmap_ref.set_bit(hasher, *loc, false);
+                            }
                             loc_iter.remove();
                             break;
                         }
@@ -186,6 +192,13 @@ impl<E: RStorage + Clock + Metrics, K: Array, V: Array, H: CHasher> Any<E, K, V,
                 }
                 Type::Commit(loc) => inactivity_floor_loc = loc,
             }
+            if let Some(ref mut bitmap_ref) = bitmap {
+                // If we reach this point and a bit hasn't been added for the operation, then it's
+                // an inactive operation and we need to tag it as such in the bitmap.
+                if bitmap_ref.bit_count() == i {
+                    bitmap_ref.append(hasher, false);
+                }
+            }
         }
 
         let db = Any {
@@ -276,7 +289,7 @@ impl<E: RStorage + Clock + Metrics, K: Array, V: Array, H: CHasher> Any<E, K, V,
     /// Updates `key` to have value `value`.  If the key already has this same value, then this is a
     /// no-op. The operation is reflected in the snapshot, but will be subject to rollback until the
     /// next successful `commit`. Returns None if the update was a no-op, and otherwise returns the
-    /// (old_position, new_position) pair for the updated key.
+    /// old_position (if any) for the updated key.
     pub async fn update(
         &mut self,
         hasher: &mut H,
@@ -574,12 +587,17 @@ mod test {
     use commonware_cryptography::{hash, sha256::Digest, Hasher as CHasher, Sha256};
     use commonware_macros::test_traced;
     use commonware_runtime::{deterministic, Runner as _};
-    use std::collections::HashMap;
+    use rand::{
+        rngs::{OsRng, StdRng},
+        RngCore, SeedableRng,
+    };
+    use std::collections::{HashMap, HashSet};
 
     /// Return an `Any` database initialized with a fixed config.
     async fn open_db<E: RStorage + Clock + Metrics, const N: usize>(
         context: E,
         hasher: &mut Sha256,
+        bitmap: Option<&mut Bitmap<Sha256, N>>,
     ) -> Any<E, Digest, Digest, Sha256> {
         let cfg = Config {
             mmr_journal_partition: "journal_partition".into(),
@@ -588,7 +606,7 @@ mod test {
             log_journal_partition: "log_journal_partition".into(),
             log_items_per_blob: 7,
         };
-        Any::<E, Digest, Digest, Sha256>::init::<N>(context, hasher, cfg, None)
+        Any::<E, Digest, Digest, Sha256>::init::<N>(context, hasher, cfg, bitmap)
             .await
             .unwrap()
     }
@@ -598,7 +616,7 @@ mod test {
         let executor = deterministic::Runner::default();
         executor.start(|context| async move {
             let mut hasher = Sha256::new();
-            let mut db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let mut db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
             assert_eq!(db.op_count(), 0);
             assert_eq!(db.oldest_retained_loc(), None);
             assert!(matches!(db.prune_inactive().await, Ok(())));
@@ -610,22 +628,25 @@ mod test {
             let root = db.root(&mut hasher);
             db.update(&mut hasher, d1, d2).await.unwrap();
             db.close().await.unwrap();
-            let mut db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let mut db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
             assert_eq!(db.root(&mut hasher), root);
             assert_eq!(db.op_count(), 0);
 
             // Test calling commit on an empty db which should make it (durably) non-empty.
-            db.commit(&mut hasher, None).await.unwrap();
+            let mut old_locs = Vec::new();
+            db.commit(&mut hasher, Some(&mut old_locs)).await.unwrap();
             assert_eq!(db.op_count(), 1); // floor op added
+            assert!(old_locs.is_empty());
             let root = db.root(&mut hasher);
             assert!(matches!(db.prune_inactive().await, Ok(())));
-            let mut db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let mut db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
             assert_eq!(db.root(&mut hasher), root);
 
             // Confirm the inactivity floor doesn't fall endlessly behind with multiple commits.
             for _ in 1..100 {
-                db.commit(&mut hasher, None).await.unwrap();
+                db.commit(&mut hasher, Some(&mut old_locs)).await.unwrap();
                 assert_eq!(db.op_count() - 1, db.inactivity_floor_loc);
+                assert!(old_locs.is_empty());
             }
         });
     }
@@ -637,7 +658,7 @@ mod test {
             // Build a db with 2 keys and make sure updates and deletions of those keys work as
             // expected.
             let mut hasher = Sha256::new();
-            let mut db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let mut db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
 
             let d1 = <Sha256 as CHasher>::Digest::decode(vec![1u8; 32].as_ref()).unwrap();
             let d2 = <Sha256 as CHasher>::Digest::decode(vec![2u8; 32].as_ref()).unwrap();
@@ -681,9 +702,11 @@ mod test {
             db.sync().await.unwrap();
 
             // Advance over 3 inactive operations.
-            db.raise_inactivity_floor(&mut hasher, 3, None)
+            let mut old_locs = Vec::new();
+            db.raise_inactivity_floor(&mut hasher, 3, Some(&mut old_locs))
                 .await
                 .unwrap();
+            assert!(old_locs.is_empty()); // no active ops were moved
             assert_eq!(db.inactivity_floor_loc, 3);
             assert_eq!(db.log.size().await.unwrap(), 6); // 4 updates, 1 deletion, 1 commit
 
@@ -718,11 +741,12 @@ mod test {
             assert_eq!(db.root(&mut hasher), root);
 
             // Make sure closing/reopening gets us back to the same state.
-            db.commit(&mut hasher, None).await.unwrap();
+            db.commit(&mut hasher, Some(&mut old_locs)).await.unwrap();
+            assert!(old_locs.is_empty());
             assert_eq!(db.log.size().await.unwrap(), 9);
             let root = db.root(&mut hasher);
             db.close().await.unwrap();
-            let mut db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let mut db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
             assert_eq!(db.log.size().await.unwrap(), 9);
             assert_eq!(db.root(&mut hasher), root);
 
@@ -745,7 +769,7 @@ mod test {
             db.commit(&mut hasher, None).await.unwrap();
             let root = db.root(&mut hasher);
             db.close().await.unwrap();
-            let mut db = open_db::<_, 32>(context, &mut hasher).await;
+            let mut db = open_db::<_, 32>(context, &mut hasher, None).await;
             assert_eq!(db.root(&mut hasher), root);
 
             // Commit will raise the inactivity floor, which won't affect state but will affect the
@@ -768,12 +792,12 @@ mod test {
         let executor = deterministic::Runner::default();
         // Build a db with 1000 keys, some of which we update and some of which we delete, and
         // confirm that the end state of the db matches that of an identically updated hashmap.
+        const ELEMENTS: u64 = 1000;
         executor.start(|context| async move {
             let mut hasher = Sha256::new();
-            let mut db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let mut db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
 
             let mut map = HashMap::<Digest, Digest>::default();
-            const ELEMENTS: u64 = 1000;
             for i in 0u64..ELEMENTS {
                 let k = hash(&i.to_be_bytes());
                 let v = hash(&(i * 1000).to_be_bytes());
@@ -809,15 +833,17 @@ mod test {
             assert_eq!(db.snapshot.len(), 857);
 
             // Test that commit will raise the activity floor.
-            db.commit(&mut hasher, None).await.unwrap();
+            let mut old_locs = Vec::new();
+            db.commit(&mut hasher, Some(&mut old_locs)).await.unwrap();
             assert_eq!(db.op_count(), 2336);
             assert_eq!(db.oldest_retained_loc().unwrap(), 1478);
             assert_eq!(db.snapshot.len(), 857);
+            assert_eq!(old_locs.len(), 858); // 857 active ops moved once, one moved twice.
 
             // Close & reopen the db, making sure the re-opened db has exactly the same state.
             let root_hash = db.root(&mut hasher);
             db.close().await.unwrap();
-            let mut db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let mut db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
             assert_eq!(root_hash, db.root(&mut hasher));
             assert_eq!(db.op_count(), 2336);
             assert_eq!(db.inactivity_floor_loc, 1478);
@@ -877,7 +903,7 @@ mod test {
         let executor = deterministic::Runner::default();
         executor.start(|context| async move {
             let mut hasher = Sha256::new();
-            let mut db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let mut db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
 
             // Insert 1000 keys then sync.
             const ELEMENTS: u64 = 1000;
@@ -904,7 +930,7 @@ mod test {
 
             // Journaled MMR recovery should read the orphaned leaf & its parents, then log
             // replaying will restore the rest.
-            let mut db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let mut db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
             assert_eq!(db.root(&mut hasher), root);
 
             // Write some additional nodes, simulate failed log commit, and test we recover to the previous commit point.
@@ -914,7 +940,7 @@ mod test {
                 db.update(&mut hasher, k, v).await.unwrap();
             }
             db.simulate_failed_commit_log(&mut hasher).await.unwrap();
-            let db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
             assert_eq!(db.root(&mut hasher), root);
         });
     }
@@ -926,7 +952,7 @@ mod test {
         let executor = deterministic::Runner::default();
         executor.start(|context| async move {
             let mut hasher = Sha256::new();
-            let mut db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let mut db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
 
             // Update the same key many times.
             const UPDATES: u64 = 100;
@@ -935,12 +961,15 @@ mod test {
                 let v = hash(&(i * 1000).to_be_bytes());
                 db.update(&mut hasher, k, v).await.unwrap();
             }
-            db.commit(&mut hasher, None).await.unwrap();
+            let mut old_locs = Vec::new();
+            db.commit(&mut hasher, Some(&mut old_locs)).await.unwrap();
+            // only 1 active op, but it gets moved twice due to the commit op.
+            assert_eq!(old_locs.len(), 2);
             let root = db.root(&mut hasher);
             db.close().await.unwrap();
 
             // Simulate a failed commit and test that the log replay doesn't leave behind old data.
-            let db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
             let iter = db.snapshot.get_iter(&k);
             assert_eq!(iter.cloned().collect::<Vec<_>>().len(), 1);
             assert_eq!(db.root(&mut hasher), root);
@@ -952,7 +981,7 @@ mod test {
         let executor = deterministic::Runner::default();
         executor.start(|context| async move {
             let mut hasher = Sha256::new();
-            let mut db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let mut db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
 
             let mut map = HashMap::<Digest, Digest>::default();
             const ELEMENTS: u64 = 10;
@@ -977,9 +1006,89 @@ mod test {
             // Close & reopen the db, making sure the re-opened db has exactly the same state.
             let root_hash = db.root(&mut hasher);
             db.close().await.unwrap();
-            let db = open_db::<_, 32>(context.clone(), &mut hasher).await;
+            let db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
             assert_eq!(root_hash, db.root(&mut hasher));
             assert!(db.get(&k).await.unwrap().is_none());
         });
     }
+
+    /// This test builds a random database which we then commit, close, and reopen with a bitmap. We
+    /// then check that the bitmap state matches that of the database's active operations.
+    #[test_traced("WARN")]
+    pub fn test_any_restore_with_bitmap() {
+        // Number of elements to initially insert into the db.
+        const ELEMENTS: u64 = 1000;
+
+        // Use a non-deterministic rng seed to ensure each run is different.
+        let rng_seed = OsRng.next_u64();
+        // Log the seed with high visibility to make failures reproducible.
+        warn!("rng_seed={}", rng_seed);
+        let mut rng = StdRng::seed_from_u64(rng_seed);
+
+        let executor = deterministic::Runner::default();
+        executor.start(|context| async move {
+            let mut hasher = Sha256::new();
+            let mut db = open_db::<_, 32>(context.clone(), &mut hasher, None).await;
+
+            for i in 0u64..ELEMENTS {
+                let k = hash(&i.to_be_bytes());
+                let v = hash(&rng.next_u32().to_be_bytes());
+                db.update(&mut hasher, k, v).await.unwrap();
+            }
+
+            // Randomly update / delete them. We use a delete frequency that is 1/7th of the update
+            // frequency.
+            for _ in 0u64..ELEMENTS * 10 {
+                let rand_key = hash(&(rng.next_u64() % ELEMENTS).to_be_bytes());
+                if rng.next_u32() % 7 == 0 {
+                    db.delete(&mut hasher, rand_key).await.unwrap();
+                    continue;
+                }
+                let v = hash(&rng.next_u32().to_be_bytes());
+                db.update(&mut hasher, rand_key, v).await.unwrap();
+                if rng.next_u32() % 20 == 0 {
+                    // Commit every ~20 updates.
+                    db.commit(&mut hasher, None).await.unwrap();
+                }
+            }
+            db.commit(&mut hasher, None).await.unwrap();
+
+            // Create a bitmap based on the current db's pruned/inactive state.
+            let mut bitmap = Bitmap::<_, 32>::new();
+            for _ in 0..db.inactivity_floor_loc {
+                bitmap.append(&mut hasher, false);
+            }
+            assert_eq!(bitmap.bit_count(), db.inactivity_floor_loc);
+
+            // Close the db, then re-open it with a bitmap.
+            let root = db.root(&mut hasher);
+            db.close().await.unwrap();
+            let db = open_db::<_, 32>(context.clone(), &mut hasher, Some(&mut bitmap)).await;
+            assert_eq!(db.root(&mut hasher), root);
+
+            // Check the bitmap state matches that of the snapshot.
+            let items = db.log.size().await.unwrap();
+            assert_eq!(bitmap.bit_count(), items);
+            let mut active_positions = HashSet::new();
+            for pos in db.inactivity_floor_loc..items {
+                let item = db.log.read(pos).await.unwrap();
+                let iter = db.snapshot.get_iter(&item.to_key());
+                let mut active = false;
+                for loc in iter {
+                    if *loc == pos {
+                        // Found an active op.
+                        active_positions.insert(pos);
+                        active = true;
+                        break;
+                    }
+                }
+                assert_eq!(bitmap.get_bit(pos), active, "position {} should match", pos);
+            }
+            for pos in db.inactivity_floor_loc..items {
+                if !active_positions.contains(&pos) {
+                    assert!(!bitmap.get_bit(pos));
+                }
+            }
+        });
+    }
 }

storage/src/mmr/bitmap.rs

@@ -311,9 +311,17 @@ impl<H: CHasher, const N: usize> Bitmap<H, N> {
     /// Convert a bit offset into the index of the chunk it belongs to within self.bitmap. Panics if
     /// the bit doesn't exist or has been pruned.
     fn chunk_index(&self, bit_offset: u64) -> usize {
-        assert!(bit_offset < self.bit_count(), "out of bounds");
+        assert!(
+            bit_offset < self.bit_count(),
+            "out of bounds: {}",
+            bit_offset
+        );
         let chunk_pos = Self::chunk_pos(bit_offset);
-        assert!(chunk_pos >= self.pruned_chunks, "bit pruned");
+        assert!(
+            chunk_pos >= self.pruned_chunks,
+            "bit pruned: {}",
+            bit_offset
+        );
 
         chunk_pos - self.pruned_chunks
     }