implements copy-on-write for staked-nodes (#19090)

Bank::staked_nodes and Bank::epoch_staked_nodes redundantly clone
staked-nodes HashMap even though an immutable reference will suffice:
https://github.com/solana-labs/solana/blob/a9014cece/runtime/src/vote_account.rs#L77

This commit implements copy-on-write semantics for staked-nodes by
wrapping the underlying HashMap in Arc<...>.

(cherry picked from commit f302774)

# Conflicts:
#	runtime/src/bank.rs
#	runtime/src/stakes.rs
#	runtime/src/vote_account.rs

gossip/src/cluster_info.rs

@@ -1749,7 +1749,7 @@ impl ClusterInfo {
                                 Some(root_bank.feature_set.clone()),
                             )
                         }
-                        None => (HashMap::new(), None),
+                        None => (Arc::default(), None),
                     };
                     let require_stake_for_gossip =
                         self.require_stake_for_gossip(feature_set.as_deref(), &stakes);
@@ -2542,7 +2542,7 @@ impl ClusterInfo {
         // feature does not roll back (if the feature happens to get enabled in
         // a minority fork).
         let (feature_set, stakes) = match bank_forks {
-            None => (None, HashMap::default()),
+            None => (None, Arc::default()),
             Some(bank_forks) => {
                 let bank = bank_forks.read().unwrap().root_bank();
                 let feature_set = bank.feature_set.clone();

ledger/src/leader_schedule_utils.rs

@@ -13,7 +13,10 @@ pub fn leader_schedule(epoch: Epoch, bank: &Bank) -> Option<LeaderSchedule> {
     bank.epoch_staked_nodes(epoch).map(|stakes| {
         let mut seed = [0u8; 32];
         seed[0..8].copy_from_slice(&epoch.to_le_bytes());
-        let mut stakes: Vec<_> = stakes.into_iter().collect();
+        let mut stakes: Vec<_> = stakes
+            .iter()
+            .map(|(pubkey, stake)| (*pubkey, *stake))
+            .collect();
         sort_stakes(&mut stakes);
         LeaderSchedule::new(
             &stakes,
@@ -91,7 +94,11 @@ mod tests {
                 .genesis_config;
         let bank = Bank::new(&genesis_config);
 
-        let pubkeys_and_stakes: Vec<_> = bank.staked_nodes().into_iter().collect();
+        let pubkeys_and_stakes: Vec<_> = bank
+            .staked_nodes()
+            .iter()
+            .map(|(pubkey, stake)| (*pubkey, *stake))
+            .collect();
         let seed = [0u8; 32];
         let leader_schedule = LeaderSchedule::new(
             &pubkeys_and_stakes,

runtime/src/bank.rs

@@ -5282,8 +5282,13 @@ impl Bank {
         self.stakes_cache.stakes().stake_delegations().clone()
     }
 
+<<<<<<< HEAD
     pub fn staked_nodes(&self) -> HashMap<Pubkey, u64> {
         self.stakes_cache.stakes().staked_nodes()
+=======
+    pub fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
+        self.stakes.read().unwrap().staked_nodes()
+>>>>>>> f302774cf (implements copy-on-write for staked-nodes (#19090))
     }
 
     /// current vote accounts for this bank along with the stake
@@ -5312,7 +5317,7 @@ impl Bank {
         &self.epoch_stakes
     }
 
-    pub fn epoch_staked_nodes(&self, epoch: Epoch) -> Option<HashMap<Pubkey, u64>> {
+    pub fn epoch_staked_nodes(&self, epoch: Epoch) -> Option<Arc<HashMap<Pubkey, u64>>> {
         Some(self.epoch_stakes.get(&epoch)?.stakes().staked_nodes())
     }
 

runtime/src/stakes.rs

@@ -1,6 +1,7 @@
 //! Stakes serve as a cache of stake and vote accounts to derive
 //! node stakes
 use {
+<<<<<<< HEAD
     crate::vote_account::{VoteAccount, VoteAccounts, VoteAccountsHashMap},
     dashmap::DashMap,
     num_derive::ToPrimitive,
@@ -113,6 +114,23 @@ impl StakesCache {
         }
     }
 }
+=======
+    crate::vote_account::{ArcVoteAccount, VoteAccounts},
+    solana_sdk::{
+        account::{AccountSharedData, ReadableAccount},
+        clock::Epoch,
+        pubkey::Pubkey,
+        stake::{
+            self,
+            state::{Delegation, StakeState},
+        },
+        sysvar::stake_history::StakeHistory,
+    },
+    solana_stake_program::stake_state,
+    solana_vote_program::vote_state::VoteState,
+    std::{borrow::Borrow, collections::HashMap, sync::Arc},
+};
+>>>>>>> f302774cf (implements copy-on-write for staked-nodes (#19090))
 
 #[derive(Default, Clone, PartialEq, Debug, Deserialize, Serialize, AbiExample)]
 pub struct Stakes {
@@ -335,7 +353,7 @@ impl Stakes {
         &self.stake_delegations
     }
 
-    pub fn staked_nodes(&self) -> HashMap<Pubkey, u64> {
+    pub fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
         self.vote_accounts.staked_nodes()
     }
 

runtime/src/vote_account.rs

@@ -1,4 +1,5 @@
 use {
+<<<<<<< HEAD
     serde::{
         de::{Deserialize, Deserializer},
         ser::{Serialize, Serializer},
@@ -13,6 +14,21 @@ use {
         cmp::Ordering,
         collections::{hash_map::Entry, HashMap},
         iter::FromIterator,
+=======
+    itertools::Itertools,
+    serde::de::{Deserialize, Deserializer},
+    serde::ser::{Serialize, Serializer},
+    solana_sdk::{
+        account::Account, account::AccountSharedData, instruction::InstructionError, pubkey::Pubkey,
+    },
+    solana_vote_program::vote_state::VoteState,
+    std::{
+        borrow::Borrow,
+        cmp::Ordering,
+        collections::{hash_map::Entry, HashMap},
+        iter::FromIterator,
+        ops::Deref,
+>>>>>>> f302774cf (implements copy-on-write for staked-nodes (#19090))
         sync::{Arc, Once, RwLock, RwLockReadGuard},
     },
 };
@@ -36,13 +52,19 @@ pub type VoteAccountsHashMap = HashMap<Pubkey, (/*stake:*/ u64, VoteAccount)>;
 
 #[derive(Debug, AbiExample)]
 pub struct VoteAccounts {
+<<<<<<< HEAD
     vote_accounts: Arc<VoteAccountsHashMap>,
+=======
+    vote_accounts: HashMap<Pubkey, (u64 /*stake*/, ArcVoteAccount)>,
+>>>>>>> f302774cf (implements copy-on-write for staked-nodes (#19090))
     // Inner Arc is meant to implement copy-on-write semantics as opposed to
     // sharing mutations (hence RwLock<Arc<...>> instead of Arc<RwLock<...>>).
     staked_nodes: RwLock<
-        HashMap<
-            Pubkey, // VoteAccount.vote_state.node_pubkey.
-            u64,    // Total stake across all vote-accounts.
+        Arc<
+            HashMap<
+                Pubkey, // VoteAccount.vote_state.node_pubkey.
+                u64,    // Total stake across all vote-accounts.
+            >,
         >,
     >,
     staked_nodes_once: Once,
@@ -69,20 +91,19 @@ impl VoteAccount {
 }
 
 impl VoteAccounts {
-    pub fn staked_nodes(&self) -> HashMap<Pubkey, u64> {
+    pub fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
         self.staked_nodes_once.call_once(|| {
-            let mut staked_nodes = HashMap::new();
-            for (stake, vote_account) in
-                self.vote_accounts.values().filter(|(stake, _)| *stake != 0)
-            {
-                if let Some(node_pubkey) = vote_account.node_pubkey() {
-                    staked_nodes
-                        .entry(node_pubkey)
-                        .and_modify(|s| *s += *stake)
-                        .or_insert(*stake);
-                }
-            }
-            *self.staked_nodes.write().unwrap() = staked_nodes
+            let staked_nodes = self
+                .vote_accounts
+                .values()
+                .filter(|(stake, _)| *stake != 0)
+                .filter_map(|(stake, vote_account)| {
+                    let node_pubkey = vote_account.node_pubkey()?;
+                    Some((node_pubkey, stake))
+                })
+                .into_grouping_map()
+                .aggregate(|acc, _node_pubkey, stake| Some(acc.unwrap_or_default() + stake));
+            *self.staked_nodes.write().unwrap() = Arc::new(staked_nodes)
         });
         self.staked_nodes.read().unwrap().clone()
     }
@@ -135,9 +156,9 @@ impl VoteAccounts {
     fn add_node_stake(&mut self, stake: u64, vote_account: &VoteAccount) {
         if stake != 0 && self.staked_nodes_once.is_completed() {
             if let Some(node_pubkey) = vote_account.node_pubkey() {
-                self.staked_nodes
-                    .write()
-                    .unwrap()
+                let mut staked_nodes = self.staked_nodes.write().unwrap();
+                let staked_nodes = Arc::make_mut(&mut staked_nodes);
+                staked_nodes
                     .entry(node_pubkey)
                     .and_modify(|s| *s += stake)
                     .or_insert(stake);
@@ -148,7 +169,9 @@ impl VoteAccounts {
     fn sub_node_stake(&mut self, stake: u64, vote_account: &VoteAccount) {
         if stake != 0 && self.staked_nodes_once.is_completed() {
             if let Some(node_pubkey) = vote_account.node_pubkey() {
-                match self.staked_nodes.write().unwrap().entry(node_pubkey) {
+                let mut staked_nodes = self.staked_nodes.write().unwrap();
+                let staked_nodes = Arc::make_mut(&mut staked_nodes);
+                match staked_nodes.entry(node_pubkey) {
                     Entry::Vacant(_) => panic!("this should not happen!"),
                     Entry::Occupied(mut entry) => match entry.get().cmp(&stake) {
                         Ordering::Less => panic!("subtraction value exceeds node's stake"),
@@ -485,7 +508,7 @@ mod tests {
             if (k + 1) % 128 == 0 {
                 assert_eq!(
                     staked_nodes(&accounts[..k + 1]),
-                    vote_accounts.staked_nodes()
+                    *vote_accounts.staked_nodes()
                 );
             }
         }
@@ -495,7 +518,7 @@ mod tests {
             let (pubkey, (_, _)) = accounts.swap_remove(index);
             vote_accounts.remove(&pubkey);
             if (k + 1) % 32 == 0 {
-                assert_eq!(staked_nodes(&accounts), vote_accounts.staked_nodes());
+                assert_eq!(staked_nodes(&accounts), *vote_accounts.staked_nodes());
             }
         }
         // Modify the stakes for some of the accounts.
@@ -510,7 +533,7 @@ mod tests {
             }
             *stake = new_stake;
             if (k + 1) % 128 == 0 {
-                assert_eq!(staked_nodes(&accounts), vote_accounts.staked_nodes());
+                assert_eq!(staked_nodes(&accounts), *vote_accounts.staked_nodes());
             }
         }
         // Remove everything.
@@ -519,7 +542,7 @@ mod tests {
             let (pubkey, (_, _)) = accounts.swap_remove(index);
             vote_accounts.remove(&pubkey);
             if accounts.len() % 32 == 0 {
-                assert_eq!(staked_nodes(&accounts), vote_accounts.staked_nodes());
+                assert_eq!(staked_nodes(&accounts), *vote_accounts.staked_nodes());
             }
         }
         assert!(vote_accounts.staked_nodes.read().unwrap().is_empty());
@@ -556,6 +579,7 @@ mod tests {
             }
         }
     }
+<<<<<<< HEAD
 
     // Asserts that returned vote-accounts are copy-on-write references.
     #[test]
@@ -590,4 +614,6 @@ mod tests {
             }
         }
     }
+=======
+>>>>>>> f302774cf (implements copy-on-write for staked-nodes (#19090))
 }