uses structural sharing for stake-delegations hash-map (#23585)

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:
https://github.com/solana-labs/solana/issues/23061#issuecomment-1063444072

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/

Cargo.lock

@@ -290,6 +290,15 @@ version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
+[[package]]
+name = "bitmaps"
+version = "2.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "031043d04099746d8db04daf1fa424b2bc8bd69d92b25962dcde24da39ab64a2"
+dependencies = [
+ "typenum",
+]
+
 [[package]]
 name = "bitvec"
 version = "0.19.5"
@@ -1927,6 +1936,22 @@ version = "0.2.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9007da9cacbd3e6343da136e98b0d2df013f553d35bdec8b518f07bea768e19c"
 
+[[package]]
+name = "im"
+version = "15.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "111c1983f3c5bb72732df25cddacee9b546d08325fb584b5ebd38148be7b0246"
+dependencies = [
+ "bitmaps",
+ "rand_core 0.5.1",
+ "rand_xoshiro",
+ "rayon",
+ "serde",
+ "sized-chunks",
+ "typenum",
+ "version_check",
+]
+
 [[package]]
 name = "index_list"
 version = "0.2.7"
@@ -3406,6 +3431,15 @@ dependencies = [
  "rand_core 0.6.3",
 ]
 
+[[package]]
+name = "rand_xoshiro"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a9fcdd2e881d02f1d9390ae47ad8e5696a9e4be7b547a1da2afbc61973217004"
+dependencies = [
+ "rand_core 0.5.1",
+]
+
 [[package]]
 name = "raptorq"
 version = "1.6.5"
@@ -4115,6 +4149,16 @@ version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2a30f10c911c0355f80f1c2faa8096efc4a58cdf8590b954d5b395efa071c711"
 
+[[package]]
+name = "sized-chunks"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "16d69225bde7a69b235da73377861095455d298f2b970996eec25ddbb42b3d1e"
+dependencies = [
+ "bitmaps",
+ "typenum",
+]
+
 [[package]]
 name = "slab"
 version = "0.4.5"
@@ -4810,6 +4854,7 @@ dependencies = [
  "bs58 0.4.0",
  "bv",
  "generic-array 0.14.5",
+ "im",
  "log",
  "memmap2 0.5.3",
  "rustc_version 0.4.0",
@@ -5594,6 +5639,7 @@ dependencies = [
  "ed25519-dalek",
  "flate2",
  "fnv",
+ "im",
  "index_list",
  "itertools 0.10.3",
  "lazy_static",

frozen-abi/Cargo.toml

@@ -21,6 +21,7 @@ thiserror = "1.0"
 
 [target.'cfg(not(target_arch = "bpf"))'.dependencies]
 generic-array = { version = "0.14.5", default-features = false, features = ["serde", "more_lengths"] }
+im = { version = "15.0.0", features = ["rayon", "serde"] }
 memmap2 = "0.5.3"
 
 [target.'cfg(not(target_arch = "bpf"))'.dev-dependencies]

frozen-abi/src/abi_example.rs

@@ -371,6 +371,21 @@ impl<
     }
 }
 
+#[cfg(not(target_arch = "bpf"))]
+impl<
+        T: Clone + std::cmp::Eq + std::hash::Hash + AbiExample,
+        S: Clone + AbiExample,
+        H: std::hash::BuildHasher + Default,
+    > AbiExample for im::HashMap<T, S, H>
+{
+    fn example() -> Self {
+        info!("AbiExample for (HashMap<T, S, H>): {}", type_name::<Self>());
+        let mut map = im::HashMap::default();
+        map.insert(T::example(), S::example());
+        map
+    }
+}
+
 impl<T: std::cmp::Ord + AbiExample, S: AbiExample> AbiExample for BTreeMap<T, S> {
     fn example() -> Self {
         info!("AbiExample for (BTreeMap<T, S>): {}", type_name::<Self>());

programs/bpf/Cargo.lock

@@ -189,6 +189,15 @@ version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
 
+[[package]]
+name = "bitmaps"
+version = "2.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "031043d04099746d8db04daf1fa424b2bc8bd69d92b25962dcde24da39ab64a2"
+dependencies = [
+ "typenum",
+]
+
 [[package]]
 name = "blake3"
 version = "0.3.8"
@@ -1334,6 +1343,22 @@ dependencies = [
  "unicode-normalization",
 ]
 
+[[package]]
+name = "im"
+version = "15.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "111c1983f3c5bb72732df25cddacee9b546d08325fb584b5ebd38148be7b0246"
+dependencies = [
+ "bitmaps",
+ "rand_core 0.5.1",
+ "rand_xoshiro",
+ "rayon",
+ "serde",
+ "sized-chunks",
+ "typenum",
+ "version_check",
+]
+
 [[package]]
 name = "index_list"
 version = "0.2.7"
@@ -2194,6 +2219,15 @@ dependencies = [
  "rand_core 0.5.1",
 ]
 
+[[package]]
+name = "rand_xoshiro"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a9fcdd2e881d02f1d9390ae47ad8e5696a9e4be7b547a1da2afbc61973217004"
+dependencies = [
+ "rand_core 0.5.1",
+]
+
 [[package]]
 name = "rayon"
 version = "1.5.1"
@@ -2663,6 +2697,16 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "65211b7b6fc3f14ff9fc7a2011a434e3e6880585bd2e9e9396315ae24cbf7852"
 
+[[package]]
+name = "sized-chunks"
+version = "0.6.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "16d69225bde7a69b235da73377861095455d298f2b970996eec25ddbb42b3d1e"
+dependencies = [
+ "bitmaps",
+ "typenum",
+]
+
 [[package]]
 name = "slab"
 version = "0.4.2"
@@ -3318,6 +3362,7 @@ dependencies = [
  "bs58 0.4.0",
  "bv",
  "generic-array 0.14.5",
+ "im",
  "log",
  "memmap2 0.5.3",
  "rustc_version 0.4.0",
@@ -3595,6 +3640,7 @@ dependencies = [
  "dir-diff",
  "flate2",
  "fnv",
+ "im",
  "index_list",
  "itertools 0.10.3",
  "lazy_static",

runtime/Cargo.toml

@@ -22,6 +22,7 @@ dashmap = { version = "4.0.2", features = ["rayon", "raw-api"] }
 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"

runtime/src/bank.rs

@@ -2545,10 +2545,10 @@ 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().collect();
         thread_pool.install(|| {
-            stakes
+            stake_delegations
-                .stake_delegations()
+                .into_par_iter()
-                .par_iter()
                 .for_each(|(stake_pubkey, delegation)| {
                     let vote_pubkey = &delegation.voter_pubkey;
                     if invalid_vote_keys.contains_key(vote_pubkey) {
@@ -6729,7 +6729,6 @@ pub(crate) mod tests {
                 genesis_sysvar_and_builtin_program_lamports, GenesisConfigInfo,
                 ValidatorVoteKeypairs,
             },
-            stake_delegations::StakeDelegations,
             status_cache::MAX_CACHE_ENTRIES,
         },
         crossbeam_channel::{bounded, unbounded},
@@ -6771,12 +6770,6 @@ pub(crate) mod tests {
         std::{result, thread::Builder, time::Duration},
     };
 
-    impl Bank {
-        fn cloned_stake_delegations(&self) -> StakeDelegations {
-            self.stakes_cache.stakes().stake_delegations().clone()
-        }
-    }
-
     fn new_sanitized_message(
         instructions: &[Instruction],
         payer: Option<&Pubkey>,
@@ -10896,7 +10889,7 @@ pub(crate) mod tests {
         } = create_genesis_config_with_leader(500, &solana_sdk::pubkey::new_rand(), 1);
         let bank = Arc::new(Bank::new_for_tests(&genesis_config));
 
-        let stake_delegations = bank.cloned_stake_delegations();
+        let stake_delegations = bank.stakes_cache.stakes().stake_delegations().clone();
         assert_eq!(stake_delegations.len(), 1); // bootstrap validator has
                                                 // to have a stake delegation
 
@@ -10932,7 +10925,7 @@ pub(crate) mod tests {
 
         bank.process_transaction(&transaction).unwrap();
 
-        let stake_delegations = bank.cloned_stake_delegations();
+        let stake_delegations = bank.stakes_cache.stakes().stake_delegations().clone();
         assert_eq!(stake_delegations.len(), 2);
         assert!(stake_delegations.get(&stake_keypair.pubkey()).is_some());
     }

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;

runtime/src/serde_snapshot/tests.rs

@@ -305,7 +305,7 @@ mod test_bank_serialize {
 
     // This some what long test harness is required to freeze the ABI of
     // Bank's serialization due to versioned nature
-    #[frozen_abi(digest = "HVyzePMkma8T54PymRW32FAgDXpSdom59K6RnPsCNJjj")]
+    #[frozen_abi(digest = "4TqK4bCbL76s9mf1gbcRVSQUgtSe68wDCk3jTwoQzr6R")]
     #[derive(Serialize, AbiExample)]
     pub struct BankAbiTestWrapperNewer {
         #[serde(serialize_with = "wrapper_newer")]

runtime/src/stake_delegations.rs

@@ -1,128 +0,0 @@
-//! Map pubkeys to stake delegations
-//!
-//! This module implements clone-on-write semantics for `StakeDelegations` to reduce unnecessary
-//! cloning of the underlying map.
-use {
-    solana_sdk::{pubkey::Pubkey, stake::state::Delegation},
-    std::{
-        collections::HashMap,
-        ops::{Deref, DerefMut},
-        sync::Arc,
-    },
-};
-
-/// A map of pubkey-to-stake-delegation with clone-on-write semantics
-#[derive(Default, Clone, PartialEq, Debug, Deserialize, Serialize, AbiExample)]
-pub struct StakeDelegations(Arc<StakeDelegationsInner>);
-
-impl Deref for StakeDelegations {
-    type Target = StakeDelegationsInner;
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
-}
-
-impl DerefMut for StakeDelegations {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        Arc::make_mut(&mut self.0)
-    }
-}
-
-/// The inner type, which maps pubkeys to stake delegations
-type StakeDelegationsInner = HashMap<Pubkey, Delegation>;
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    /// Ensure that StakeDelegations is indeed clone-on-write
-    #[test]
-    fn test_stake_delegations_is_cow() {
-        let voter_pubkey = Pubkey::new_unique();
-        let stake = rand::random();
-        let activation_epoch = rand::random();
-        let warmup_cooldown_rate = rand::random();
-        let delegation =
-            Delegation::new(&voter_pubkey, stake, activation_epoch, warmup_cooldown_rate);
-
-        let pubkey = Pubkey::new_unique();
-
-        let mut stake_delegations = StakeDelegations::default();
-        stake_delegations.insert(pubkey, delegation);
-
-        // Test: Clone the stake delegations and **do not modify**.  Assert the underlying maps are
-        // the same instance.
-        {
-            let stake_delegations2 = stake_delegations.clone();
-            assert_eq!(stake_delegations, stake_delegations2);
-            assert!(
-                Arc::ptr_eq(&stake_delegations.0, &stake_delegations2.0),
-                "Inner Arc must point to the same HashMap"
-            );
-            assert!(
-                std::ptr::eq(stake_delegations.deref(), stake_delegations2.deref()),
-                "Deref must point to the same HashMap"
-            );
-        }
-
-        // Test: Clone the stake delegations and then modify (remove the K-V, then re-add the same
-        // one, so the stake delegations are still logically equal).  Assert the underlying maps
-        // are unique instances.
-        {
-            let mut stake_delegations2 = stake_delegations.clone();
-            stake_delegations2.clear();
-            assert_ne!(stake_delegations, stake_delegations2);
-            stake_delegations2.insert(pubkey, delegation);
-            assert_eq!(stake_delegations, stake_delegations2);
-            assert!(
-                !Arc::ptr_eq(&stake_delegations.0, &stake_delegations2.0),
-                "Inner Arc must point to different HashMaps"
-            );
-            assert!(
-                !std::ptr::eq(stake_delegations.deref(), stake_delegations2.deref()),
-                "Deref must point to different HashMaps"
-            );
-        }
-    }
-
-    /// Ensure that StakeDelegations serializes and deserializes between the inner and outer types
-    #[test]
-    fn test_stake_delegations_serde() {
-        let voter_pubkey = Pubkey::new_unique();
-        let stake = rand::random();
-        let activation_epoch = rand::random();
-        let warmup_cooldown_rate = rand::random();
-        let delegation =
-            Delegation::new(&voter_pubkey, stake, activation_epoch, warmup_cooldown_rate);
-
-        let pubkey = Pubkey::new_unique();
-
-        let mut stake_delegations_outer = StakeDelegations::default();
-        stake_delegations_outer.insert(pubkey, delegation);
-
-        let mut stake_delegations_inner = StakeDelegationsInner::default();
-        stake_delegations_inner.insert(pubkey, delegation);
-
-        // Test: Assert that serializing the outer and inner types produces the same data
-        assert_eq!(
-            bincode::serialize(&stake_delegations_outer).unwrap(),
-            bincode::serialize(&stake_delegations_inner).unwrap(),
-        );
-
-        // Test: Assert that serializing the outer type then deserializing to the inner type
-        // produces the same values
-        {
-            let data = bincode::serialize(&stake_delegations_outer).unwrap();
-            let deserialized_inner: StakeDelegationsInner = bincode::deserialize(&data).unwrap();
-            assert_eq!(&deserialized_inner, stake_delegations_outer.deref());
-        }
-
-        // Test: Assert that serializing the inner type then deserializing to the outer type
-        // produces the same values
-        {
-            let data = bincode::serialize(&stake_delegations_inner).unwrap();
-            let deserialized_outer: StakeDelegations = bincode::deserialize(&data).unwrap();
-            assert_eq!(deserialized_outer.deref(), &stake_delegations_inner);
-        }
-    }
-}

runtime/src/stakes.rs

@@ -2,11 +2,11 @@
 //! node stakes
 use {
     crate::{
-        stake_delegations::StakeDelegations,
         stake_history::StakeHistory,
         vote_account::{VoteAccount, VoteAccounts, VoteAccountsHashMap},
     },
     dashmap::DashMap,
+    im::HashMap as ImHashMap,
     num_derive::ToPrimitive,
     num_traits::ToPrimitive,
     rayon::{
@@ -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(crate) fn stake_delegations(&self) -> &ImHashMap<Pubkey, Delegation> {
         &self.stake_delegations
     }