calculate_capitalization uses hash calculation

accounts-bench/src/main.rs

@@ -119,6 +119,7 @@ fn main() {
                 solana_sdk::clock::Slot::default(),
                 &ancestors,
                 None,
+                false,
             );
             time_store.stop();
             if results != results_store {

runtime/src/accounts.rs

@@ -619,8 +619,13 @@ impl Accounts {
             .collect()
     }
 
-    pub fn calculate_capitalization(&self, ancestors: &Ancestors) -> u64 {
-        self.accounts_db.unchecked_scan_accounts(
+    pub fn calculate_capitalization(&self, ancestors: &Ancestors, slot: Slot) -> u64 {
+        let cap = self
+            .accounts_db
+            .update_accounts_hash_with_index_option(false, false, slot, ancestors, None, true)
+            .1;
+        // debug code for the moment
+        let cap2 = self.accounts_db.unchecked_scan_accounts(
             "calculate_capitalization_scan_elapsed",
             ancestors,
             |total_capitalization: &mut u64, (_pubkey, loaded_account, _slot)| {
@@ -632,7 +637,9 @@ impl Accounts {
                     );
                 }
             },
-        )
+        );
+        assert_eq!(cap, cap2);
+        cap2
     }
 
     #[must_use]

runtime/src/accounts_cache.rs

@@ -48,6 +48,10 @@ impl SlotCacheInner {
         );
     }
 
+    pub fn get_all_pubkeys(&self) -> Vec<Pubkey> {
+        self.cache.iter().map(|item| *item.key()).collect()
+    }
+
     pub fn insert(
         &self,
         pubkey: &Pubkey,

runtime/src/accounts_db.rs

@@ -4328,15 +4328,16 @@ impl AccountsDb {
     }
 
     pub fn update_accounts_hash(&self, slot: Slot, ancestors: &Ancestors) -> (Hash, u64) {
-        self.update_accounts_hash_with_index_option(true, false, slot, ancestors, None)
+        self.update_accounts_hash_with_index_option(true, false, slot, ancestors, None, false)
     }
 
     pub fn update_accounts_hash_test(&self, slot: Slot, ancestors: &Ancestors) -> (Hash, u64) {
-        self.update_accounts_hash_with_index_option(true, true, slot, ancestors, None)
+        self.update_accounts_hash_with_index_option(true, true, slot, ancestors, None, false)
     }
 
     /// Scan through all the account storage in parallel
     fn scan_account_storage_no_bank<F, B>(
+        accounts_cache: Option<&AccountsCache>,
         snapshot_storages: &SortedStorages,
         scan_func: F,
     ) -> Vec<B>
@@ -4364,6 +4365,23 @@ impl AccountsDb {
                             });
                         }
                     }
+                    if let Some(cache) = accounts_cache {
+                        if let Some(slot_cache) = cache.slot_cache(slot) {
+                            let keys = slot_cache.get_all_pubkeys();
+                            for key in keys {
+                                if let Some(cached_account) = slot_cache.get_cloned(&key) {
+                                    let mut accessor = LoadedAccountAccessor::Cached(Some((
+                                        key,
+                                        Cow::Owned(cached_account),
+                                    )));
+                                    let account = accessor.get_loaded_account().unwrap();
+                                    scan_func(account, &mut retval, slot);
+                                };
+                            }
+                        }
+                    } else {
+                        error!("no write cache");
+                    }
                 }
                 retval
             })
@@ -4376,14 +4394,15 @@ impl AccountsDb {
         slot: Slot,
         ancestors: &Ancestors,
         check_hash: bool,
+        accounts_cache: Option<&AccountsCache>,
     ) -> Result<(Hash, u64), BankHashVerificationError> {
         if !use_index {
             let mut collect_time = Measure::start("collect");
             let (combined_maps, slots) = self.get_snapshot_storages(slot, Some(ancestors));
             collect_time.stop();
 
             let mut sort_time = Measure::start("sort_storages");
-            let storages = SortedStorages::new_with_slots(&combined_maps, &slots);
+            let storages = SortedStorages::new_with_slots(&combined_maps, &slots, Some(slot));
             sort_time.stop();
 
             let timings = HashStats {
@@ -4397,6 +4416,7 @@ impl AccountsDb {
                 Some(&self.thread_pool_clean),
                 timings,
                 check_hash,
+                accounts_cache,
             )
         } else {
             self.calculate_accounts_hash(slot, ancestors, check_hash)
@@ -4410,15 +4430,27 @@ impl AccountsDb {
         slot: Slot,
         ancestors: &Ancestors,
         expected_capitalization: Option<u64>,
+        use_write_cache: bool,
     ) -> (Hash, u64) {
         let check_hash = false;
+        let accounts_cache = if use_write_cache {
+            Some(&self.accounts_cache)
+        } else {
+            None
+        };
         let (hash, total_lamports) = self
-            .calculate_accounts_hash_helper(use_index, slot, ancestors, check_hash)
+            .calculate_accounts_hash_helper(use_index, slot, ancestors, check_hash, accounts_cache)
             .unwrap(); // unwrap here will never fail since check_hash = false
         if debug_verify {
             // calculate the other way (store or non-store) and verify results match.
             let (hash_other, total_lamports_other) = self
-                .calculate_accounts_hash_helper(!use_index, slot, ancestors, check_hash)
+                .calculate_accounts_hash_helper(
+                    !use_index,
+                    slot,
+                    ancestors,
+                    check_hash,
+                    accounts_cache,
+                )
                 .unwrap(); // unwrap here will never fail since check_hash = false
 
             let success = hash == hash_other
@@ -4432,12 +4464,13 @@ impl AccountsDb {
         (hash, total_lamports)
     }
 
-    fn scan_snapshot_stores(
+    fn scan_snapshot_stores_with_cache(
         storage: &SortedStorages,
         mut stats: &mut crate::accounts_hash::HashStats,
         bins: usize,
         bin_range: &Range<usize>,
         check_hash: bool,
+        accounts_cache: Option<&AccountsCache>,
     ) -> Result<Vec<Vec<Vec<CalculateHashIntermediate>>>, BankHashVerificationError> {
         let bin_calculator = PubkeyBinCalculator16::new(bins);
         assert!(bin_range.start < bins && bin_range.end <= bins && bin_range.start < bin_range.end);
@@ -4446,6 +4479,7 @@ impl AccountsDb {
         let mismatch_found = AtomicU64::new(0);
 
         let result: Vec<Vec<Vec<CalculateHashIntermediate>>> = Self::scan_account_storage_no_bank(
+            accounts_cache,
             storage,
             |loaded_account: LoadedAccount,
              accum: &mut Vec<Vec<CalculateHashIntermediate>>,
@@ -4513,6 +4547,7 @@ impl AccountsDb {
         thread_pool: Option<&ThreadPool>,
         mut stats: HashStats,
         check_hash: bool,
+        accounts_cache: Option<&AccountsCache>,
     ) -> Result<(Hash, u64), BankHashVerificationError> {
         let mut scan_and_hash = move || {
             // When calculating hashes, it is helpful to break the pubkeys found into bins based on the pubkey value.
@@ -4539,12 +4574,13 @@ impl AccountsDb {
                     end: (pass + 1) * bins_per_pass,
                 };
 
-                let result = Self::scan_snapshot_stores(
+                let result = Self::scan_snapshot_stores_with_cache(
                     &storages,
                     &mut stats,
                     PUBKEY_BINS_FOR_CALCULATING_HASHES,
                     &bounds,
                     check_hash,
+                    accounts_cache,
                 )?;
 
                 let (hash, lamports, for_next_pass) = AccountsHash::rest_of_hash_calculation(
@@ -4574,7 +4610,7 @@ impl AccountsDb {
         use BankHashVerificationError::*;
 
         let (calculated_hash, calculated_lamports) =
-            self.calculate_accounts_hash_helper(true, slot, ancestors, true)?;
+            self.calculate_accounts_hash_helper(true, slot, ancestors, true, None)?;
 
         if calculated_lamports != total_lamports {
             warn!(
@@ -5324,7 +5360,7 @@ impl AccountsDb {
     pub fn get_snapshot_storages(
         &self,
         snapshot_slot: Slot,
-        _ancestors: Option<&Ancestors>,
+        ancestors: Option<&Ancestors>,
     ) -> (SnapshotStorages, Vec<Slot>) {
         let mut m = Measure::start("get slots");
         let slots = self
@@ -5344,7 +5380,12 @@ impl AccountsDb {
                     slots
                         .iter()
                         .filter_map(|slot| {
-                            if *slot <= snapshot_slot && self.accounts_index.is_root(*slot) {
+                            if *slot <= snapshot_slot
+                                && (self.accounts_index.is_root(*slot)
+                                    || ancestors
+                                        .map(|ancestors| ancestors.contains_key(&slot))
+                                        .unwrap_or_default())
+                            {
                                 self.storage.0.get(&slot).map_or_else(
                                     || None,
                                     |item| {
@@ -5851,6 +5892,18 @@ pub mod tests {
         SortedStorages::new(&[])
     }
 
+    impl AccountsDb {
+        fn scan_snapshot_stores(
+            storage: &SortedStorages,
+            stats: &mut crate::accounts_hash::HashStats,
+            bins: usize,
+            bin_range: &Range<usize>,
+            check_hash: bool,
+        ) -> Result<Vec<Vec<Vec<CalculateHashIntermediate>>>, BankHashVerificationError> {
+            Self::scan_snapshot_stores_with_cache(storage, stats, bins, bin_range, check_hash, None)
+        }
+    }
+
     #[test]
     #[should_panic(
         expected = "bin_range.start < bins && bin_range.end <= bins &&\\n    bin_range.start < bin_range.end"
@@ -6186,6 +6239,7 @@ pub mod tests {
             None,
             HashStats::default(),
             false,
+            None,
         )
         .unwrap();
         let expected_hash = Hash::from_str("GKot5hBsd81kMupNCXHaqbhv3huEbxAFMLnpcX2hniwn").unwrap();
@@ -6207,6 +6261,7 @@ pub mod tests {
             None,
             HashStats::default(),
             false,
+            None,
         )
         .unwrap();
 
@@ -6254,6 +6309,7 @@ pub mod tests {
 
         let calls = AtomicU64::new(0);
         let result = AccountsDb::scan_account_storage_no_bank(
+            None,
             &get_storage_refs(&storages),
             |loaded_account: LoadedAccount, accum: &mut Vec<u64>, slot: Slot| {
                 calls.fetch_add(1, Ordering::Relaxed);
@@ -8171,10 +8227,10 @@ pub mod tests {
         db.add_root(some_slot);
         let check_hash = true;
         assert!(db
-            .calculate_accounts_hash_helper(false, some_slot, &ancestors, check_hash)
+            .calculate_accounts_hash_helper(false, some_slot, &ancestors, check_hash, None)
             .is_err());
         assert!(db
-            .calculate_accounts_hash_helper(true, some_slot, &ancestors, check_hash)
+            .calculate_accounts_hash_helper(true, some_slot, &ancestors, check_hash, None)
             .is_err());
     }
 
@@ -8194,9 +8250,9 @@ pub mod tests {
         db.add_root(some_slot);
         let check_hash = true;
         assert_eq!(
-            db.calculate_accounts_hash_helper(false, some_slot, &ancestors, check_hash)
+            db.calculate_accounts_hash_helper(false, some_slot, &ancestors, check_hash, None)
                 .unwrap(),
-            db.calculate_accounts_hash_helper(true, some_slot, &ancestors, check_hash)
+            db.calculate_accounts_hash_helper(true, some_slot, &ancestors, check_hash, None)
                 .unwrap(),
         );
     }

runtime/src/bank.rs

@@ -4563,7 +4563,9 @@ impl Bank {
     }
 
     pub fn calculate_capitalization(&self) -> u64 {
-        self.rc.accounts.calculate_capitalization(&self.ancestors)
+        self.rc
+            .accounts
+            .calculate_capitalization(&self.ancestors, self.slot())
     }
 
     pub fn calculate_and_verify_capitalization(&self) -> bool {
@@ -4612,6 +4614,7 @@ impl Bank {
                 self.slot(),
                 &self.ancestors,
                 Some(self.capitalization()),
+                false,
             );
         assert_eq!(total_lamports, self.capitalization());
         hash

runtime/src/snapshot_utils.rs

@@ -1006,6 +1006,7 @@ pub fn process_accounts_package_pre(
             thread_pool,
             crate::accounts_hash::HashStats::default(),
             false,
+            None,
         )
         .unwrap();
 

runtime/src/sorted_storages.rs

@@ -45,21 +45,25 @@ impl<'a> SortedStorages<'a> {
                 storage.slot() // this must be unique. Will be enforced in new_with_slots
             })
             .collect::<Vec<_>>();
-        Self::new_with_slots(source, &slots)
+        Self::new_with_slots(source, &slots, None)
     }
 
     // source[i] is in slot slots[i]
     // assumptions:
     // 1. slots vector contains unique slot #s.
     // 2. slots and source are the same len
-    pub fn new_with_slots(source: &'a [SnapshotStorage], slots: &[Slot]) -> Self {
+    pub fn new_with_slots(
+        source: &'a [SnapshotStorage],
+        slots: &[Slot],
+        slot: Option<Slot>,
+    ) -> Self {
         assert_eq!(
             source.len(),
             slots.len(),
             "source and slots are different lengths"
         );
-        let mut min = Slot::MAX;
-        let mut max = Slot::MIN;
+        let mut min = slot.unwrap_or(Slot::MAX);
+        let mut max = slot.map(|slot| slot + 1).unwrap_or(Slot::MIN);
         let slot_count = source.len();
         let mut time = Measure::start("get slot");
         slots.iter().for_each(|slot| {
@@ -132,18 +136,18 @@ pub mod tests {
     #[test]
     #[should_panic(expected = "slots are not unique")]
     fn test_sorted_storages_duplicate_slots() {
-        SortedStorages::new_with_slots(&[Vec::new(), Vec::new()], &[0, 0]);
+        SortedStorages::new_with_slots(&[Vec::new(), Vec::new()], &[0, 0], None);
     }
 
     #[test]
     #[should_panic(expected = "source and slots are different lengths")]
     fn test_sorted_storages_mismatched_lengths() {
-        SortedStorages::new_with_slots(&[Vec::new()], &[0, 0]);
+        SortedStorages::new_with_slots(&[Vec::new()], &[0, 0], None);
     }
 
     #[test]
     fn test_sorted_storages_none() {
-        let result = SortedStorages::new_with_slots(&[], &[]);
+        let result = SortedStorages::new_with_slots(&[], &[], None);
         assert_eq!(result.range, Range::default());
         assert_eq!(result.slot_count, 0);
         assert_eq!(result.storages.len(), 0);
@@ -156,7 +160,7 @@ pub mod tests {
         let vec_check = vec.clone();
         let slot = 4;
         let vecs = [vec];
-        let result = SortedStorages::new_with_slots(&vecs, &[slot]);
+        let result = SortedStorages::new_with_slots(&vecs, &[slot], None);
         assert_eq!(
             result.range,
             Range {
@@ -175,7 +179,7 @@ pub mod tests {
         let vec_check = vec.clone();
         let slots = [4, 7];
         let vecs = [vec.clone(), vec];
-        let result = SortedStorages::new_with_slots(&vecs, &slots);
+        let result = SortedStorages::new_with_slots(&vecs, &slots, None);
         assert_eq!(
             result.range,
             Range {