add ancient get_unique_accounts_from_storage_for_combining_ancient_slots (#30176)

runtime/src/accounts_db.rs

@@ -17257,7 +17257,10 @@ pub mod tests {
         }
     }
 
-    fn get_all_accounts(db: &AccountsDb, slots: Range<Slot>) -> Vec<(Pubkey, AccountSharedData)> {
+    pub(crate) fn get_all_accounts(
+        db: &AccountsDb,
+        slots: Range<Slot>,
+    ) -> Vec<(Pubkey, AccountSharedData)> {
         slots
             .filter_map(|slot| {
                 let storage = db.storage.get_slot_storage_entry(slot);
@@ -17273,7 +17276,7 @@ pub mod tests {
             .collect::<Vec<_>>()
     }
 
-    fn compare_all_accounts(
+    pub(crate) fn compare_all_accounts(
         one: &[(Pubkey, AccountSharedData)],
         two: &[(Pubkey, AccountSharedData)],
     ) {

runtime/src/ancient_append_vecs.rs

@@ -5,7 +5,7 @@
 //! Otherwise, an ancient append vec is the same as any other append vec
 use {
     crate::{
-        accounts_db::{AccountStorageEntry, AccountsDb},
+        accounts_db::{AccountStorageEntry, AccountsDb, GetUniqueAccountsResult},
         append_vec::{AppendVec, StoredAccountMeta},
     },
     rand::{thread_rng, Rng},
@@ -225,6 +225,26 @@ impl AccountsDb {
         }
         infos
     }
+
+    /// for each slot in 'ancient_slots', collect all accounts in that slot
+    /// return the collection of accounts by slot
+    #[allow(dead_code)]
+    fn get_unique_accounts_from_storage_for_combining_ancient_slots<'a>(
+        &self,
+        ancient_slots: &'a [SlotInfo],
+    ) -> Vec<(&'a SlotInfo, GetUniqueAccountsResult<'a>)> {
+        let mut accounts_to_combine = Vec::with_capacity(ancient_slots.len());
+
+        for info in ancient_slots {
+            let unique_accounts = self.get_unique_accounts_from_storage_for_shrink(
+                &info.storage,
+                &self.shrink_ancient_stats.shrink_stats,
+            );
+            accounts_to_combine.push((info, unique_accounts));
+        }
+
+        accounts_to_combine
+    }
 }
 
 /// a set of accounts need to be stored.
@@ -318,8 +338,8 @@ pub mod tests {
             accounts_db::{
                 get_temp_accounts_paths,
                 tests::{
-                    create_db_with_storages_and_index, create_storages_and_update_index,
-                    remove_account_for_tests,
+                    compare_all_accounts, create_db_with_storages_and_index,
+                    create_storages_and_update_index, get_all_accounts, remove_account_for_tests,
                 },
             },
             append_vec::{AccountMeta, StoredAccountMeta, StoredMeta},
@@ -329,10 +349,109 @@ pub mod tests {
             hash::Hash,
             pubkey::Pubkey,
         },
+        std::ops::Range,
         strum::IntoEnumIterator,
         strum_macros::EnumIter,
     };
 
+    fn get_sample_storages(
+        slots: usize,
+    ) -> (
+        AccountsDb,
+        Vec<Arc<AccountStorageEntry>>,
+        Range<u64>,
+        Vec<SlotInfo>,
+    ) {
+        let alive = true;
+        let (db, slot1) = create_db_with_storages_and_index(alive, slots, None);
+        let original_stores = (0..slots)
+            .filter_map(|slot| db.storage.get_slot_storage_entry((slot as Slot) + slot1))
+            .collect::<Vec<_>>();
+        let slot_infos = original_stores
+            .iter()
+            .map(|storage| SlotInfo {
+                storage: Arc::clone(storage),
+                slot: storage.slot(),
+                capacity: 0,
+                alive_bytes: 0,
+                should_shrink: false,
+            })
+            .collect();
+        (
+            db,
+            original_stores,
+            slot1..(slot1 + slots as Slot),
+            slot_infos,
+        )
+    }
+
+    fn unique_to_accounts<'a>(
+        one: impl Iterator<Item = &'a GetUniqueAccountsResult<'a>>,
+    ) -> Vec<(Pubkey, AccountSharedData)> {
+        one.flat_map(|result| {
+            result
+                .stored_accounts
+                .iter()
+                .map(|result| (*result.pubkey(), result.to_account_shared_data()))
+        })
+        .collect()
+    }
+
+    pub(crate) fn compare_all_vec_accounts<'a>(
+        one: impl Iterator<Item = &'a GetUniqueAccountsResult<'a>>,
+        two: impl Iterator<Item = &'a GetUniqueAccountsResult<'a>>,
+    ) {
+        compare_all_accounts(&unique_to_accounts(one), &unique_to_accounts(two));
+    }
+
+    #[test]
+    fn test_get_unique_accounts_from_storage_for_combining_ancient_slots() {
+        for num_slots in 0..3 {
+            for reverse in [false, true] {
+                let (db, storages, slots, mut infos) = get_sample_storages(num_slots);
+                let original_results = storages
+                    .iter()
+                    .map(|store| db.get_unique_accounts_from_storage(store))
+                    .collect::<Vec<_>>();
+                if reverse {
+                    // reverse the contents for further testing
+                    infos = infos.into_iter().rev().collect();
+                }
+                let results =
+                    db.get_unique_accounts_from_storage_for_combining_ancient_slots(&infos);
+
+                let all_accounts = get_all_accounts(&db, slots);
+                assert_eq!(all_accounts.len(), num_slots);
+
+                compare_all_vec_accounts(
+                    original_results.iter(),
+                    results.iter().map(|(_, accounts)| accounts),
+                );
+                compare_all_accounts(
+                    &all_accounts,
+                    &unique_to_accounts(results.iter().map(|(_, accounts)| accounts)),
+                );
+
+                let map = |info: &SlotInfo| {
+                    (
+                        info.storage.append_vec_id(),
+                        info.slot,
+                        info.capacity,
+                        info.alive_bytes,
+                        info.should_shrink,
+                    )
+                };
+                assert_eq!(
+                    infos.iter().map(map).collect::<Vec<_>>(),
+                    results
+                        .into_iter()
+                        .map(|(info, _)| map(info))
+                        .collect::<Vec<_>>()
+                );
+            }
+        }
+    }
+
     #[test]
     fn test_accounts_to_store_simple() {
         let map = vec![];