From a14e24a230e75453cad48a4891ad2ea2d73481d9 Mon Sep 17 00:00:00 2001
From: HaoranYi <haoran.yi@solana.com>
Date: Thu, 7 Mar 2024 18:18:29 +0000
Subject: [PATCH] refactor compute delete hash output and fix a test

---
 accounts-db/src/accounts_db.rs      | 34 +++++++++++++++++++++--------
 runtime/src/bank.rs                 |  8 +++++--
 runtime/src/serde_snapshot/tests.rs | 13 ++++++++++-
 3 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/accounts-db/src/accounts_db.rs b/accounts-db/src/accounts_db.rs
index bc5d090c076c99..af37f6fbe1f27a 100644
--- a/accounts-db/src/accounts_db.rs
+++ b/accounts-db/src/accounts_db.rs
@@ -2386,6 +2386,17 @@ pub struct PubkeyHashAccount {
 
 type LTHashCacheMap = RwLock<HashMap<Pubkey, (Option<AccountSharedData>, Option<AccountLTHash>)>>;
 
+pub struct AccountsDeltaHashCalculationOutput {
+    /// `delta_hash` computed from all the accounts that have been written to.
+    pub delta_hash: AccountsDeltaHash,
+
+    /// Contains the accounts that have been written to.
+    ///
+    /// Note that `accounts` collected in the vec may be in arbitrary order as a
+    /// result of parallel scanning in accounts_db.
+    pub accounts: Vec<(Pubkey, AccountHash)>,
+}
+
 impl AccountsDb {
     pub const DEFAULT_ACCOUNTS_HASH_CACHE_DIR: &'static str = "accounts_hash_cache";
 
@@ -8075,14 +8086,16 @@ impl AccountsDb {
         }
     }
 
-    /// Wrapper function to calculate accounts delta hash for `slot` (only used for testing and benchmarking.)
+    /// Wrapper function to calculate accounts delta hash for `slot` (only used
+    /// for testing and benchmarking.)
     ///
-    /// As part of calculating the accounts delta hash, get a list of accounts modified this slot
-    /// (aka dirty pubkeys) and add them to `self.uncleaned_pubkeys` for future cleaning.
-    pub fn calculate_accounts_delta_hash(
-        &self,
-        slot: Slot,
-    ) -> (AccountsDeltaHash, Vec<(Pubkey, AccountHash)>) {
+    /// As part of calculating the accounts delta hash, get a list of accounts
+    /// modified this slot (aka dirty pubkeys) and add them to
+    /// `self.uncleaned_pubkeys` for future cleaning.
+    ///
+
+    #[cfg(feature = "dev-context-only-utils")]
+    pub fn calculate_accounts_delta_hash(&self, slot: Slot) -> AccountsDeltaHashCalculationOutput {
         self.calculate_accounts_delta_hash_internal(slot, None, HashMap::default())
     }
 
@@ -8144,7 +8157,7 @@ impl AccountsDb {
         slot: Slot,
         ignore: Option<Pubkey>,
         mut skipped_rewrites: HashMap<Pubkey, AccountHash>,
-    ) -> (AccountsDeltaHash, Vec<(Pubkey, AccountHash)>) {
+    ) -> AccountsDeltaHashCalculationOutput {
         let (mut hashes, scan_us, mut accumulate) = self.get_pubkey_hash_for_slot(slot);
         let original_pubkey_hash = hashes.clone();
         let dirty_keys = hashes.iter().map(|(pubkey, _hash)| *pubkey).collect();
@@ -8185,7 +8198,10 @@ impl AccountsDb {
             .skipped_rewrites_num
             .fetch_add(num_skipped_rewrites, Ordering::Relaxed);
 
-        (accounts_delta_hash, original_pubkey_hash)
+        AccountsDeltaHashCalculationOutput {
+            delta_hash: accounts_delta_hash,
+            accounts: original_pubkey_hash,
+        }
     }
 
     /// Set the accounts delta hash for `slot` in the `accounts_delta_hashes` map
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index c18b77893ab547..8cbb1dfb9498f5 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -75,7 +75,8 @@ use {
         accounts::{AccountAddressFilter, Accounts, PubkeyAccountSlot},
         accounts_db::{
             AccountShrinkThreshold, AccountStorageEntry, AccountsDb, AccountsDbConfig,
-            CalcAccountsHashDataSource, VerifyAccountsHashAndLamportsConfig,
+            AccountsDeltaHashCalculationOutput, CalcAccountsHashDataSource,
+            VerifyAccountsHashAndLamportsConfig,
         },
         accounts_hash::{
             AccountHash, AccountLTHash, AccountsHash, CalcAccountsHashConfig, HashStats,
@@ -6333,7 +6334,10 @@ impl Bank {
         .then_some(sysvar::epoch_rewards::id());
 
         let mut delta_hash_timer = Measure::start("delta_hash_compute");
-        let (accounts_delta_hash, pubkey_hash) = self
+        let AccountsDeltaHashCalculationOutput {
+            delta_hash: accounts_delta_hash,
+            accounts: pubkey_hash,
+        } = self
             .rc
             .accounts
             .accounts_db
diff --git a/runtime/src/serde_snapshot/tests.rs b/runtime/src/serde_snapshot/tests.rs
index 510069c92662fc..5556d8bc47a619 100644
--- a/runtime/src/serde_snapshot/tests.rs
+++ b/runtime/src/serde_snapshot/tests.rs
@@ -267,7 +267,18 @@ mod serde_snapshot_tests {
         ));
         check_accounts_local(&daccounts, &pubkeys, 100);
         let daccounts_delta_hash = daccounts.accounts_db.calculate_accounts_delta_hash(slot);
-        assert_eq!(accounts_delta_hash, daccounts_delta_hash);
+        assert_eq!(
+            accounts_delta_hash.delta_hash,
+            daccounts_delta_hash.delta_hash
+        );
+
+        // Convert the accounts into HashMap before compare since the accounts
+        // may be of different order due to parallel scanning.
+        let h1: std::collections::HashMap<_, _> =
+            accounts_delta_hash.accounts.into_iter().collect();
+        let h2: std::collections::HashMap<_, _> =
+            daccounts_delta_hash.accounts.into_iter().collect();
+        assert_eq!(h1, h2);
         let daccounts_hash = daccounts.accounts_db.get_accounts_hash(slot).unwrap().0;
         assert_eq!(accounts_hash, daccounts_hash);
     }