Add base_slot to get_snapshot_storages() (solana-labs#19348)

core/tests/snapshots.rs

@@ -282,7 +282,7 @@ mod tests {
             snapshot_path,
             last_bank.src.slot_deltas(&last_bank.src.roots()),
             &snapshot_config.snapshot_package_output_path,
-            last_bank.get_snapshot_storages(),
+            last_bank.get_snapshot_storages(None),
             ArchiveFormat::TarBzip2,
             snapshot_version,
             None,
@@ -364,7 +364,7 @@ mod tests {
 
         // Take snapshot of zeroth bank
         let bank0 = bank_forks.get(0).unwrap();
-        let storages = bank0.get_snapshot_storages();
+        let storages = bank0.get_snapshot_storages(None);
         snapshot_utils::add_bank_snapshot(snapshot_path, bank0, &storages, snapshot_version)
             .unwrap();
 
@@ -424,7 +424,7 @@ mod tests {
             if slot == saved_slot as u64 {
                 // Find the relevant snapshot storages
                 let snapshot_storage_files: HashSet<_> = bank_forks[slot]
-                    .get_snapshot_storages()
+                    .get_snapshot_storages(None)
                     .into_iter()
                     .flatten()
                     .map(|s| s.get_path())
@@ -763,7 +763,7 @@ mod tests {
             &bank_snapshot_info,
             &snapshot_config.snapshot_path,
             &snapshot_config.snapshot_package_output_path,
-            bank.get_snapshot_storages(),
+            bank.get_snapshot_storages(None),
             snapshot_config.archive_format,
             snapshot_config.snapshot_version,
             None,
@@ -788,7 +788,7 @@ mod tests {
             .find(|elem| elem.slot == slot)
             .ok_or_else(|| Error::new(ErrorKind::Other, "did not find snapshot with this path"))?;
         let storages = {
-            let mut storages = bank.get_snapshot_storages();
+            let mut storages = bank.get_snapshot_storages(Some(incremental_snapshot_base_slot));
             snapshot_utils::filter_snapshot_storages_for_incremental_snapshot(
                 &mut storages,
                 incremental_snapshot_base_slot,

runtime/src/accounts_db.rs

@@ -4834,7 +4834,7 @@ impl AccountsDb {
             };
 
             let mut collect_time = Measure::start("collect");
-            let (combined_maps, slots) = self.get_snapshot_storages(slot, Some(ancestors));
+            let (combined_maps, slots) = self.get_snapshot_storages(slot, None, Some(ancestors));
             collect_time.stop();
 
             let mut sort_time = Measure::start("sort_storages");
@@ -5884,6 +5884,7 @@ impl AccountsDb {
     pub fn get_snapshot_storages(
         &self,
         snapshot_slot: Slot,
+        snapshot_base_slot: Option<Slot>,
         ancestors: Option<&Ancestors>,
     ) -> (SnapshotStorages, Vec<Slot>) {
         let mut m = Measure::start("get slots");
@@ -5905,6 +5906,8 @@ impl AccountsDb {
                         .iter()
                         .filter_map(|slot| {
                             if *slot <= snapshot_slot
+                                && snapshot_base_slot
+                                    .map_or(true, |snapshot_base_slot| *slot > snapshot_base_slot)
                                 && (self.accounts_index.is_root(*slot)
                                     || ancestors
                                         .map(|ancestors| ancestors.contains_key(slot))
@@ -6635,7 +6638,7 @@ pub mod tests {
         accounts.store_uncached(slot, &to_store[..]);
         accounts.add_root(slot);
 
-        let (storages, slots) = accounts.get_snapshot_storages(slot, None);
+        let (storages, slots) = accounts.get_snapshot_storages(slot, None, None);
         assert_eq!(storages.len(), slots.len());
         storages
             .iter()
@@ -9203,7 +9206,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).0.is_empty());
+        assert!(db.get_snapshot_storages(0, None, None).0.is_empty());
     }
 
     #[test]
@@ -9218,10 +9221,13 @@ pub mod tests {
 
         db.add_root(base_slot);
         db.store_uncached(base_slot, &[(&key, &account)]);
-        assert!(db.get_snapshot_storages(before_slot, None).0.is_empty());
+        assert!(db
+            .get_snapshot_storages(before_slot, None, None)
+            .0
+            .is_empty());
 
-        assert_eq!(1, db.get_snapshot_storages(base_slot, None).0.len());
-        assert_eq!(1, db.get_snapshot_storages(after_slot, None).0.len());
+        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());
     }
 
     #[test]
@@ -9241,10 +9247,13 @@ pub mod tests {
             .unwrap()
             .clear();
         db.add_root(base_slot);
-        assert!(db.get_snapshot_storages(after_slot, None).0.is_empty());
+        assert!(db
+            .get_snapshot_storages(after_slot, None, None)
+            .0
+            .is_empty());
 
         db.store_uncached(base_slot, &[(&key, &account)]);
-        assert_eq!(1, db.get_snapshot_storages(after_slot, None).0.len());
+        assert_eq!(1, db.get_snapshot_storages(after_slot, None, None).0.len());
     }
 
     #[test]
@@ -9257,10 +9266,13 @@ pub mod tests {
         let after_slot = base_slot + 1;
 
         db.store_uncached(base_slot, &[(&key, &account)]);
-        assert!(db.get_snapshot_storages(after_slot, None).0.is_empty());
+        assert!(db
+            .get_snapshot_storages(after_slot, None, None)
+            .0
+            .is_empty());
 
         db.add_root(base_slot);
-        assert_eq!(1, db.get_snapshot_storages(after_slot, None).0.len());
+        assert_eq!(1, db.get_snapshot_storages(after_slot, None, None).0.len());
     }
 
     #[test]
@@ -9274,7 +9286,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).0.len());
+        assert_eq!(1, db.get_snapshot_storages(after_slot, None, None).0.len());
 
         db.storage
             .get_slot_stores(0)
@@ -9285,7 +9297,32 @@ pub mod tests {
             .next()
             .unwrap()
             .remove_account(0, true);
-        assert!(db.get_snapshot_storages(after_slot, None).0.is_empty());
+        assert!(db
+            .get_snapshot_storages(after_slot, None, None)
+            .0
+            .is_empty());
+    }
+
+    #[test]
+    fn test_get_snapshot_storages_with_base_slot() {
+        let db = AccountsDb::new(Vec::new(), &ClusterType::Development);
+
+        let key = Pubkey::default();
+        let account = AccountSharedData::new(1, 0, &key);
+
+        let slot = 10;
+        db.store_uncached(slot, &[(&key, &account)]);
+        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()
+        );
     }
 
     #[test]
@@ -9449,7 +9486,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).0;
+        let snapshot_stores = accounts.get_snapshot_storages(current_slot, None, None).0;
         let total_accounts: usize = snapshot_stores
             .iter()
             .flatten()

runtime/src/bank.rs

@@ -483,13 +483,6 @@ impl BankRc {
             bank_id_generator: Arc::new(AtomicU64::new(0)),
         }
     }
-
-    pub fn get_snapshot_storages(&self, slot: Slot) -> SnapshotStorages {
-        self.accounts
-            .accounts_db
-            .get_snapshot_storages(slot, None)
-            .0
-    }
 }
 
 #[derive(Default, Debug, AbiExample)]
@@ -4904,8 +4897,12 @@ impl Bank {
         )
     }
 
-    pub fn get_snapshot_storages(&self) -> SnapshotStorages {
-        self.rc.get_snapshot_storages(self.slot())
+    pub fn get_snapshot_storages(&self, base_slot: Option<Slot>) -> SnapshotStorages {
+        self.rc
+            .accounts
+            .accounts_db
+            .get_snapshot_storages(self.slot(), base_slot, None)
+            .0
     }
 
     #[must_use]

runtime/src/serde_snapshot/tests.rs

@@ -28,7 +28,9 @@ fn copy_append_vecs<P: AsRef<Path>>(
     accounts_db: &AccountsDb,
     output_dir: P,
 ) -> std::io::Result<UnpackedAppendVecMap> {
-    let storage_entries = accounts_db.get_snapshot_storages(Slot::max_value(), None).0;
+    let storage_entries = accounts_db
+        .get_snapshot_storages(Slot::max_value(), None, None)
+        .0;
     let mut unpacked_append_vec_map = UnpackedAppendVecMap::new();
     for storage in storage_entries.iter().flatten() {
         let storage_path = storage.get_path();
@@ -151,7 +153,7 @@ fn test_accounts_serialize_style(serde_style: SerdeStyle) {
         &mut writer,
         &*accounts.accounts_db,
         0,
-        &accounts.accounts_db.get_snapshot_storages(0, None).0,
+        &accounts.accounts_db.get_snapshot_storages(0, None, None).0,
     )
     .unwrap();
 
@@ -203,7 +205,7 @@ fn test_bank_serialize_style(serde_style: SerdeStyle) {
     bank2.squash();
     bank2.force_flush_accounts_cache();
 
-    let snapshot_storages = bank2.get_snapshot_storages();
+    let snapshot_storages = bank2.get_snapshot_storages(None);
     let mut buf = vec![];
     let mut writer = Cursor::new(&mut buf);
     crate::serde_snapshot::bank_to_stream(
@@ -259,7 +261,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).0;
+    let snapshot_storages = accounts.get_snapshot_storages(slot, None, None).0;
     accountsdb_to_stream(
         SerdeStyle::Newer,
         &mut writer,
@@ -321,7 +323,7 @@ mod test_bank_serialize {
             .rc
             .accounts
             .accounts_db
-            .get_snapshot_storages(0, None)
+            .get_snapshot_storages(0, None, None)
             .0;
         // ensure there is a single snapshot storage example for ABI digesting
         assert_eq!(snapshot_storages.len(), 1);

runtime/src/snapshot_utils.rs

@@ -1541,7 +1541,7 @@ pub fn snapshot_bank(
     archive_format: &ArchiveFormat,
     hash_for_testing: Option<Hash>,
 ) -> Result<()> {
-    let storages = root_bank.get_snapshot_storages();
+    let storages = root_bank.get_snapshot_storages(None);
     let mut add_snapshot_time = Measure::start("add-snapshot-ms");
     add_bank_snapshot(snapshots_dir, root_bank, &storages, snapshot_version)?;
     add_snapshot_time.stop();
@@ -1592,7 +1592,7 @@ pub fn bank_to_full_snapshot_archive(
     bank.rehash(); // Bank accounts may have been manually modified by the caller
 
     let temp_dir = tempfile::tempdir_in(snapshots_dir)?;
-    let storages = bank.get_snapshot_storages();
+    let storages = bank.get_snapshot_storages(None);
     let bank_snapshot_info = add_bank_snapshot(&temp_dir, bank, &storages, snapshot_version)?;
 
     package_process_and_archive_full_snapshot(
@@ -1636,7 +1636,7 @@ pub fn bank_to_incremental_snapshot_archive(
 
     let temp_dir = tempfile::tempdir_in(snapshots_dir)?;
     let storages = {
-        let mut storages = bank.get_snapshot_storages();
+        let mut storages = bank.get_snapshot_storages(Some(full_snapshot_slot));
         filter_snapshot_storages_for_incremental_snapshot(&mut storages, full_snapshot_slot);
         storages
     };

runtime/src/system_instruction_processor.rs

@@ -1386,9 +1386,9 @@ mod tests {
             bank.squash();
             bank.force_flush_accounts_cache();
             // do clean and assert that it actually did its job
-            assert_eq!(3, bank.get_snapshot_storages().len());
+            assert_eq!(3, bank.get_snapshot_storages(None).len());
             bank.clean_accounts(false, false, None);
-            assert_eq!(2, bank.get_snapshot_storages().len());
+            assert_eq!(2, bank.get_snapshot_storages(None).len());
         });
     }