tests for old_locs returned from commit

Author: roberto-bayardo

storage/src/adb/any.rs

@@ -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,11 @@ 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
+    /// reflecting its activity status. The bitmap is expected to already have a number of bits
+    /// corresponding to the portion of the database below the inactivity floor, and this method
+    /// will panic otherwise.
     pub async fn init<const N: usize>(
         context: E,
         hasher: &mut H,
@@ -159,6 +164,9 @@ impl<E: RStorage + Clock + Metrics, K: Array, V: Array, H: CHasher> Any<E, K, V,
         let pruned_to_pos = mmr.pruned_to_pos();
         let start_leaf_num = leaf_pos_to_num(pruned_to_pos).unwrap();
 
+        if let Some(ref bitmap) = bitmap {
+            assert_eq!(start_leaf_num, bitmap.bit_count());
+        }
         for i in start_leaf_num..log_size {
             let op: Operation<K, V> = log.read(i).await?;
             match op.to_type() {
@@ -167,6 +175,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 +197,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 {
@@ -273,10 +291,10 @@ impl<E: RStorage + Clock + Metrics, K: Array, V: Array, H: CHasher> Any<E, K, V,
             .map(|pos| leaf_pos_to_num(pos).unwrap())
     }
 
-    /// Updates `key` to have value `value`.  If the key already has this same value, then this is a
+    /// 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,
@@ -305,8 +323,7 @@ impl<E: RStorage + Clock + Metrics, K: Array, V: Array, H: CHasher> Any<E, K, V,
 
     /// Delete `key` and its value from the db. Deleting a key that already has no value is a no-op.
     /// The operation is reflected in the snapshot, but will be subject to rollback until the next
-    /// successful `commit`. Returns an (inactive_position, active_position) tuple for the key, or
-    /// none if the update was a no-op.
+    /// successful `commit`. Returns the location of the deleted value for the key (if any).
     pub async fn delete(&mut self, hasher: &mut H, key: K) -> Result<Option<u64>, Error> {
         let mut loc_iter = self.snapshot.remove_iter(&key);
         for loc in &mut loc_iter {
@@ -406,7 +423,8 @@ impl<E: RStorage + Clock + Metrics, K: Array, V: Array, H: CHasher> Any<E, K, V,
 
     /// Commit any pending operations to the db, ensuring they are persisted to disk & recoverable
     /// upon return from this function. Also raises the inactivity floor according to the schedule,
-    /// and prunes those operations below it.
+    /// and prunes those operations below it. If an old_locs vector is provided, then any locations
+    /// that were flipped from active to inactive during this operation are added to it.
     pub async fn commit(
         &mut self,
         hasher: &mut H,
@@ -574,12 +592,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 +611,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 +621,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 +633,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 +663,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 +707,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 +746,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 +774,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 +797,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 +838,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 +908,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 +935,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 +945,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 +957,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 +966,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 +986,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 +1011,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));
+                }
+            }
+        });
+    }
 }