From 6a34b11dd0b169469363ea703dabb0ef9eef1632 Mon Sep 17 00:00:00 2001 From: Sagar Dhawan Date: Fri, 1 Mar 2019 17:31:59 -0800 Subject: [PATCH] Sum up all stakes for a delegate when calculating stake (#3045) --- src/staking_utils.rs | 84 +++++++++++++++++++++++++++++--------------- 1 file changed, 55 insertions(+), 29 deletions(-) diff --git a/src/staking_utils.rs b/src/staking_utils.rs index cb62dad011..c20f1c438d 100644 --- a/src/staking_utils.rs +++ b/src/staking_utils.rs @@ -8,31 +8,43 @@ use solana_sdk::vote_program::VoteState; pub fn get_supermajority_slot(bank: &Bank, epoch_height: u64) -> Option { // Find the amount of stake needed for supermajority let stakes_and_lockouts = epoch_stakes_and_lockouts(bank, epoch_height); - let total_stake: u64 = stakes_and_lockouts.values().map(|s| s.0).sum(); + let total_stake: u64 = stakes_and_lockouts.iter().map(|s| s.0).sum(); let supermajority_stake = total_stake * 2 / 3; // Filter out the states that don't have a max lockout - find_supermajority_slot(supermajority_stake, stakes_and_lockouts.values()) + find_supermajority_slot(supermajority_stake, stakes_and_lockouts.iter()) } /// Collect the node Pubkey and staker account balance for nodes /// that have non-zero balance in their corresponding staking accounts pub fn node_stakes(bank: &Bank) -> HashMap { - node_stakes_extractor(bank, |stake, _| stake) + sum_node_stakes(&node_stakes_extractor(bank, |stake, _| 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 { - node_stakes_at_epoch_extractor(bank, epoch_height, |stake, _| stake) + 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>) -> HashMap { + stakes + .iter() + .map(|(delegate, stakes)| (*delegate, stakes.iter().sum())) + .collect() } /// Return the checkpointed stakes that should be used to generate a leader schedule. /// state_extractor takes (stake, vote_state) and maps to an output. -fn node_stakes_at_epoch_extractor( +fn node_stakes_at_epoch_extractor( bank: &Bank, epoch_height: u64, state_extractor: F, -) -> HashMap +) -> HashMap> where F: Fn(u64, &VoteState) -> T, { @@ -42,11 +54,11 @@ where /// Return the checkpointed stakes that should be used to generate a leader schedule. /// state_extractor takes (stake, vote_state) and maps to an output -fn node_stakes_at_slot_extractor( +fn node_stakes_at_slot_extractor( bank: &Bank, current_slot_height: u64, state_extractor: F, -) -> HashMap +) -> HashMap> where F: Fn(u64, &VoteState) -> T, { @@ -67,26 +79,33 @@ where /// Collect the node Pubkey and staker account balance for nodes /// that have non-zero balance in their corresponding staker accounts. /// state_extractor takes (stake, vote_state) and maps to an output -fn node_stakes_extractor(bank: &Bank, state_extractor: F) -> HashMap +fn node_stakes_extractor(bank: &Bank, state_extractor: F) -> HashMap> where F: Fn(u64, &VoteState) -> T, { - bank.vote_states(|account_id, _| bank.get_balance(&account_id) > 0) - .iter() - .map(|(account_id, state)| { - ( + let mut map: HashMap> = 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)| { + if map.contains_key(&state.delegate_id) { + let entry = map.get_mut(&state.delegate_id).unwrap(); + entry.push(state_extractor(bank.get_balance(&account_id), &state)); + } else { + map.insert( state.delegate_id, - state_extractor(bank.get_balance(&account_id), &state), - ) - }) - .collect() + vec![state_extractor(bank.get_balance(&account_id), &state)], + ); + } + }); + map } -fn epoch_stakes_and_lockouts( - bank: &Bank, - epoch_height: u64, -) -> HashMap)> { - node_stakes_at_epoch_extractor(bank, epoch_height, |stake, state| (stake, state.root_slot)) +fn epoch_stakes_and_lockouts(bank: &Bank, epoch_height: u64) -> Vec<(u64, Option)> { + node_stakes_at_epoch_extractor(bank, epoch_height, |stake, states| { + (stake, states.root_slot) + }) + .into_iter() + .flat_map(|(_, stake_and_states)| stake_and_states) + .collect() } fn find_supermajority_slot<'a, I>(supermajority_stake: u64, stakes_and_lockouts: I) -> Option @@ -148,7 +167,7 @@ mod tests { assert_eq!(bank.slot_height(), bank.stakers_slot_offset()); let mut expected = HashMap::new(); - expected.insert(pubkey, bootstrap_tokens - 1); + expected.insert(pubkey, vec![bootstrap_tokens - 1]); let bank = new_from_parent(&Arc::new(bank)); assert_eq!( node_stakes_at_slot_extractor(&bank, bank.slot_height(), |s, _| s), @@ -170,16 +189,12 @@ mod tests { // Validator has no token staked, so they get filtered out. Only the bootstrap leader // created by the genesis block will get included - let expected: Vec<_> = epoch_stakes_and_lockouts(&bank, 0) - .values() - .cloned() - .collect(); + 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).values().cloned()); + let result: HashSet<_> = HashSet::from_iter(epoch_stakes_and_lockouts(&bank, 0)); let expected: HashSet<_> = HashSet::from_iter(vec![(1, None), (499, None)]); assert_eq!(result, expected); } @@ -230,4 +245,15 @@ mod tests { Some(3) ); } + + #[test] + fn test_sum_node_stakes() { + let mut stakes = HashMap::new(); + stakes.insert(Pubkey::default(), vec![1, 2, 3, 4, 5]); + assert_eq!(sum_node_stakes(&stakes).len(), 1); + assert_eq!( + sum_node_stakes(&stakes).get(&Pubkey::default()), + Some(&15_u64) + ); + } }