bank: add current_epoch_staked_nodes() (solana-labs#2650)

* bank: add current_epoch_staked_nodes()

Add current_epoch_staked_nodes() which returns the staked nodes for the
current epoch. Remove Bank::staked_nodes() which used to return the
bank's view of staked nodes updated to the last vote processed.

The updated call sites don't really need super up to date stake info,
and with this change we can stop updating staked node info for every
vote on every bank. Instead we now compute it once per epoch.

* bank: current_epoch_stakes: explain why self.epoch + 1

core/src/banking_stage/latest_unprocessed_votes.rs

@@ -132,7 +132,7 @@ pub(crate) fn weighted_random_order_by_stake<'a>(
     pubkeys: impl Iterator<Item = &'a Pubkey>,
 ) -> impl Iterator<Item = Pubkey> + 'static {
     // Efraimidis and Spirakis algo for weighted random sample without replacement
-    let staked_nodes = bank.staked_nodes();
+    let staked_nodes = bank.current_epoch_staked_nodes();
     let mut pubkey_with_weight: Vec<(f64, Pubkey)> = pubkeys
         .filter_map(|&pubkey| {
             let stake = staked_nodes.get(&pubkey).copied().unwrap_or(0);
@@ -420,6 +420,7 @@ mod tests {
         solana_perf::packet::{Packet, PacketBatch, PacketFlags},
         solana_runtime::{
             bank::Bank,
+            epoch_stakes::EpochStakes,
             genesis_utils::{self, ValidatorVoteKeypairs},
         },
         solana_sdk::{hash::Hash, signature::Signer, system_transaction::transfer},
@@ -845,7 +846,12 @@ mod tests {
     #[test]
     fn test_forwardable_packets() {
         let latest_unprocessed_votes = LatestUnprocessedVotes::new();
-        let bank = Arc::new(Bank::default_for_tests());
+        let mut bank = Bank::default_for_tests();
+        bank.set_epoch_stakes_for_test(
+            bank.epoch().saturating_add(1),
+            EpochStakes::new_for_tests(HashMap::new(), bank.epoch().saturating_add(1)),
+        );
+        let bank = Arc::new(bank);
         let mut forward_packet_batches_by_accounts =
             ForwardPacketBatchesByAccounts::new_with_default_batch_limits();
 

core/src/repair/quic_endpoint.rs

@@ -724,7 +724,7 @@ async fn prune_connection_cache(
     debug_assert!(prune_cache_pending.load(Ordering::Relaxed));
     let staked_nodes = {
         let root_bank = bank_forks.read().unwrap().root_bank();
-        root_bank.staked_nodes()
+        root_bank.current_epoch_staked_nodes()
     };
     {
         let mut cache = cache.lock().await;

core/src/staked_nodes_updater_service.rs

@@ -32,7 +32,7 @@ impl StakedNodesUpdaterService {
                 while !exit.load(Ordering::Relaxed) {
                     let stakes = {
                         let root_bank = bank_forks.read().unwrap().root_bank();
-                        root_bank.staked_nodes()
+                        root_bank.current_epoch_staked_nodes()
                     };
                     let overrides = staked_nodes_overrides.read().unwrap().clone();
                     *staked_nodes.write().unwrap() = StakedNodes::new(stakes, overrides);

gossip/src/cluster_info.rs

@@ -1895,7 +1895,7 @@ impl ClusterInfo {
                         Some(ref bank_forks) => {
                             let root_bank = bank_forks.read().unwrap().root_bank();
                             (
-                                root_bank.staked_nodes(),
+                                root_bank.current_epoch_staked_nodes(),
                                 Some(root_bank.feature_set.clone()),
                             )
                         }
@@ -2704,7 +2704,7 @@ impl ClusterInfo {
             Some(bank_forks) => {
                 let bank = bank_forks.read().unwrap().root_bank();
                 let feature_set = bank.feature_set.clone();
-                (Some(feature_set), bank.staked_nodes())
+                (Some(feature_set), bank.current_epoch_staked_nodes())
             }
         };
         self.process_packets(

ledger/src/leader_schedule_utils.rs

@@ -99,7 +99,7 @@ mod tests {
         let bank = Bank::new_for_tests(&genesis_config);
 
         let pubkeys_and_stakes: Vec<_> = bank
-            .staked_nodes()
+            .current_epoch_staked_nodes()
             .iter()
             .map(|(pubkey, stake)| (*pubkey, *stake))
             .collect();

runtime/src/bank.rs

@@ -5905,10 +5905,6 @@ impl Bank {
             });
     }
 
-    pub fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
-        self.stakes_cache.stakes().staked_nodes()
-    }
-
     /// current vote accounts for this bank along with the stake
     ///   attributed to each account
     pub fn vote_accounts(&self) -> Arc<VoteAccountsHashMap> {
@@ -5923,6 +5919,15 @@ impl Bank {
         Some(vote_account.clone())
     }
 
+    /// Get the EpochStakes for the current Bank::epoch
+    pub fn current_epoch_stakes(&self) -> &EpochStakes {
+        // The stakes for a given epoch (E) in self.epoch_stakes are keyed by leader schedule epoch
+        // (E + 1) so the stakes for the current epoch are stored at self.epoch_stakes[E + 1]
+        self.epoch_stakes
+            .get(&self.epoch.saturating_add(1))
+            .expect("Current epoch stakes must exist")
+    }
+
     /// Get the EpochStakes for a given epoch
     pub fn epoch_stakes(&self, epoch: Epoch) -> Option<&EpochStakes> {
         self.epoch_stakes.get(&epoch)
@@ -5932,6 +5937,11 @@ impl Bank {
         &self.epoch_stakes
     }
 
+    /// Get the staked nodes map for the current Bank::epoch
+    pub fn current_epoch_staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
+        self.current_epoch_stakes().stakes().staked_nodes()
+    }
+
     pub fn epoch_staked_nodes(&self, epoch: Epoch) -> Option<Arc<HashMap<Pubkey, u64>>> {
         Some(self.epoch_stakes.get(&epoch)?.stakes().staked_nodes())
     }

turbine/src/quic_endpoint.rs

@@ -575,7 +575,7 @@ async fn prune_connection_cache(
     debug_assert!(prune_cache_pending.load(Ordering::Relaxed));
     let staked_nodes = {
         let root_bank = bank_forks.read().unwrap().root_bank();
-        root_bank.staked_nodes()
+        root_bank.current_epoch_staked_nodes()
     };
     {
         let mut cache = cache.lock().await;