uses structural sharing for stake-delegations hash-map

StakeDelegations is using Arc to implement copy-on-write semantics: https://github.com/solana-labs/solana/blob/58c0db970/runtime/src/stake_delegations.rs#L14-L16 However a single delegation change will still clone the entire hash-map, resulting in excessive memory use as observed in: solana-labs#23061 (comment) This commit instead uses immutable hash-map implementing structural sharing: > which means that if two data structures are mostly copies of each > other, most of the memory they take up will be shared between them. https://docs.rs/im/latest/im/
behzadnouri · Mar 10, 2022 · c058e2f · c058e2f
1 parent 58c0db9
commit c058e2f
Show file tree

Hide file tree

Showing 7 changed files with 102 additions and 136 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/programs/bpf/Cargo.lock b/programs/bpf/Cargo.lock
diff --git a/runtime/Cargo.toml b/runtime/Cargo.toml
@@ -22,6 +22,7 @@ crossbeam-channel = "0.5"
 dir-diff = "0.3.2"
 flate2 = "1.0.22"
 fnv = "1.0.7"
+im = { version = "15.0.0", features = ["rayon", "serde"] }
 index_list = "0.2.7"
 itertools = "0.10.3"
 lazy_static = "1.4.0"

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
@@ -2525,9 +2525,13 @@ impl Bank {
         let invalid_stake_keys: DashMap<Pubkey, InvalidCacheEntryReason> = DashMap::new();
         let invalid_vote_keys: DashMap<Pubkey, InvalidCacheEntryReason> = DashMap::new();
 
+        let stake_delegations: Vec<_> = stakes
+            .stake_delegations()
+            .iter()
+            .map(|(&pubkey, &delegation)| (pubkey, delegation))
+            .collect();
         thread_pool.install(|| {
-            stakes
-                .stake_delegations()
+            stake_delegations
                 .par_iter()
                 .for_each(|(stake_pubkey, delegation)| {
                     let vote_pubkey = &delegation.voter_pubkey;
@@ -6712,10 +6716,10 @@ pub(crate) mod tests {
                 genesis_sysvar_and_builtin_program_lamports, GenesisConfigInfo,
                 ValidatorVoteKeypairs,
             },
-            stake_delegations::StakeDelegations,
             status_cache::MAX_CACHE_ENTRIES,
         },
         crossbeam_channel::{bounded, unbounded},
+        im::HashMap as ImHashMap,
         solana_program_runtime::invoke_context::InvokeContext,
         solana_sdk::{
             account::Account,
@@ -6755,7 +6759,7 @@ pub(crate) mod tests {
     };
 
     impl Bank {
-        fn cloned_stake_delegations(&self) -> StakeDelegations {
+        fn cloned_stake_delegations(&self) -> ImHashMap<Pubkey, Delegation> {
             self.stakes_cache.stakes().stake_delegations().clone()
         }
     }

diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs
@@ -52,7 +52,6 @@ pub mod snapshot_hash;
 pub mod snapshot_package;
 pub mod snapshot_utils;
 pub mod sorted_storages;
-pub mod stake_delegations;
 pub mod stake_history;
 pub mod stake_weighted_timestamp;
 pub mod stakes;

diff --git a/runtime/src/stake_delegations.rs b/runtime/src/stake_delegations.rs
diff --git a/runtime/src/stakes.rs b/runtime/src/stakes.rs
@@ -2,7 +2,6 @@
 //! node stakes
 use {
     crate::{
-        stake_delegations::StakeDelegations,
         stake_history::StakeHistory,
         vote_account::{VoteAccount, VoteAccounts, VoteAccountsHashMap},
     },
@@ -28,6 +27,7 @@ use {
         collections::HashMap,
         sync::{Arc, RwLock, RwLockReadGuard},
     },
+    im::HashMap as ImHashMap,
 };
 
 #[derive(Debug, Clone, PartialEq, ToPrimitive)]
@@ -151,7 +151,7 @@ pub struct Stakes {
     vote_accounts: VoteAccounts,
 
     /// stake_delegations
-    stake_delegations: StakeDelegations,
+    stake_delegations: ImHashMap<Pubkey, Delegation>,
 
     /// unused
     unused: u64,
@@ -337,7 +337,7 @@ impl Stakes {
         &self.vote_accounts
     }
 
-    pub fn stake_delegations(&self) -> &StakeDelegations {
+    pub fn stake_delegations(&self) -> &ImHashMap<Pubkey, Delegation> {
         &self.stake_delegations
     }