removes RwLock+Once in favor of OnceCell in caching staked-nodes (#26313)

VoteAccounts uses std::sync::{RwLock, Once} to lazily compute and cache
staked_nodes:
https://github.com/solana-labs/solana/blob/032bee13a/runtime/src/vote_account.rs#L89-L104

This commit instead switches to using once_cell::sync::OnceCell which
provides this exact intended functionality by design.

runtime/src/bank.rs

@@ -2533,7 +2533,7 @@ impl Bank {
                     .stakes_cache
                     .stakes()
                     .vote_accounts()
-                    .delegated_stakes_iter()
+                    .delegated_stakes()
                     .map(|(pubkey, stake)| (*pubkey, stake))
                     .collect();
                 info!(

runtime/src/stakes.rs

@@ -57,7 +57,7 @@ impl StakesCache {
         self.0.read().unwrap()
     }
 
-    pub fn check_and_store(&self, pubkey: &Pubkey, account: &impl ReadableAccount) {
+    pub(crate) fn check_and_store(&self, pubkey: &Pubkey, account: &impl ReadableAccount) {
         // TODO: If the account is already cached as a vote or stake account
         // but the owner changes, then this needs to evict the account from
         // the cache. see:
@@ -110,12 +110,12 @@ impl StakesCache {
         }
     }
 
-    pub fn activate_epoch(&self, next_epoch: Epoch, thread_pool: &ThreadPool) {
+    pub(crate) fn activate_epoch(&self, next_epoch: Epoch, thread_pool: &ThreadPool) {
         let mut stakes = self.0.write().unwrap();
         stakes.activate_epoch(next_epoch, thread_pool)
     }
 
-    pub fn handle_invalid_keys(
+    pub(crate) fn handle_invalid_keys(
         &self,
         invalid_stake_keys: DashMap<Pubkey, InvalidCacheEntryReason>,
         invalid_vote_keys: DashMap<Pubkey, InvalidCacheEntryReason>,
@@ -231,7 +231,7 @@ impl Stakes<StakeAccount> {
         })
     }
 
-    pub fn history(&self) -> &StakeHistory {
+    pub(crate) fn history(&self) -> &StakeHistory {
         &self.stake_history
     }
 
@@ -303,7 +303,7 @@ impl Stakes<StakeAccount> {
     }
 
     /// Sum the lamports of the vote accounts and the delegated stake
-    pub fn vote_balance_and_staked(&self) -> u64 {
+    pub(crate) fn vote_balance_and_staked(&self) -> u64 {
         let get_stake = |stake_account: &StakeAccount| stake_account.delegation().stake;
         let get_lamports = |(_, vote_account): (_, &VoteAccount)| vote_account.lamports();
 
@@ -465,7 +465,7 @@ pub(crate) mod serde_stakes_enum_compat {
 }
 
 #[cfg(test)]
-pub mod tests {
+pub(crate) mod tests {
     use {
         super::*,
         rand::Rng,
@@ -476,7 +476,7 @@ pub mod tests {
     };
 
     //  set up some dummies for a staked node     ((     vote      )  (     stake     ))
-    pub fn create_staked_node_accounts(
+    pub(crate) fn create_staked_node_accounts(
         stake: u64,
     ) -> ((Pubkey, AccountSharedData), (Pubkey, AccountSharedData)) {
         let vote_pubkey = solana_sdk::pubkey::new_rand();
@@ -489,7 +489,10 @@ pub mod tests {
     }
 
     //   add stake to a vote_pubkey                               (   stake    )
-    pub fn create_stake_account(stake: u64, vote_pubkey: &Pubkey) -> (Pubkey, AccountSharedData) {
+    pub(crate) fn create_stake_account(
+        stake: u64,
+        vote_pubkey: &Pubkey,
+    ) -> (Pubkey, AccountSharedData) {
         let stake_pubkey = solana_sdk::pubkey::new_rand();
         (
             stake_pubkey,
@@ -503,7 +506,7 @@ pub mod tests {
         )
     }
 
-    pub fn create_warming_staked_node_accounts(
+    fn create_warming_staked_node_accounts(
         stake: u64,
         epoch: Epoch,
     ) -> ((Pubkey, AccountSharedData), (Pubkey, AccountSharedData)) {
@@ -517,7 +520,7 @@ pub mod tests {
     }
 
     // add stake to a vote_pubkey                               (   stake    )
-    pub fn create_warming_stake_account(
+    fn create_warming_stake_account(
         stake: u64,
         epoch: Epoch,
         vote_pubkey: &Pubkey,
@@ -847,7 +850,7 @@ pub mod tests {
     fn test_vote_balance_and_staked_normal() {
         let stakes_cache = StakesCache::default();
         impl Stakes<StakeAccount> {
-            pub fn vote_balance_and_warmed_staked(&self) -> u64 {
+            fn vote_balance_and_warmed_staked(&self) -> u64 {
                 let vote_balance: u64 = self
                     .vote_accounts
                     .iter()
@@ -855,7 +858,7 @@ pub mod tests {
                     .sum();
                 let warmed_stake: u64 = self
                     .vote_accounts
-                    .delegated_stakes_iter()
+                    .delegated_stakes()
                     .map(|(_pubkey, stake)| stake)
                     .sum();
                 vote_balance + warmed_stake

runtime/src/vote_account.rs

@@ -12,7 +12,7 @@ use {
         cmp::Ordering,
         collections::{hash_map::Entry, HashMap},
         iter::FromIterator,
-        sync::{Arc, Once, RwLock},
+        sync::Arc,
     },
     thiserror::Error,
 };
@@ -37,13 +37,12 @@ struct VoteAccountInner {
 
 pub type VoteAccountsHashMap = HashMap<Pubkey, (/*stake:*/ u64, VoteAccount)>;
 
-#[derive(Debug, AbiExample, Deserialize)]
+#[derive(Clone, Debug, AbiExample, Deserialize)]
 #[serde(from = "Arc<VoteAccountsHashMap>")]
 pub struct VoteAccounts {
     vote_accounts: Arc<VoteAccountsHashMap>,
-    // 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<
+    // Inner Arc is meant to implement copy-on-write semantics.
+    staked_nodes: OnceCell<
         Arc<
             HashMap<
                 Pubkey, // VoteAccount.vote_state.node_pubkey.
@@ -51,7 +50,6 @@ pub struct VoteAccounts {
             >,
         >,
     >,
-    staked_nodes_once: Once,
 }
 
 impl VoteAccount {
@@ -86,33 +84,34 @@ impl VoteAccounts {
         self.vote_accounts.len()
     }
 
-    pub fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
-        self.staked_nodes_once.call_once(|| {
-            let staked_nodes = self
-                .vote_accounts
+    pub fn staked_nodes(&self) -> Arc<HashMap</*node_pubkey:*/ Pubkey, /*stake:*/ u64>> {
+        self.staked_nodes
+            .get_or_init(|| {
+                Arc::new(
+                    self.vote_accounts
                         .values()
-                .filter(|(stake, _)| *stake != 0)
+                        .filter(|(stake, _)| *stake != 0u64)
                         .filter_map(|(stake, vote_account)| {
-                    let node_pubkey = vote_account.node_pubkey()?;
-                    Some((node_pubkey, stake))
+                            Some((vote_account.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()
+                        .aggregate(|acc, _node_pubkey, stake| {
+                            Some(acc.unwrap_or_default() + stake)
+                        }),
+                )
+            })
+            .clone()
     }
 
-    pub fn get(&self, pubkey: &Pubkey) -> Option<&VoteAccount> {
-        self.vote_accounts
-            .get(pubkey)
-            .map(|(_stake, vote_account)| vote_account)
+    pub(crate) fn get(&self, pubkey: &Pubkey) -> Option<&VoteAccount> {
+        let (_stake, vote_account) = self.vote_accounts.get(pubkey)?;
+        Some(vote_account)
     }
 
     pub fn get_delegated_stake(&self, pubkey: &Pubkey) -> u64 {
         self.vote_accounts
             .get(pubkey)
-            .map(|(stake, ..)| *stake)
+            .map(|(stake, _vote_account)| *stake)
             .unwrap_or_default()
     }
 
@@ -122,10 +121,10 @@ impl VoteAccounts {
             .map(|(vote_pubkey, (_stake, vote_account))| (vote_pubkey, vote_account))
     }
 
-    pub(crate) fn delegated_stakes_iter(&self) -> impl Iterator<Item = (&Pubkey, u64)> {
+    pub(crate) fn delegated_stakes(&self) -> impl Iterator<Item = (&Pubkey, u64)> {
         self.vote_accounts
             .iter()
-            .map(|(vote_pubkey, (stake, ..))| (vote_pubkey, *stake))
+            .map(|(vote_pubkey, (stake, _vote_account))| (vote_pubkey, *stake))
     }
 
     pub(crate) fn find_max_by_delegated_stake(&self) -> Option<&VoteAccount> {
@@ -172,24 +171,31 @@ impl VoteAccounts {
     }
 
     fn add_node_stake(&mut self, stake: u64, vote_account: &VoteAccount) {
-        if stake != 0 && self.staked_nodes_once.is_completed() {
+        if stake == 0u64 {
+            return;
+        }
+        let staked_nodes = match self.staked_nodes.get_mut() {
+            None => return,
+            Some(staked_nodes) => staked_nodes,
+        };
         if let Some(node_pubkey) = vote_account.node_pubkey() {
-                let mut staked_nodes = self.staked_nodes.write().unwrap();
-                let staked_nodes = Arc::make_mut(&mut staked_nodes);
-                staked_nodes
+            Arc::make_mut(staked_nodes)
                 .entry(node_pubkey)
                 .and_modify(|s| *s += stake)
                 .or_insert(stake);
         }
     }
-    }
 
     fn sub_node_stake(&mut self, stake: u64, vote_account: &VoteAccount) {
-        if stake != 0 && self.staked_nodes_once.is_completed() {
+        if stake == 0u64 {
+            return;
+        }
+        let staked_nodes = match self.staked_nodes.get_mut() {
+            None => return,
+            Some(staked_nodes) => staked_nodes,
+        };
         if let Some(node_pubkey) = vote_account.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) {
+            match Arc::make_mut(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"),
@@ -202,7 +208,6 @@ impl VoteAccounts {
         }
     }
 }
-}
 
 impl Serialize for VoteAccount {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
@@ -260,29 +265,7 @@ impl Default for VoteAccounts {
     fn default() -> Self {
         Self {
             vote_accounts: Arc::default(),
-            staked_nodes: RwLock::default(),
-            staked_nodes_once: Once::new(),
-        }
-    }
-}
-
-impl Clone for VoteAccounts {
-    fn clone(&self) -> Self {
-        if self.staked_nodes_once.is_completed() {
-            let staked_nodes = self.staked_nodes.read().unwrap().clone();
-            let other = Self {
-                vote_accounts: self.vote_accounts.clone(),
-                staked_nodes: RwLock::new(staked_nodes),
-                staked_nodes_once: Once::new(),
-            };
-            other.staked_nodes_once.call_once(|| {});
-            other
-        } else {
-            Self {
-                vote_accounts: self.vote_accounts.clone(),
-                staked_nodes: RwLock::default(),
-                staked_nodes_once: Once::new(),
-            }
+            staked_nodes: OnceCell::new(),
         }
     }
 }
@@ -292,7 +275,6 @@ impl PartialEq<VoteAccounts> for VoteAccounts {
         let Self {
             vote_accounts,
             staked_nodes: _,
-            staked_nodes_once: _,
         } = self;
         vote_accounts == &other.vote_accounts
     }
@@ -302,8 +284,7 @@ impl From<Arc<VoteAccountsHashMap>> for VoteAccounts {
     fn from(vote_accounts: Arc<VoteAccountsHashMap>) -> Self {
         Self {
             vote_accounts,
-            staked_nodes: RwLock::default(),
-            staked_nodes_once: Once::new(),
+            staked_nodes: OnceCell::new(),
         }
     }
 }
@@ -548,7 +529,7 @@ mod tests {
                 assert_eq!(staked_nodes(&accounts), *vote_accounts.staked_nodes());
             }
         }
-        assert!(vote_accounts.staked_nodes.read().unwrap().is_empty());
+        assert!(vote_accounts.staked_nodes.get().unwrap().is_empty());
     }
 
     // Asserts that returned staked-nodes are copy-on-write references.