Stores capitalization in account hashes map (#30635)

runtime/src/accounts_db.rs

@@ -1417,7 +1417,7 @@ pub struct AccountsDb {
     pub thread_pool_clean: ThreadPool,
 
     accounts_delta_hashes: Mutex<HashMap<Slot, AccountsDeltaHash>>,
-    accounts_hashes: Mutex<HashMap<Slot, AccountsHash>>,
+    accounts_hashes: Mutex<HashMap<Slot, (AccountsHash, /*capitalization*/ u64)>>,
     bank_hash_stats: Mutex<HashMap<Slot, BankHashStats>>,
 
     pub stats: AccountsStats,
@@ -7346,14 +7346,18 @@ impl AccountsDb {
                 expected_capitalization,
             )
             .unwrap(); // unwrap here will never fail since check_hash = false
-        self.set_accounts_hash(slot, accounts_hash);
+        self.set_accounts_hash(slot, (accounts_hash, total_lamports));
         (accounts_hash, total_lamports)
     }
 
     /// Set the accounts hash for `slot` in the `accounts_hashes` map
     ///
     /// returns the previous accounts hash for `slot`
-    fn set_accounts_hash(&self, slot: Slot, accounts_hash: AccountsHash) -> Option<AccountsHash> {
+    fn set_accounts_hash(
+        &self,
+        slot: Slot,
+        accounts_hash: (AccountsHash, /*capitalization*/ u64),
+    ) -> Option<(AccountsHash, /*capitalization*/ u64)> {
         self.accounts_hashes
             .lock()
             .unwrap()
@@ -7365,12 +7369,13 @@ impl AccountsDb {
         &mut self,
         slot: Slot,
         accounts_hash: SerdeAccountsHash,
-    ) -> Option<AccountsHash> {
+        capitalization: u64,
-        self.set_accounts_hash(slot, accounts_hash.into())
+    ) -> Option<(AccountsHash, /*capitalization*/ u64)> {
+        self.set_accounts_hash(slot, (accounts_hash.into(), capitalization))
     }
 
     /// Get the accounts hash for `slot` in the `accounts_hashes` map
-    pub fn get_accounts_hash(&self, slot: Slot) -> Option<AccountsHash> {
+    pub fn get_accounts_hash(&self, slot: Slot) -> Option<(AccountsHash, /*capitalization*/ u64)> {
         self.accounts_hashes.lock().unwrap().get(&slot).cloned()
     }
 
@@ -7674,7 +7679,7 @@ impl AccountsDb {
 
         if config.ignore_mismatch {
             Ok(())
-        } else if let Some(found_accounts_hash) = self.get_accounts_hash(slot) {
+        } else if let Some((found_accounts_hash, _)) = self.get_accounts_hash(slot) {
             if calculated_accounts_hash == found_accounts_hash {
                 Ok(())
             } else {
@@ -9535,7 +9540,7 @@ pub mod tests {
 
         // used by serde_snapshot tests
         pub fn set_accounts_hash_for_tests(&self, slot: Slot, accounts_hash: AccountsHash) {
-            self.set_accounts_hash(slot, accounts_hash);
+            self.set_accounts_hash(slot, (accounts_hash, u64::default()));
         }
 
         // used by serde_snapshot tests
@@ -11878,8 +11883,8 @@ pub mod tests {
                 accounts.get_accounts_delta_hash(latest_slot).unwrap(),
             );
             assert_eq!(
-                daccounts.get_accounts_hash(latest_slot).unwrap(),
+                daccounts.get_accounts_hash(latest_slot).unwrap().0,
-                accounts.get_accounts_hash(latest_slot).unwrap(),
+                accounts.get_accounts_hash(latest_slot).unwrap().0,
             );
 
             daccounts.print_count_and_status("daccounts");
@@ -12610,7 +12615,8 @@ pub mod tests {
 
         db.store_for_tests(some_slot, &[(&key, &account)]);
         db.add_root_and_flush_write_cache(some_slot);
-        db.update_accounts_hash_for_tests(some_slot, &ancestors, true, true);
+        let (_, capitalization) =
+            db.update_accounts_hash_for_tests(some_slot, &ancestors, true, true);
 
         let config = VerifyAccountsHashAndLamportsConfig::new_for_test(
             &ancestors,
@@ -12630,7 +12636,10 @@ pub mod tests {
             Err(MissingAccountsHash)
         );
 
-        db.set_accounts_hash(some_slot, AccountsHash(Hash::new(&[0xca; HASH_BYTES])));
+        db.set_accounts_hash(
+            some_slot,
+            (AccountsHash(Hash::new(&[0xca; HASH_BYTES])), capitalization),
+        );
 
         assert_matches!(
             db.verify_accounts_hash_and_lamports(some_slot, 1, config),

runtime/src/bank.rs

@@ -7190,7 +7190,11 @@ impl Bank {
     }
 
     pub fn get_accounts_hash(&self) -> Option<AccountsHash> {
-        self.rc.accounts.accounts_db.get_accounts_hash(self.slot())
+        self.rc
+            .accounts
+            .accounts_db
+            .get_accounts_hash(self.slot())
+            .map(|(accounts_hash, _)| accounts_hash)
     }
 
     pub fn get_snapshot_hash(&self) -> SnapshotHash {

runtime/src/serde_snapshot.rs

@@ -550,6 +550,7 @@ where
         accounts_update_notifier,
         exit,
         bank_fields.epoch_accounts_hash,
+        bank_fields.capitalization,
     )?;
 
     let bank_rc = BankRc::new(Accounts::new_empty(accounts_db), bank_fields.slot);
@@ -672,6 +673,7 @@ fn reconstruct_accountsdb_from_fields<E>(
     accounts_update_notifier: Option<AccountsUpdateNotifier>,
     exit: &Arc<AtomicBool>,
     epoch_accounts_hash: Option<Hash>,
+    capitalization: u64,
 ) -> Result<(AccountsDb, ReconstructedAccountsDbInfo), Error>
 where
     E: SerializableStorage + std::marker::Sync,
@@ -739,8 +741,11 @@ where
         old_accounts_delta_hash.is_none(),
         "There should not already be an AccountsDeltaHash at slot {snapshot_slot}: {old_accounts_delta_hash:?}",
         );
-    let old_accounts_hash = accounts_db
+    let old_accounts_hash = accounts_db.set_accounts_hash_from_snapshot(
-        .set_accounts_hash_from_snapshot(snapshot_slot, snapshot_bank_hash_info.accounts_hash);
+        snapshot_slot,
+        snapshot_bank_hash_info.accounts_hash,
+        capitalization,
+    );
     assert!(
         old_accounts_hash.is_none(),
         "There should not already be an AccountsHash at slot {snapshot_slot}: {old_accounts_hash:?}",

runtime/src/serde_snapshot/newer.rs

@@ -282,7 +282,7 @@ impl<'a> TypeContext<'a> for Context {
         let accounts_hash = serializable_db
             .accounts_db
             .get_accounts_hash(slot)
-            .map(Into::into)
+            .map(|(accounts_hash, _)| accounts_hash.into())
             .unwrap_or_default();
         let stats = serializable_db
             .accounts_db

runtime/src/serde_snapshot/tests.rs

@@ -124,6 +124,7 @@ where
         None,
         &Arc::default(),
         None,
+        u64::default(),
     )
     .map(|(accounts_db, _)| accounts_db)
 }
@@ -221,8 +222,8 @@ fn test_accounts_serialize_style(serde_style: SerdeStyle) {
     check_accounts(&daccounts, &pubkeys, 100);
     let daccounts_delta_hash = daccounts.accounts_db.calculate_accounts_delta_hash(slot);
     assert_eq!(accounts_delta_hash, daccounts_delta_hash);
-    let daccounts_hash = daccounts.accounts_db.get_accounts_hash(slot);
+    let daccounts_hash = daccounts.accounts_db.get_accounts_hash(slot).unwrap().0;
-    assert_eq!(Some(accounts_hash), daccounts_hash);
+    assert_eq!(accounts_hash, daccounts_hash);
 }
 
 fn test_bank_serialize_style(