AccountsDb::get_snapshot_storages() takes a Range for slots (#29054)

runtime/src/accounts_db.rs

@@ -7378,8 +7378,7 @@ impl AccountsDb {
                 }
 
                 let mut collect_time = Measure::start("collect");
-                let (combined_maps, slots) =
-                    self.get_snapshot_storages(slot, None, config.ancestors);
+                let (combined_maps, slots) = self.get_snapshot_storages(..=slot, config.ancestors);
                 collect_time.stop();
 
                 let mut sort_time = Measure::start("sort_storages");
@@ -8677,10 +8676,10 @@ impl AccountsDb {
         }
     }
 
+    /// Get storages to use for snapshots, for the requested slots
     pub fn get_snapshot_storages(
         &self,
-        snapshot_slot: Slot,
-        snapshot_base_slot: Option<Slot>,
+        requested_slots: impl RangeBounds<Slot> + Sync,
         ancestors: Option<&Ancestors>,
     ) -> (SnapshotStorages, Vec<Slot>) {
         let mut m = Measure::start("get slots");
@@ -8701,9 +8700,7 @@ impl AccountsDb {
                     slots
                         .iter()
                         .filter_map(|slot| {
-                            if *slot <= snapshot_slot
-                                && snapshot_base_slot
-                                    .map_or(true, |snapshot_base_slot| *slot > snapshot_base_slot)
+                            if requested_slots.contains(slot)
                                 && (self.accounts_index.is_alive_root(*slot)
                                     || ancestors
                                         .map(|ancestors| ancestors.contains_key(slot))
@@ -10266,7 +10263,7 @@ pub mod tests {
         accounts.store_uncached(slot, &to_store[..]);
         accounts.add_root(slot);
 
-        let (storages, slots) = accounts.get_snapshot_storages(slot, None, None);
+        let (storages, slots) = accounts.get_snapshot_storages(..=slot, None);
         assert_eq!(storages.len(), slots.len());
         storages
             .iter()
@@ -13074,7 +13071,7 @@ pub mod tests {
     #[test]
     fn test_get_snapshot_storages_empty() {
         let db = AccountsDb::new(Vec::new(), &ClusterType::Development);
-        assert!(db.get_snapshot_storages(0, None, None).0.is_empty());
+        assert!(db.get_snapshot_storages(..=0, None).0.is_empty());
     }
 
     #[test]
@@ -13089,13 +13086,10 @@ pub mod tests {
 
         db.add_root(base_slot);
         db.store_uncached(base_slot, &[(&key, &account)]);
-        assert!(db
-            .get_snapshot_storages(before_slot, None, None)
-            .0
-            .is_empty());
+        assert!(db.get_snapshot_storages(..=before_slot, None).0.is_empty());
 
-        assert_eq!(1, db.get_snapshot_storages(base_slot, None, None).0.len());
-        assert_eq!(1, db.get_snapshot_storages(after_slot, None, None).0.len());
+        assert_eq!(1, db.get_snapshot_storages(..=base_slot, None).0.len());
+        assert_eq!(1, db.get_snapshot_storages(..=after_slot, None).0.len());
     }
 
     #[test]
@@ -13115,13 +13109,10 @@ pub mod tests {
             .unwrap()
             .clear();
         db.add_root(base_slot);
-        assert!(db
-            .get_snapshot_storages(after_slot, None, None)
-            .0
-            .is_empty());
+        assert!(db.get_snapshot_storages(..=after_slot, None).0.is_empty());
 
         db.store_uncached(base_slot, &[(&key, &account)]);
-        assert_eq!(1, db.get_snapshot_storages(after_slot, None, None).0.len());
+        assert_eq!(1, db.get_snapshot_storages(..=after_slot, None).0.len());
     }
 
     #[test]
@@ -13134,13 +13125,10 @@ pub mod tests {
         let after_slot = base_slot + 1;
 
         db.store_uncached(base_slot, &[(&key, &account)]);
-        assert!(db
-            .get_snapshot_storages(after_slot, None, None)
-            .0
-            .is_empty());
+        assert!(db.get_snapshot_storages(..=after_slot, None).0.is_empty());
 
         db.add_root(base_slot);
-        assert_eq!(1, db.get_snapshot_storages(after_slot, None, None).0.len());
+        assert_eq!(1, db.get_snapshot_storages(..=after_slot, None).0.len());
     }
 
     #[test]
@@ -13154,7 +13142,7 @@ pub mod tests {
 
         db.store_uncached(base_slot, &[(&key, &account)]);
         db.add_root(base_slot);
-        assert_eq!(1, db.get_snapshot_storages(after_slot, None, None).0.len());
+        assert_eq!(1, db.get_snapshot_storages(..=after_slot, None).0.len());
 
         db.storage
             .get_slot_stores(0)
@@ -13165,10 +13153,7 @@ pub mod tests {
             .next()
             .unwrap()
             .remove_account(0, true);
-        assert!(db
-            .get_snapshot_storages(after_slot, None, None)
-            .0
-            .is_empty());
+        assert!(db.get_snapshot_storages(..=after_slot, None).0.is_empty());
     }
 
     #[test]
@@ -13183,14 +13168,9 @@ pub mod tests {
         db.add_root(slot);
         assert_eq!(
             0,
-            db.get_snapshot_storages(slot + 1, Some(slot), None).0.len()
-        );
-        assert_eq!(
-            1,
-            db.get_snapshot_storages(slot + 1, Some(slot - 1), None)
-                .0
-                .len()
+            db.get_snapshot_storages(slot + 1..=slot + 1, None).0.len()
         );
+        assert_eq!(1, db.get_snapshot_storages(slot..=slot + 1, None).0.len());
     }
 
     #[test]
@@ -13354,7 +13334,7 @@ pub mod tests {
         accounts.store_uncached(current_slot, &[(&pubkey2, &zero_lamport_account)]);
         accounts.store_uncached(current_slot, &[(&pubkey3, &zero_lamport_account)]);
 
-        let snapshot_stores = accounts.get_snapshot_storages(current_slot, None, None).0;
+        let snapshot_stores = accounts.get_snapshot_storages(..=current_slot, None).0;
         let total_accounts: usize = snapshot_stores
             .iter()
             .flatten()

runtime/src/bank.rs

@@ -6944,11 +6944,19 @@ impl Bank {
             .check_complete()
     }
 
+    /// Get this bank's storages to use for snapshots.
+    ///
+    /// If a base slot is provided, return only the storages that are *higher* than this slot.
     pub fn get_snapshot_storages(&self, base_slot: Option<Slot>) -> SnapshotStorages {
+        // if a base slot is provided, request storages starting at the slot *after*
+        let start_slot = base_slot.map_or(0, |slot| slot.saturating_add(1));
+        // we want to *include* the storage at our slot
+        let requested_slots = start_slot..=self.slot();
+
         self.rc
             .accounts
             .accounts_db
-            .get_snapshot_storages(self.slot(), base_slot, None)
+            .get_snapshot_storages(requested_slots, None)
             .0
     }
 

runtime/src/serde_snapshot/tests.rs

@@ -25,6 +25,7 @@ use {
     },
     std::{
         io::{BufReader, Cursor},
+        ops::RangeFull,
         path::Path,
         sync::{Arc, RwLock},
     },
@@ -36,9 +37,7 @@ fn copy_append_vecs<P: AsRef<Path>>(
     accounts_db: &AccountsDb,
     output_dir: P,
 ) -> std::io::Result<StorageAndNextAppendVecId> {
-    let storage_entries = accounts_db
-        .get_snapshot_storages(Slot::max_value(), None, None)
-        .0;
+    let storage_entries = accounts_db.get_snapshot_storages(RangeFull, None).0;
     let storage: AccountStorageMap = AccountStorageMap::with_capacity(storage_entries.len());
     let mut next_append_vec_id = 0;
     for storage_entry in storage_entries.into_iter().flatten() {
@@ -187,7 +186,7 @@ fn test_accounts_serialize_style(serde_style: SerdeStyle) {
         &mut writer,
         &accounts.accounts_db,
         0,
-        &accounts.accounts_db.get_snapshot_storages(0, None, None).0,
+        &accounts.accounts_db.get_snapshot_storages(..=0, None).0,
     )
     .unwrap();
 
@@ -430,7 +429,7 @@ pub(crate) fn reconstruct_accounts_db_via_serialization(
     slot: Slot,
 ) -> AccountsDb {
     let mut writer = Cursor::new(vec![]);
-    let snapshot_storages = accounts.get_snapshot_storages(slot, None, None).0;
+    let snapshot_storages = accounts.get_snapshot_storages(..=slot, None).0;
     accountsdb_to_stream(
         SerdeStyle::Newer,
         &mut writer,
@@ -714,7 +713,7 @@ mod test_bank_serialize {
             .rc
             .accounts
             .accounts_db
-            .get_snapshot_storages(0, None, None)
+            .get_snapshot_storages(..=0, None)
             .0;
         // ensure there is a single snapshot storage example for ABI digesting
         assert_eq!(snapshot_storages.len(), 1);

runtime/src/snapshot_minimizer.rs

@@ -276,7 +276,7 @@ impl<'a> SnapshotMinimizer<'a> {
     ) -> (Vec<Slot>, Vec<Arc<AccountStorageEntry>>) {
         let snapshot_storages = self
             .accounts_db()
-            .get_snapshot_storages(self.starting_slot, None, None)
+            .get_snapshot_storages(..=self.starting_slot, None)
             .0;
 
         let dead_slots = Mutex::new(Vec::new());
@@ -680,7 +680,7 @@ mod tests {
         };
         minimizer.minimize_accounts_db();
 
-        let snapshot_storages = accounts.get_snapshot_storages(current_slot, None, None).0;
+        let snapshot_storages = accounts.get_snapshot_storages(..=current_slot, None).0;
         assert_eq!(snapshot_storages.len(), 3);
 
         let mut account_count = 0;