ancient append vec tests (#28635)

runtime/src/accounts_db.rs

@@ -136,6 +136,10 @@ const CACHE_VIRTUAL_WRITE_VERSION: StoredMetaWriteVersion = 0;
 pub(crate) const CACHE_VIRTUAL_OFFSET: Offset = 0;
 const CACHE_VIRTUAL_STORED_SIZE: StoredSize = 0;
 
+// When getting accounts for shrinking from the index, this is the # of accounts to lookup per thread.
+// This allows us to split up accounts index accesses across multiple threads.
+const SHRINK_COLLECT_CHUNK_SIZE: usize = 50;
+
 /// temporary enum during feature activation of
 /// ignore slot when calculating an account hash #28420
 #[derive(Debug, Clone, Copy)]
@@ -3800,9 +3804,8 @@ impl AccountsDb {
             .fetch_add(len as u64, Ordering::Relaxed);
         let all_are_zero_lamports_collect = Mutex::new(true);
         self.thread_pool_clean.install(|| {
-            let chunk_size = 50; // # accounts/thread
             stored_accounts
-                .par_chunks(chunk_size)
+                .par_chunks(SHRINK_COLLECT_CHUNK_SIZE)
                 .for_each(|stored_accounts| {
                     let LoadAccountsIndexForShrink {
                         alive_total,
@@ -9654,6 +9657,7 @@ pub mod tests {
             inline_spl_token,
         },
         assert_matches::assert_matches,
+        itertools::Itertools,
         rand::{prelude::SliceRandom, thread_rng, Rng},
         solana_sdk::{
             account::{
@@ -17038,6 +17042,116 @@ pub mod tests {
         }
     }
 
+    #[test]
+    fn test_shrink_collect_simple() {
+        solana_logger::setup();
+        let account_counts = [
+            1,
+            SHRINK_COLLECT_CHUNK_SIZE,
+            SHRINK_COLLECT_CHUNK_SIZE + 1,
+            SHRINK_COLLECT_CHUNK_SIZE * 2,
+        ];
+        let max_num_accounts = *account_counts.iter().max().unwrap();
+        let pubkeys = (0..max_num_accounts)
+            .map(|_| solana_sdk::pubkey::new_rand())
+            .collect::<Vec<_>>();
+        // write accounts, maybe remove from index
+        // check shrink_collect results
+        for lamports in [0, 1] {
+            for space in [0, 8] {
+                if lamports == 0 && space != 0 {
+                    // illegal - zero lamport accounts are written with 0 space
+                    continue;
+                }
+                for account_count in account_counts {
+                    for alive in [false, true] {
+                        debug!("space: {space}, lamports: {lamports}, alive: {alive}, account_count: {account_count}");
+                        let size = 100000;
+                        let db = AccountsDb::new_single_for_tests();
+                        let slot5 = 5;
+                        let inserted_store = db.create_and_insert_store(slot5, size, "test");
+                        let account = AccountSharedData::new(
+                            lamports,
+                            space,
+                            AccountSharedData::default().owner(),
+                        );
+                        for pubkey in pubkeys.iter().take(account_count) {
+                            // store in append vec and index
+                            db.store_uncached(slot5, &[(pubkey, &account)]);
+                            if !alive {
+                                // remove from index so pubkey is 'dead'
+                                db.accounts_index.purge_exact(
+                                    pubkey,
+                                    &([slot5].into_iter().collect::<HashSet<_>>()),
+                                    &mut Vec::default(),
+                                );
+                            }
+                        }
+                        let storages = vec![inserted_store];
+                        let mut stored_accounts = Vec::default();
+                        assert_eq!(storages.len(), 1);
+                        let shrink_collect = db.shrink_collect(
+                            storages.iter(),
+                            &mut stored_accounts,
+                            &ShrinkStats::default(),
+                        );
+                        if alive {
+                            assert!(shrink_collect.unrefed_pubkeys.is_empty());
+                            assert_eq!(shrink_collect.alive_accounts.len(), account_count);
+                            assert_eq!(
+                                shrink_collect
+                                    .alive_accounts
+                                    .iter()
+                                    .map(|(pubkey, _)| *pubkey)
+                                    .sorted()
+                                    .collect::<Vec<_>>(),
+                                pubkeys[..account_count]
+                                    .iter()
+                                    .sorted()
+                                    .cloned()
+                                    .collect::<Vec<_>>()
+                            );
+                            assert_eq!(
+                                shrink_collect.aligned_total,
+                                PAGE_SIZE
+                                    * if account_count >= 100 {
+                                        4
+                                    } else if account_count >= 50 {
+                                        2
+                                    } else {
+                                        1
+                                    }
+                            );
+                            // 136 is size of serialized account, with zero length data, and with padding to next aligned entry
+                            assert_eq!(shrink_collect.alive_total, (136 + space) * account_count);
+                        } else {
+                            assert!(shrink_collect.alive_accounts.is_empty());
+                            assert_eq!(shrink_collect.unrefed_pubkeys.len(), account_count);
+                            assert_eq!(
+                                shrink_collect
+                                    .unrefed_pubkeys
+                                    .iter()
+                                    .sorted()
+                                    .cloned()
+                                    .collect::<Vec<_>>(),
+                                pubkeys[..account_count].iter().sorted().collect::<Vec<_>>()
+                            );
+                            assert_eq!(shrink_collect.aligned_total, 0);
+                            assert_eq!(shrink_collect.alive_total, 0);
+                        }
+                        assert_eq!(shrink_collect.original_bytes, 102400);
+                        assert_eq!(shrink_collect.store_ids, vec![storages[0].append_vec_id()]);
+                        assert_eq!(shrink_collect.total_starting_accounts, account_count);
+                        assert_eq!(
+                            shrink_collect.all_are_zero_lamports,
+                            lamports == 0 || !alive
+                        );
+                    }
+                }
+            }
+        }
+    }
+
     #[test]
     fn test_should_move_to_ancient_append_vec() {
         solana_logger::setup();