Make StakeDelegations clone-on-write (#21542)

runtime/src/bank.rs

@@ -125,7 +125,7 @@ use solana_sdk::{
     },
 };
 use solana_stake_program::stake_state::{
-    self, Delegation, InflationPointCalculationEvent, PointValue, StakeState,
+    self, InflationPointCalculationEvent, PointValue, StakeState,
 };
 use solana_vote_program::{
     vote_instruction::VoteInstruction,
@@ -5784,11 +5784,6 @@ impl Bank {
         }
     }
 
-    /// current stake delegations for this bank
-    pub fn cloned_stake_delegations(&self) -> HashMap<Pubkey, Delegation> {
-        self.stakes.read().unwrap().stake_delegations().clone()
-    }
-
     pub fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
         self.stakes.read().unwrap().staked_nodes()
     }
@@ -6494,6 +6489,7 @@ pub(crate) mod tests {
             create_genesis_config_with_leader, create_genesis_config_with_vote_accounts,
             GenesisConfigInfo, ValidatorVoteKeypairs,
         },
+        stake_delegations::StakeDelegations,
         status_cache::MAX_CACHE_ENTRIES,
     };
     use crossbeam_channel::{bounded, unbounded};
@@ -6533,6 +6529,12 @@ pub(crate) mod tests {
     };
     use std::{result, thread::Builder, time::Duration};
 
+    impl Bank {
+        fn cloned_stake_delegations(&self) -> StakeDelegations {
+            self.stakes.read().unwrap().stake_delegations().clone()
+        }
+    }
+
     fn new_sanitized_message(
         instructions: &[Instruction],
         payer: Option<&Pubkey>,

runtime/src/lib.rs

@@ -48,6 +48,7 @@ pub mod snapshot_hash;
 pub mod snapshot_package;
 pub mod snapshot_utils;
 pub mod sorted_storages;
+pub mod stake_delegations;
 pub mod stake_weighted_timestamp;
 pub mod stakes;
 pub mod status_cache;

runtime/src/serde_snapshot/tests.rs

@@ -312,7 +312,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 = "5vYNtKztrNKfb6TtRktXdLCbU73rJSWhLwguCHLvFujZ")]
+    #[frozen_abi(digest = "FBhQnLvFHaCkXN8ZL8MaR6yQUhLgyvWqdoFeBRwsmbBo")]
     #[derive(Serialize, AbiExample)]
     pub struct BankAbiTestWrapperFuture {
         #[serde(serialize_with = "wrapper_future")]

runtime/src/stake_delegations.rs

@@ -0,0 +1,86 @@
+//! 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::*;
+
+    #[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"
+            );
+        }
+    }
+}

runtime/src/stakes.rs

@@ -1,7 +1,10 @@
 //! Stakes serve as a cache of stake and vote accounts to derive
 //! node stakes
 use {
-    crate::vote_account::{VoteAccount, VoteAccounts, VoteAccountsHashMap},
+    crate::{
+        stake_delegations::StakeDelegations,
+        vote_account::{VoteAccount, VoteAccounts, VoteAccountsHashMap},
+    },
     rayon::{
         iter::{IntoParallelRefIterator, ParallelIterator},
         ThreadPool,
@@ -27,7 +30,7 @@ pub struct Stakes {
     vote_accounts: VoteAccounts,
 
     /// stake_delegations
-    stake_delegations: HashMap<Pubkey, Delegation>,
+    stake_delegations: StakeDelegations,
 
     /// unused
     unused: u64,
@@ -271,7 +274,7 @@ impl Stakes {
         &self.vote_accounts
     }
 
-    pub fn stake_delegations(&self) -> &HashMap<Pubkey, Delegation> {
+    pub fn stake_delegations(&self) -> &StakeDelegations {
         &self.stake_delegations
     }