Cleanup staking utils to divide functionality between delegate and normal node utitliites. Also replaces vote_states() with more generalized vote_accounts() in Bank. (#3070)

core/src/broadcast_stage.rs

@@ -49,7 +49,7 @@ impl Broadcast {
         let mut broadcast_table = cluster_info
             .read()
             .unwrap()
-            .sorted_tvu_peers(&staking_utils::node_stakes(&bank));
+            .sorted_tvu_peers(&staking_utils::delegated_stakes(&bank));
         // Layer 1, leader nodes are limited to the fanout size.
         broadcast_table.truncate(DATA_PLANE_FANOUT);
         inc_new_counter_info!("broadcast_service-num_peers", broadcast_table.len() + 1);

core/src/cluster_info.rs

@@ -877,9 +877,9 @@ impl ClusterInfo {
                 loop {
                     let start = timestamp();
                     let stakes: HashMap<_, _> = match bank_forks {
-                        Some(ref bank_forks) => {
+                        Some(ref bank_forks) => staking_utils::delegated_stakes(
-                            staking_utils::node_stakes(&bank_forks.read().unwrap().working_bank())
+                            &bank_forks.read().unwrap().working_bank(),
-                        }
+                        ),
                         None => HashMap::new(),
                     };
                     let _ = Self::run_gossip(&obj, &stakes, &blob_sender);

core/src/leader_schedule_utils.rs

@@ -5,7 +5,8 @@ use solana_sdk::pubkey::Pubkey;
 
 /// Return the leader schedule for the given epoch.
 fn leader_schedule(epoch_height: u64, bank: &Bank) -> LeaderSchedule {
-    let stakes = staking_utils::node_stakes_at_epoch(bank, epoch_height);
+    let stakes = staking_utils::delegated_stakes_at_epoch(bank, epoch_height)
+        .expect("epoch state must exist");
     let mut seed = [0u8; 32];
     seed[0..8].copy_from_slice(&epoch_height.to_le_bytes());
     let mut stakes: Vec<_> = stakes.into_iter().collect();
@@ -110,7 +111,7 @@ mod tests {
             GenesisBlock::new_with_leader(BOOTSTRAP_LEADER_TOKENS, pubkey, BOOTSTRAP_LEADER_TOKENS);
         let bank = Bank::new(&genesis_block);
 
-        let ids_and_stakes: Vec<_> = staking_utils::node_stakes(&bank).into_iter().collect();
+        let ids_and_stakes: Vec<_> = staking_utils::delegated_stakes(&bank).into_iter().collect();
         let seed = [0u8; 32];
         let leader_schedule =
             LeaderSchedule::new(&ids_and_stakes, seed, genesis_block.slots_per_epoch);

core/src/retransmit_stage.rs

@@ -40,7 +40,7 @@ fn retransmit(
             .to_owned(),
     );
     let (neighbors, children) = compute_retransmit_peers(
-        &staking_utils::node_stakes(&bank_forks.read().unwrap().working_bank()),
+        &staking_utils::delegated_stakes(&bank_forks.read().unwrap().working_bank()),
         cluster_info,
         DATA_PLANE_FANOUT,
         NEIGHBORHOOD_SIZE,

core/src/staking_utils.rs

@@ -1,7 +1,9 @@
 use hashbrown::HashMap;
 use solana_runtime::bank::Bank;
+use solana_sdk::account::Account;
 use solana_sdk::pubkey::Pubkey;
 use solana_vote_api::vote_state::VoteState;
+use std::borrow::Borrow;
 
 /// Looks through vote accounts, and finds the latest slot that has achieved
 /// supermajority lockout
@@ -15,97 +17,97 @@ pub fn get_supermajority_slot(bank: &Bank, epoch_height: u64) -> Option<u64> {
     find_supermajority_slot(supermajority_stake, stakes_and_lockouts.iter())
 }
 
-/// Collect the node Pubkey and staker account balance for nodes
+pub fn vote_account_balances(bank: &Bank) -> HashMap<Pubkey, u64> {
-/// that have non-zero balance in their corresponding staking accounts
+    let node_staked_accounts = node_staked_accounts(bank);
-pub fn node_stakes(bank: &Bank) -> HashMap<Pubkey, u64> {
+    node_staked_accounts
-    sum_node_stakes(&node_stakes_extractor(bank, |stake, _| stake))
+        .map(|(id, stake, _)| (id, stake))
-}
-
-/// Return the checkpointed stakes that should be used to generate a leader schedule.
-pub fn node_stakes_at_epoch(bank: &Bank, epoch_height: u64) -> HashMap<Pubkey, u64> {
-    sum_node_stakes(&node_stakes_at_epoch_extractor(
-        bank,
-        epoch_height,
-        |stake, _| stake,
-    ))
-}
-
-/// Sum up all the staking accounts for each delegate
-fn sum_node_stakes(stakes: &HashMap<Pubkey, Vec<u64>>) -> HashMap<Pubkey, u64> {
-    stakes
-        .iter()
-        .map(|(delegate, stakes)| (*delegate, stakes.iter().sum()))
         .collect()
 }
 
-/// Return the checkpointed stakes that should be used to generate a leader schedule.
+/// Collect the delegate account balance and vote states for delegates have non-zero balance in
-/// state_extractor takes (stake, vote_state) and maps to an output.
+/// any of their managed staking accounts
-fn node_stakes_at_epoch_extractor<F, T: Clone>(
+pub fn delegated_stakes(bank: &Bank) -> HashMap<Pubkey, u64> {
+    let node_staked_accounts = node_staked_accounts(bank);
+    let node_staked_vote_states = to_vote_state(node_staked_accounts);
+    to_delegated_stakes(node_staked_vote_states)
+}
+
+/// At the specified epoch, collect the node account balance and vote states for nodes that
+/// have non-zero balance in their corresponding staking accounts
+pub fn vote_account_balances_at_epoch(
     bank: &Bank,
     epoch_height: u64,
-    state_extractor: F,
+) -> Option<HashMap<Pubkey, u64>> {
-) -> HashMap<Pubkey, Vec<T>>
+    let node_staked_accounts = node_staked_accounts_at_epoch(bank, epoch_height);
-where
+    node_staked_accounts.map(|epoch_state| epoch_state.map(|(id, stake, _)| (*id, stake)).collect())
-    F: Fn(u64, &VoteState) -> T,
-{
-    let epoch_slot_height = epoch_height * bank.slots_per_epoch();
-    node_stakes_at_slot_extractor(bank, epoch_slot_height, state_extractor)
 }
 
-/// Return the checkpointed stakes that should be used to generate a leader schedule.
+/// At the specified epoch, collect the delgate account balance and vote states for delegates
-/// state_extractor takes (stake, vote_state) and maps to an output
+/// that have non-zero balance in any of their managed staking accounts
-fn node_stakes_at_slot_extractor<F, T: Clone>(
+pub fn delegated_stakes_at_epoch(bank: &Bank, epoch_height: u64) -> Option<HashMap<Pubkey, u64>> {
+    let node_staked_accounts = node_staked_accounts_at_epoch(bank, epoch_height);
+    let node_staked_vote_states = node_staked_accounts.map(to_vote_state);
+    node_staked_vote_states.map(to_delegated_stakes)
+}
+
+/// Collect the node account balance and vote states for nodes have non-zero balance in
+/// their corresponding staking accounts
+fn node_staked_accounts(bank: &Bank) -> impl Iterator<Item = (Pubkey, u64, Account)> {
+    bank.vote_accounts().filter_map(|(account_id, account)| {
+        filter_zero_balances(&account).map(|stake| (account_id, stake, account))
+    })
+}
+
+fn node_staked_accounts_at_epoch(
     bank: &Bank,
-    current_slot_height: u64,
+    epoch_height: u64,
-    state_extractor: F,
+) -> Option<impl Iterator<Item = (&Pubkey, u64, &Account)>> {
-) -> HashMap<Pubkey, Vec<T>>
+    bank.epoch_vote_accounts(epoch_height).map(|epoch_state| {
-where
+        epoch_state.into_iter().filter_map(|(account_id, account)| {
-    F: Fn(u64, &VoteState) -> T,
+            filter_zero_balances(account).map(|stake| (account_id, stake, account))
-{
+        })
-    let slot_height = current_slot_height.saturating_sub(bank.stakers_slot_offset());
+    })
-
-    let parents = bank.parents();
-    let mut banks = vec![bank];
-    banks.extend(parents.iter().map(|x| x.as_ref()));
-
-    let bank = banks
-        .iter()
-        .find(|bank| bank.slot() <= slot_height)
-        .unwrap_or_else(|| banks.last().unwrap());
-
-    node_stakes_extractor(bank, state_extractor)
 }
 
-/// Collect the node Pubkey and staker account balance for nodes
+fn filter_zero_balances(account: &Account) -> Option<u64> {
-/// that have non-zero balance in their corresponding staker accounts.
+    let balance = Bank::read_balance(&account);
-/// state_extractor takes (stake, vote_state) and maps to an output
+    if balance > 0 {
-fn node_stakes_extractor<F, T: Clone>(bank: &Bank, state_extractor: F) -> HashMap<Pubkey, Vec<T>>
+        Some(balance)
-where
+    } else {
-    F: Fn(u64, &VoteState) -> T,
+        None
-{
+    }
-    let mut map: HashMap<Pubkey, Vec<T>> = HashMap::new();
+}
-    let vote_states = bank.vote_states(|account_id, _| bank.get_balance(&account_id) > 0);
+
-    vote_states.into_iter().for_each(|(account_id, state)| {
+fn to_vote_state(
-        if map.contains_key(&state.delegate_id) {
+    node_staked_accounts: impl Iterator<Item = (impl Borrow<Pubkey>, u64, impl Borrow<Account>)>,
-            let entry = map.get_mut(&state.delegate_id).unwrap();
+) -> impl Iterator<Item = (u64, VoteState)> {
-            entry.push(state_extractor(bank.get_balance(&account_id), &state));
+    node_staked_accounts.filter_map(|(_, stake, account)| {
-        } else {
+        VoteState::deserialize(&account.borrow().userdata)
-            map.insert(
+            .ok()
-                state.delegate_id,
+            .map(|vote_state| (stake, vote_state))
-                vec![state_extractor(bank.get_balance(&account_id), &state)],
+    })
-            );
+}
-        }
+
+fn to_delegated_stakes(
+    node_staked_accounts: impl Iterator<Item = (u64, VoteState)>,
+) -> HashMap<Pubkey, u64> {
+    let mut map: HashMap<Pubkey, u64> = HashMap::new();
+    node_staked_accounts.for_each(|(stake, state)| {
+        let delegate = &state.delegate_id;
+        map.entry(*delegate)
+            .and_modify(|s| *s += stake)
+            .or_insert(stake);
     });
     map
 }
 
 fn epoch_stakes_and_lockouts(bank: &Bank, epoch_height: u64) -> Vec<(u64, Option<u64>)> {
-    node_stakes_at_epoch_extractor(bank, epoch_height, |stake, states| {
+    let node_staked_accounts =
-        (stake, states.root_slot)
+        node_staked_accounts_at_epoch(bank, epoch_height).expect("Bank state for epoch is missing");
-    })
+    let node_staked_vote_states = to_vote_state(node_staked_accounts);
-    .into_iter()
+    node_staked_vote_states
-    .flat_map(|(_, stake_and_states)| stake_and_states)
+        .map(|(stake, states)| (stake, states.root_slot))
-    .collect()
+        .collect()
 }
 
 fn find_supermajority_slot<'a, I>(supermajority_stake: u64, stakes_and_lockouts: I) -> Option<u64>
@@ -138,21 +140,13 @@ mod tests {
     use crate::voting_keypair::tests as voting_keypair_tests;
     use hashbrown::HashSet;
     use solana_sdk::genesis_block::GenesisBlock;
-    use solana_sdk::hash::Hash;
     use solana_sdk::pubkey::Pubkey;
     use solana_sdk::signature::{Keypair, KeypairUtil};
     use std::iter::FromIterator;
     use std::sync::Arc;
 
-    fn register_ticks(bank: &Bank, n: u64) -> (u64, u64, u64) {
+    fn new_from_parent(parent: &Arc<Bank>, slot: u64) -> Bank {
-        for _ in 0..n {
+        Bank::new_from_parent(parent, Pubkey::default(), slot)
-            bank.register_tick(&Hash::default());
-        }
-        (bank.tick_index(), bank.slot_index(), bank.epoch_height())
-    }
-
-    fn new_from_parent(parent: &Arc<Bank>) -> Bank {
-        Bank::new_from_parent(parent, Pubkey::default(), parent.slot() + 1)
     }
 
     #[test]
@@ -162,18 +156,20 @@ mod tests {
         let (genesis_block, _) =
             GenesisBlock::new_with_leader(bootstrap_tokens, pubkey, bootstrap_tokens);
         let bank = Bank::new(&genesis_block);
-        let bank = new_from_parent(&Arc::new(bank));
-        let ticks_per_offset = bank.stakers_slot_offset() * bank.ticks_per_slot();
-        register_ticks(&bank, ticks_per_offset);
-        assert_eq!(bank.slot_height(), bank.stakers_slot_offset());
 
+        // Epoch doesn't exist
         let mut expected = HashMap::new();
-        expected.insert(pubkey, vec![bootstrap_tokens - 2]);
+        assert_eq!(vote_account_balances_at_epoch(&bank, 10), None);
-        let bank = new_from_parent(&Arc::new(bank));
+
-        assert_eq!(
+        // First epoch has the bootstrap leader
-            node_stakes_at_slot_extractor(&bank, bank.slot_height(), |s, _| s),
+        expected.insert(genesis_block.bootstrap_leader_vote_account_id, 1);
-            expected
+        let expected = Some(expected);
-        );
+        assert_eq!(vote_account_balances_at_epoch(&bank, 0), expected);
+
+        // Second epoch carries same information
+        let bank = new_from_parent(&Arc::new(bank), 1);
+        assert_eq!(vote_account_balances_at_epoch(&bank, 0), expected);
+        assert_eq!(vote_account_balances_at_epoch(&bank, 1), expected);
     }
 
     #[test]
@@ -185,17 +181,27 @@ mod tests {
         let bank_voter = Keypair::new();
 
         // Give the validator some stake but don't setup a staking account
+        // Validator has no tokens staked, so they get filtered out. Only the bootstrap leader
+        // created by the genesis block will get included
         bank.transfer(1, &mint_keypair, validator.pubkey(), genesis_block.hash())
             .unwrap();
 
-        // Validator has no token staked, so they get filtered out. Only the bootstrap leader
+        // Make a mint vote account. Because the mint has nonzero stake, this
-        // created by the genesis block will get included
+        // should show up in the active set
-        let expected: Vec<_> = epoch_stakes_and_lockouts(&bank, 0);
-        assert_eq!(expected, vec![(1, None)]);
-
         voting_keypair_tests::new_vote_account_with_vote(&mint_keypair, &bank_voter, &bank, 499, 0);
 
-        let result: HashSet<_> = HashSet::from_iter(epoch_stakes_and_lockouts(&bank, 0));
+        // Have to wait until the epoch at (stakers_slot_offset / slots_per_epoch) + 1
+        // for the new votes to take effect. Earlier epochs were generated by genesis
+        let epoch = (bank.stakers_slot_offset() / bank.slots_per_epoch()) + 1;
+        let epoch_slot = epoch * bank.slots_per_epoch();
+        let epoch_slot_offset = epoch_slot - bank.stakers_slot_offset();
+
+        let bank = new_from_parent(&Arc::new(bank), epoch_slot_offset);
+
+        let result: Vec<_> = epoch_stakes_and_lockouts(&bank, 0);
+        assert_eq!(result, vec![(1, None)]);
+
+        let result: HashSet<_> = HashSet::from_iter(epoch_stakes_and_lockouts(&bank, epoch));
         let expected: HashSet<_> = HashSet::from_iter(vec![(1, None), (499, None)]);
         assert_eq!(result, expected);
     }
@@ -248,13 +254,22 @@ mod tests {
     }
 
     #[test]
-    fn test_sum_node_stakes() {
+    fn test_to_delegated_stakes() {
-        let mut stakes = HashMap::new();
+        let mut stakes = Vec::new();
-        stakes.insert(Pubkey::default(), vec![1, 2, 3, 4, 5]);
+        let delegate1 = Keypair::new().pubkey();
-        assert_eq!(sum_node_stakes(&stakes).len(), 1);
+        let delegate2 = Keypair::new().pubkey();
-        assert_eq!(
+
-            sum_node_stakes(&stakes).get(&Pubkey::default()),
+        // Delegate 1 has stake of 3
-            Some(&15_u64)
+        for i in 0..3 {
-        );
+            stakes.push((i, VoteState::new(delegate1)));
+        }
+
+        // Delegate 1 has stake of 5
+        stakes.push((5, VoteState::new(delegate2)));
+
+        let result = to_delegated_stakes(stakes.into_iter());
+        assert_eq!(result.len(), 2);
+        assert_eq!(result[&delegate1], 3);
+        assert_eq!(result[&delegate2], 5);
     }
 }

runtime/src/bank.rs

@@ -179,7 +179,6 @@ impl Bank {
         bank.epoch_vote_accounts = {
             let mut epoch_vote_accounts = parent.epoch_vote_accounts.clone();
             let epoch = bank.epoch_from_stakers_slot_offset();
-
             // update epoch_vote_states cache
             //  if my parent didn't populate for this epoch, we've
             //  crossed a boundary
@@ -788,33 +787,8 @@ impl Bank {
     }
 
     ///  vote accounts for the specific epoch
-    pub fn epoch_vote_accounts<F, T>(&self, epoch: u64, filter: F) -> Option<HashMap<Pubkey, T>>
+    pub fn epoch_vote_accounts(&self, epoch: u64) -> Option<&HashMap<Pubkey, Account>> {
-    where
+        self.epoch_vote_accounts.get(&epoch)
-        F: Fn(&Pubkey, &Account) -> Option<(Pubkey, T)>,
-    {
-        self.epoch_vote_accounts.get(&epoch).map(|accounts| {
-            accounts
-                .iter()
-                .filter_map(|(pubkey, account)| filter(pubkey, account))
-                .collect()
-        })
-    }
-
-    pub fn vote_states<F>(&self, cond: F) -> HashMap<Pubkey, VoteState>
-    where
-        F: Fn(&Pubkey, &VoteState) -> bool,
-    {
-        self.accounts()
-            .get_vote_accounts(self.accounts_id)
-            .filter_map(|(p, account)| {
-                if let Ok(vote_state) = VoteState::deserialize(&account.userdata) {
-                    if cond(&p, &vote_state) {
-                        return Some((p, vote_state));
-                    }
-                }
-                None
-            })
-            .collect()
     }
 
     /// Return the number of slots since the last epoch boundary.
@@ -1497,30 +1471,31 @@ mod tests {
 
         let parent = Arc::new(Bank::new(&genesis_block));
 
-        let vote_accounts0 = parent.epoch_vote_accounts(0, |pubkey, account| {
+        let vote_accounts0: Option<HashMap<_, _>> = parent.epoch_vote_accounts(0).map(|accounts| {
-            if let Ok(vote_state) = VoteState::deserialize(&account.userdata) {
+            accounts
-                if vote_state.delegate_id == leader_id {
+                .iter()
-                    Some((*pubkey, true))
+                .filter_map(|(pubkey, account)| {
-                } else {
+                    if let Ok(vote_state) = VoteState::deserialize(&account.userdata) {
-                    None
+                        if vote_state.delegate_id == leader_id {
-                }
+                            Some((*pubkey, true))
-            } else {
+                        } else {
-                None
+                            None
-            }
+                        }
+                    } else {
+                        None
+                    }
+                })
+                .collect()
         });
         assert!(vote_accounts0.is_some());
         assert!(vote_accounts0.iter().len() != 0);
 
-        fn all(key: &Pubkey, _account: &Account) -> Option<(Pubkey, Option<()>)> {
-            Some((*key, None))
-        }
-
         let mut i = 1;
         loop {
             if i > STAKERS_SLOT_OFFSET / SLOTS_PER_EPOCH {
                 break;
             }
-            assert!(parent.epoch_vote_accounts(i, all).is_some());
+            assert!(parent.epoch_vote_accounts(i).is_some());
             i += 1;
         }
 
@@ -1531,7 +1506,7 @@ mod tests {
             SLOTS_PER_EPOCH - (STAKERS_SLOT_OFFSET % SLOTS_PER_EPOCH),
         );
 
-        assert!(child.epoch_vote_accounts(i, all).is_some());
+        assert!(child.epoch_vote_accounts(i).is_some());
 
         // child crosses epoch boundary but isn't the first slot in the epoch
         let child = Bank::new_from_parent(
@@ -1539,7 +1514,7 @@ mod tests {
             leader_id,
             SLOTS_PER_EPOCH - (STAKERS_SLOT_OFFSET % SLOTS_PER_EPOCH) + 1,
         );
-        assert!(child.epoch_vote_accounts(i, all).is_some());
+        assert!(child.epoch_vote_accounts(i).is_some());
     }
 
 }