diff --git a/core/src/bank_weight_fork_choice.rs b/core/src/bank_weight_fork_choice.rs new file mode 100644 index 0000000000..d5500f4bc4 --- /dev/null +++ b/core/src/bank_weight_fork_choice.rs @@ -0,0 +1,152 @@ +use crate::{ + consensus::{ComputedBankState, Tower}, + fork_choice::ForkChoice, + progress_map::{ForkStats, ProgressMap}, +}; +use solana_ledger::bank_forks::BankForks; +use solana_runtime::bank::Bank; +use solana_sdk::timing; +use std::time::Instant; +use std::{ + collections::{HashMap, HashSet}, + sync::{Arc, RwLock}, +}; + +#[derive(Default)] +pub struct BankWeightForkChoice {} + +impl ForkChoice for BankWeightForkChoice { + fn compute_bank_stats( + &mut self, + bank: &Bank, + _tower: &Tower, + progress: &mut ProgressMap, + computed_bank_stats: &ComputedBankState, + ) { + let bank_slot = bank.slot(); + // Only time progress map should be missing a bank slot + // is if this node was the leader for this slot as those banks + // are not replayed in replay_active_banks() + let parent_weight = bank + .parent() + .and_then(|b| progress.get(&b.slot())) + .map(|x| x.fork_stats.fork_weight) + .unwrap_or(0); + + let stats = progress + .get_fork_stats_mut(bank_slot) + .expect("All frozen banks must exist in the Progress map"); + + let ComputedBankState { bank_weight, .. } = computed_bank_stats; + stats.weight = *bank_weight; + stats.fork_weight = stats.weight + parent_weight; + } + + // Returns: + // 1) The heaviest overall bank + // 2) The heavest bank on the same fork as the last vote (doesn't require a + // switching proof to vote for) + fn select_forks( + &self, + frozen_banks: &[Arc], + tower: &Tower, + progress: &ProgressMap, + ancestors: &HashMap>, + _bank_forks: &RwLock, + ) -> (Arc, Option>) { + let tower_start = Instant::now(); + assert!(!frozen_banks.is_empty()); + let num_frozen_banks = frozen_banks.len(); + + trace!("frozen_banks {}", frozen_banks.len()); + let num_old_banks = frozen_banks + .iter() + .filter(|b| b.slot() < tower.root().unwrap_or(0)) + .count(); + + let last_vote = tower.last_vote().slots.last().cloned(); + let mut heaviest_bank_on_same_fork = None; + let mut heaviest_same_fork_weight = 0; + let stats: Vec<&ForkStats> = frozen_banks + .iter() + .map(|bank| { + // Only time progress map should be missing a bank slot + // is if this node was the leader for this slot as those banks + // are not replayed in replay_active_banks() + let stats = progress + .get_fork_stats(bank.slot()) + .expect("All frozen banks must exist in the Progress map"); + + if let Some(last_vote) = last_vote { + if ancestors + .get(&bank.slot()) + .expect("Entry in frozen banks must exist in ancestors") + .contains(&last_vote) + { + // Descendant of last vote cannot be locked out + assert!(!stats.is_locked_out); + + // ancestors(slot) should not contain the slot itself, + // so we should never get the same bank as the last vote + assert_ne!(bank.slot(), last_vote); + // highest weight, lowest slot first. frozen_banks is sorted + // from least slot to greatest slot, so if two banks have + // the same fork weight, the lower slot will be picked + if stats.fork_weight > heaviest_same_fork_weight { + heaviest_bank_on_same_fork = Some(bank.clone()); + heaviest_same_fork_weight = stats.fork_weight; + } + } + } + + stats + }) + .collect(); + let num_not_recent = stats.iter().filter(|s| !s.is_recent).count(); + let num_has_voted = stats.iter().filter(|s| s.has_voted).count(); + let num_empty = stats.iter().filter(|s| s.is_empty).count(); + let num_threshold_failure = stats.iter().filter(|s| !s.vote_threshold).count(); + let num_votable_threshold_failure = stats + .iter() + .filter(|s| s.is_recent && !s.has_voted && !s.vote_threshold) + .count(); + + let mut candidates: Vec<_> = frozen_banks.iter().zip(stats.iter()).collect(); + + //highest weight, lowest slot first + candidates.sort_by_key(|b| (b.1.fork_weight, 0i64 - b.0.slot() as i64)); + let rv = candidates + .last() + .expect("frozen banks was nonempty so candidates must also be nonempty"); + let ms = timing::duration_as_ms(&tower_start.elapsed()); + let weights: Vec<(u128, u64, u64)> = candidates + .iter() + .map(|x| (x.1.weight, x.0.slot(), x.1.block_height)) + .collect(); + debug!( + "@{:?} tower duration: {:?} len: {}/{} weights: {:?}", + timing::timestamp(), + ms, + candidates.len(), + stats.iter().filter(|s| !s.has_voted).count(), + weights, + ); + datapoint_debug!( + "replay_stage-select_forks", + ("frozen_banks", num_frozen_banks as i64, i64), + ("not_recent", num_not_recent as i64, i64), + ("has_voted", num_has_voted as i64, i64), + ("old_banks", num_old_banks as i64, i64), + ("empty_banks", num_empty as i64, i64), + ("threshold_failure", num_threshold_failure as i64, i64), + ( + "votable_threshold_failure", + num_votable_threshold_failure as i64, + i64 + ), + ("tower_duration", ms as i64, i64), + ); + + (rv.0.clone(), heaviest_bank_on_same_fork) + } +} diff --git a/core/src/consensus.rs b/core/src/consensus.rs index ce7ac5836e..c1bdd97cfe 100644 --- a/core/src/consensus.rs +++ b/core/src/consensus.rs @@ -21,7 +21,7 @@ use solana_vote_program::{ use std::{ collections::{BTreeMap, HashMap, HashSet}, ops::Bound::{Included, Unbounded}, - sync::Arc, + sync::{Arc, RwLock}, }; #[derive(PartialEq, Clone, Debug)] @@ -79,6 +79,14 @@ impl StakeLockout { } } +pub(crate) struct ComputedBankState { + pub stake_lockouts: HashMap, + pub total_staked: u64, + pub bank_weight: u128, + pub lockout_intervals: LockoutIntervals, + pub pubkey_votes: Vec<(Pubkey, Slot)>, +} + pub struct Tower { node_pubkey: Pubkey, threshold_depth: usize, @@ -102,10 +110,14 @@ impl Default for Tower { } impl Tower { - pub fn new(node_pubkey: &Pubkey, vote_account_pubkey: &Pubkey, bank_forks: &BankForks) -> Self { + pub fn new( + node_pubkey: &Pubkey, + vote_account_pubkey: &Pubkey, + root: Slot, + heaviest_bank: &Bank, + ) -> Self { let mut tower = Self::new_with_key(node_pubkey); - - tower.initialize_lockouts_from_bank_forks(&bank_forks, vote_account_pubkey); + tower.initialize_lockouts_from_bank_forks(vote_account_pubkey, root, heaviest_bank); tower } @@ -126,27 +138,28 @@ impl Tower { } } - pub fn collect_vote_lockouts( - &self, + pub(crate) fn collect_vote_lockouts( + node_pubkey: &Pubkey, bank_slot: u64, vote_accounts: F, ancestors: &HashMap>, all_pubkeys: &mut PubkeyReferences, - ) -> (HashMap, u64, u128, LockoutIntervals) + ) -> ComputedBankState where F: Iterator, { let mut stake_lockouts = HashMap::new(); - let mut total_stake = 0; - let mut total_weight = 0; + let mut total_staked = 0; + let mut bank_weight = 0; // Tree of intervals of lockouts of the form [slot, slot + slot.lockout], // keyed by end of the range let mut lockout_intervals = BTreeMap::new(); + let mut pubkey_votes = vec![]; for (key, (lamports, account)) in vote_accounts { if lamports == 0 { continue; } - trace!("{} {} with stake {}", self.node_pubkey, key, lamports); + trace!("{} {} with stake {}", node_pubkey, key, lamports); let vote_state = VoteState::from(&account); if vote_state.is_none() { datapoint_warn!( @@ -169,7 +182,7 @@ impl Tower { .push((vote.slot, key)); } - if key == self.node_pubkey || vote_state.node_pubkey == self.node_pubkey { + if key == *node_pubkey || vote_state.node_pubkey == *node_pubkey { debug!("vote state {:?}", vote_state); debug!( "observed slot {}", @@ -188,10 +201,15 @@ impl Tower { } let start_root = vote_state.root_slot; + // Add the latest vote to update the `heaviest_subtree_fork_choice` + if let Some(latest_vote) = vote_state.votes.back() { + pubkey_votes.push((key, latest_vote.slot)); + } + vote_state.process_slot_vote_unchecked(bank_slot); for vote in &vote_state.votes { - total_weight += vote.lockout() as u128 * lamports as u128; + bank_weight += vote.lockout() as u128 * lamports as u128; Self::update_ancestor_lockouts(&mut stake_lockouts, &vote, ancestors); } @@ -202,7 +220,7 @@ impl Tower { slot: root, }; trace!("ROOT: {}", vote.slot); - total_weight += vote.lockout() as u128 * lamports as u128; + bank_weight += vote.lockout() as u128 * lamports as u128; Self::update_ancestor_lockouts(&mut stake_lockouts, &vote, ancestors); } } @@ -211,7 +229,7 @@ impl Tower { confirmation_count: MAX_LOCKOUT_HISTORY as u32, slot: root, }; - total_weight += vote.lockout() as u128 * lamports as u128; + bank_weight += vote.lockout() as u128 * lamports as u128; Self::update_ancestor_lockouts(&mut stake_lockouts, &vote, ancestors); } @@ -232,9 +250,16 @@ impl Tower { // Update all the parents of this last vote with the stake of this vote account Self::update_ancestor_stakes(&mut stake_lockouts, vote.slot, lamports, ancestors); } - total_stake += lamports; + total_staked += lamports; + } + + ComputedBankState { + stake_lockouts, + total_staked, + bank_weight, + lockout_intervals, + pubkey_votes, } - (stake_lockouts, total_stake, total_weight, lockout_intervals) } pub fn is_slot_confirmed( @@ -509,7 +534,7 @@ impl Tower { } /// Update lockouts for all the ancestors - fn update_ancestor_lockouts( + pub(crate) fn update_ancestor_lockouts( stake_lockouts: &mut HashMap, vote: &Lockout, ancestors: &HashMap>, @@ -529,6 +554,28 @@ impl Tower { } } + pub(crate) fn find_heaviest_bank( + bank_forks: &RwLock, + node_pubkey: &Pubkey, + ) -> Option> { + let ancestors = bank_forks.read().unwrap().ancestors(); + let mut bank_weights: Vec<_> = bank_forks + .read() + .unwrap() + .frozen_banks() + .values() + .map(|b| { + ( + Self::bank_weight(node_pubkey, b, &ancestors), + b.parents().len(), + b.clone(), + ) + }) + .collect(); + bank_weights.sort_by_key(|b| (b.0, b.1)); + bank_weights.pop().map(|b| b.2) + } + /// Update stake for all the ancestors. /// Note, stake is the same for all the ancestor. fn update_ancestor_stakes( @@ -537,9 +584,8 @@ impl Tower { lamports: u64, ancestors: &HashMap>, ) { - // If there's no ancestors, that means this slot must be from before the current root, - // in which case the lockouts won't be calculated in bank_weight anyways, so ignore - // this slot + // If there's no ancestors, that means this slot must be from + // before the current root, so ignore this slot let vote_slot_ancestors = ancestors.get(&slot); if vote_slot_ancestors.is_none() { return; @@ -552,8 +598,13 @@ impl Tower { } } - fn bank_weight(&self, bank: &Bank, ancestors: &HashMap>) -> u128 { - let (_, _, bank_weight, _) = self.collect_vote_lockouts( + fn bank_weight( + node_pubkey: &Pubkey, + bank: &Bank, + ancestors: &HashMap>, + ) -> u128 { + let ComputedBankState { bank_weight, .. } = Self::collect_vote_lockouts( + node_pubkey, bank.slot(), bank.vote_accounts().into_iter(), ancestors, @@ -562,47 +613,28 @@ impl Tower { bank_weight } - fn find_heaviest_bank(&self, bank_forks: &BankForks) -> Option> { - let ancestors = bank_forks.ancestors(); - let mut bank_weights: Vec<_> = bank_forks - .frozen_banks() - .values() - .map(|b| { - ( - self.bank_weight(b, &ancestors), - b.parents().len(), - b.clone(), - ) - }) - .collect(); - bank_weights.sort_by_key(|b| (b.0, b.1)); - bank_weights.pop().map(|b| b.2) - } - fn initialize_lockouts_from_bank_forks( &mut self, - bank_forks: &BankForks, vote_account_pubkey: &Pubkey, + root: Slot, + heaviest_bank: &Bank, ) { - if let Some(bank) = self.find_heaviest_bank(bank_forks) { - let root = bank_forks.root(); - if let Some((_stake, vote_account)) = bank.vote_accounts().get(vote_account_pubkey) { - let mut vote_state = VoteState::deserialize(&vote_account.data) - .expect("vote_account isn't a VoteState?"); - vote_state.root_slot = Some(root); - vote_state.votes.retain(|v| v.slot > root); - trace!( - "{} lockouts initialized to {:?}", - self.node_pubkey, - vote_state - ); - - assert_eq!( - vote_state.node_pubkey, self.node_pubkey, - "vote account's node_pubkey doesn't match", - ); - self.lockouts = vote_state; - } + if let Some((_stake, vote_account)) = heaviest_bank.vote_accounts().get(vote_account_pubkey) + { + let mut vote_state = VoteState::deserialize(&vote_account.data) + .expect("vote_account isn't a VoteState?"); + vote_state.root_slot = Some(root); + vote_state.votes.retain(|v| v.slot > root); + trace!( + "{} lockouts initialized to {:?}", + self.node_pubkey, + vote_state + ); + assert_eq!( + vote_state.node_pubkey, self.node_pubkey, + "vote account's node_pubkey doesn't match", + ); + self.lockouts = vote_state; } } @@ -626,10 +658,13 @@ impl Tower { pub mod test { use super::*; use crate::{ + bank_weight_fork_choice::BankWeightForkChoice, cluster_info_vote_listener::VoteTracker, cluster_slots::ClusterSlots, + fork_choice::SelectVoteAndResetForkResult, + heaviest_subtree_fork_choice::HeaviestSubtreeForkChoice, progress_map::ForkProgress, - replay_stage::{HeaviestForkFailures, ReplayStage, SelectVoteAndResetForkResult}, + replay_stage::{HeaviestForkFailures, ReplayStage}, }; use solana_ledger::bank_forks::BankForks; use solana_runtime::{ @@ -645,7 +680,7 @@ pub mod test { signature::{Keypair, Signer}, }; use solana_vote_program::{ - vote_state::{Vote, VoteStateVersions}, + vote_state::{Vote, VoteStateVersions, MAX_LOCKOUT_HISTORY}, vote_transaction, }; use std::{ @@ -662,18 +697,26 @@ pub mod test { pub vote_pubkeys: Vec, pub bank_forks: RwLock, pub progress: ProgressMap, + pub heaviest_subtree_fork_choice: HeaviestSubtreeForkChoice, } impl VoteSimulator { pub(crate) fn new(num_keypairs: usize) -> Self { - let (validator_keypairs, node_pubkeys, vote_pubkeys, bank_forks, progress) = - Self::init_state(num_keypairs); + let ( + validator_keypairs, + node_pubkeys, + vote_pubkeys, + bank_forks, + progress, + heaviest_subtree_fork_choice, + ) = Self::init_state(num_keypairs); Self { validator_keypairs, node_pubkeys, vote_pubkeys, bank_forks: RwLock::new(bank_forks), progress, + heaviest_subtree_fork_choice, } } pub(crate) fn fill_bank_forks( @@ -716,6 +759,8 @@ pub mod test { } } new_bank.freeze(); + self.heaviest_subtree_fork_choice + .add_new_leaf_slot(new_bank.slot(), Some(new_bank.parent_slot())); self.bank_forks.write().unwrap().insert(new_bank); walk.forward(); } @@ -750,6 +795,8 @@ pub mod test { &ClusterSlots::default(), &self.bank_forks, &mut PubkeyReferences::default(), + &mut self.heaviest_subtree_fork_choice, + &mut BankWeightForkChoice::default(), ); let vote_bank = self @@ -766,7 +813,7 @@ pub mod test { heaviest_fork_failures, .. } = ReplayStage::select_vote_and_reset_forks( - &Some(vote_bank.clone()), + &vote_bank.clone(), &None, &ancestors, &descendants, @@ -795,6 +842,7 @@ pub mod test { &None, &mut PubkeyReferences::default(), None, + &mut self.heaviest_subtree_fork_choice, ) } @@ -880,6 +928,7 @@ pub mod test { Vec, BankForks, ProgressMap, + HeaviestSubtreeForkChoice, ) { let keypairs: HashMap<_, _> = std::iter::repeat_with(|| { let node_keypair = Keypair::new(); @@ -902,8 +951,16 @@ pub mod test { .map(|keys| keys.vote_keypair.pubkey()) .collect(); - let (bank_forks, progress) = initialize_state(&keypairs, 10_000); - (keypairs, node_pubkeys, vote_pubkeys, bank_forks, progress) + let (bank_forks, progress, heaviest_subtree_fork_choice) = + initialize_state(&keypairs, 10_000); + ( + keypairs, + node_pubkeys, + vote_pubkeys, + bank_forks, + progress, + heaviest_subtree_fork_choice, + ) } } @@ -911,7 +968,7 @@ pub mod test { pub(crate) fn initialize_state( validator_keypairs_map: &HashMap, stake: u64, - ) -> (BankForks, ProgressMap) { + ) -> (BankForks, ProgressMap, HeaviestSubtreeForkChoice) { let validator_keypairs: Vec<_> = validator_keypairs_map.values().collect(); let GenesisConfigInfo { genesis_config, @@ -931,7 +988,10 @@ pub mod test { 0, ForkProgress::new(bank0.last_blockhash(), None, None, 0, 0), ); - (BankForks::new(0, bank0), progress) + let bank_forks = BankForks::new(0, bank0); + let heaviest_subtree_fork_choice = + HeaviestSubtreeForkChoice::new_from_bank_forks(&bank_forks); + (bank_forks, progress, heaviest_subtree_fork_choice) } fn gen_stakes(stake_votes: &[(u64, &[u64])]) -> Vec<(Pubkey, (u64, Account))> { @@ -1291,20 +1351,32 @@ pub mod test { #[test] fn test_collect_vote_lockouts_sums() { //two accounts voting for slot 0 with 1 token staked - let accounts = gen_stakes(&[(1, &[0]), (1, &[0])]); - let tower = Tower::new_for_tests(0, 0.67); + let mut accounts = gen_stakes(&[(1, &[0]), (1, &[0])]); + accounts.sort_by_key(|(pk, _)| *pk); + let account_latest_votes: Vec<(Pubkey, Slot)> = + accounts.iter().map(|(pubkey, _)| (*pubkey, 0)).collect(); + let ancestors = vec![(1, vec![0].into_iter().collect()), (0, HashSet::new())] .into_iter() .collect(); - let (staked_lockouts, total_staked, bank_weight, _) = tower.collect_vote_lockouts( + let ComputedBankState { + stake_lockouts, + total_staked, + bank_weight, + mut pubkey_votes, + .. + } = Tower::collect_vote_lockouts( + &Pubkey::default(), 1, accounts.into_iter(), &ancestors, &mut PubkeyReferences::default(), ); - assert_eq!(staked_lockouts[&0].stake, 2); - assert_eq!(staked_lockouts[&0].lockout, 2 + 2 + 4 + 4); + assert_eq!(stake_lockouts[&0].stake, 2); + assert_eq!(stake_lockouts[&0].lockout, 2 + 2 + 4 + 4); assert_eq!(total_staked, 2); + pubkey_votes.sort(); + assert_eq!(pubkey_votes, account_latest_votes); // Each acccount has 1 vote in it. After simulating a vote in collect_vote_lockouts, // the account will have 2 votes, with lockout 2 + 4 = 6. So expected weight for @@ -1316,7 +1388,12 @@ pub mod test { fn test_collect_vote_lockouts_root() { let votes: Vec = (0..MAX_LOCKOUT_HISTORY as u64).collect(); //two accounts voting for slots 0..MAX_LOCKOUT_HISTORY with 1 token staked - let accounts = gen_stakes(&[(1, &votes), (1, &votes)]); + let mut accounts = gen_stakes(&[(1, &votes), (1, &votes)]); + accounts.sort_by_key(|(pk, _)| *pk); + let account_latest_votes: Vec<(Pubkey, Slot)> = accounts + .iter() + .map(|(pubkey, _)| (*pubkey, (MAX_LOCKOUT_HISTORY - 1) as Slot)) + .collect(); let mut tower = Tower::new_for_tests(0, 0.67); let mut ancestors = HashMap::new(); for i in 0..(MAX_LOCKOUT_HISTORY + 1) { @@ -1337,18 +1414,27 @@ pub mod test { + root_weight; let expected_bank_weight = 2 * vote_account_expected_weight; assert_eq!(tower.lockouts.root_slot, Some(0)); - let (staked_lockouts, _total_staked, bank_weight, _) = tower.collect_vote_lockouts( + let ComputedBankState { + stake_lockouts, + bank_weight, + mut pubkey_votes, + .. + } = Tower::collect_vote_lockouts( + &Pubkey::default(), MAX_LOCKOUT_HISTORY as u64, accounts.into_iter(), &ancestors, &mut PubkeyReferences::default(), ); for i in 0..MAX_LOCKOUT_HISTORY { - assert_eq!(staked_lockouts[&(i as u64)].stake, 2); + assert_eq!(stake_lockouts[&(i as u64)].stake, 2); } + // should be the sum of all the weights for root - assert!(staked_lockouts[&0].lockout > (2 * (1 << MAX_LOCKOUT_HISTORY))); + assert!(stake_lockouts[&0].lockout > (2 * (1 << MAX_LOCKOUT_HISTORY))); assert_eq!(bank_weight, expected_bank_weight); + pubkey_votes.sort(); + assert_eq!(pubkey_votes, account_latest_votes); } #[test] @@ -1381,7 +1467,7 @@ pub mod test { ); tower.record_vote(i, Hash::default()); } - assert!(!tower.check_vote_stake_threshold(MAX_LOCKOUT_HISTORY as u64 + 1, &stakes, 2)); + assert!(!tower.check_vote_stake_threshold(MAX_LOCKOUT_HISTORY as u64 + 1, &stakes, 2,)); } #[test] @@ -1609,48 +1695,7 @@ pub mod test { tower.record_vote(0, Hash::default()); tower.record_vote(1, Hash::default()); tower.record_vote(2, Hash::default()); - assert!(tower.check_vote_stake_threshold(6, &stakes, 2)); - } - - #[test] - fn test_lockout_is_updated_for_entire_branch() { - let mut stake_lockouts = HashMap::new(); - let vote = Lockout { - slot: 2, - confirmation_count: 1, - }; - let set: HashSet = vec![0u64, 1u64].into_iter().collect(); - let mut ancestors = HashMap::new(); - ancestors.insert(2, set); - let set: HashSet = vec![0u64].into_iter().collect(); - ancestors.insert(1, set); - Tower::update_ancestor_lockouts(&mut stake_lockouts, &vote, &ancestors); - assert_eq!(stake_lockouts[&0].lockout, 2); - assert_eq!(stake_lockouts[&1].lockout, 2); - assert_eq!(stake_lockouts[&2].lockout, 2); - } - - #[test] - fn test_lockout_is_updated_for_slot_or_lower() { - let mut stake_lockouts = HashMap::new(); - let set: HashSet = vec![0u64, 1u64].into_iter().collect(); - let mut ancestors = HashMap::new(); - ancestors.insert(2, set); - let set: HashSet = vec![0u64].into_iter().collect(); - ancestors.insert(1, set); - let vote = Lockout { - slot: 2, - confirmation_count: 1, - }; - Tower::update_ancestor_lockouts(&mut stake_lockouts, &vote, &ancestors); - let vote = Lockout { - slot: 1, - confirmation_count: 2, - }; - Tower::update_ancestor_lockouts(&mut stake_lockouts, &vote, &ancestors); - assert_eq!(stake_lockouts[&0].lockout, 2 + 4); - assert_eq!(stake_lockouts[&1].lockout, 2 + 4); - assert_eq!(stake_lockouts[&2].lockout, 2); + assert!(tower.check_vote_stake_threshold(6, &stakes, 2,)); } #[test] @@ -1729,7 +1774,7 @@ pub mod test { let total_stake = 4; let threshold_size = 0.67; let threshold_stake = (f64::ceil(total_stake as f64 * threshold_size)) as u64; - let tower_votes: Vec = (0..VOTE_THRESHOLD_DEPTH as u64).collect(); + let tower_votes: Vec = (0..VOTE_THRESHOLD_DEPTH as u64).collect(); let accounts = gen_stakes(&[ (threshold_stake, &[(VOTE_THRESHOLD_DEPTH - 2) as u64]), (total_stake - threshold_stake, &tower_votes[..]), @@ -1746,29 +1791,35 @@ pub mod test { for vote in &tower_votes { tower.record_vote(*vote, Hash::default()); } - let (staked_lockouts, total_staked, _, _) = tower.collect_vote_lockouts( + let ComputedBankState { + stake_lockouts, + total_staked, + .. + } = Tower::collect_vote_lockouts( + &Pubkey::default(), vote_to_evaluate, accounts.clone().into_iter(), &ancestors, &mut PubkeyReferences::default(), ); - assert!(tower.check_vote_stake_threshold(vote_to_evaluate, &staked_lockouts, total_staked)); + assert!(tower.check_vote_stake_threshold(vote_to_evaluate, &stake_lockouts, total_staked,)); // CASE 2: Now we want to evaluate a vote for slot VOTE_THRESHOLD_DEPTH + 1. This slot // will expire the vote in one of the vote accounts, so we should have insufficient // stake to pass the threshold let vote_to_evaluate = VOTE_THRESHOLD_DEPTH as u64 + 1; - let (staked_lockouts, total_staked, _, _) = tower.collect_vote_lockouts( + let ComputedBankState { + stake_lockouts, + total_staked, + .. + } = Tower::collect_vote_lockouts( + &Pubkey::default(), vote_to_evaluate, accounts.into_iter(), &ancestors, &mut PubkeyReferences::default(), ); - assert!(!tower.check_vote_stake_threshold( - vote_to_evaluate, - &staked_lockouts, - total_staked - )); + assert!(!tower.check_vote_stake_threshold(vote_to_evaluate, &stake_lockouts, total_staked,)); } fn vote_and_check_recent(num_votes: usize) { diff --git a/core/src/fork_choice.rs b/core/src/fork_choice.rs new file mode 100644 index 0000000000..117fcf027e --- /dev/null +++ b/core/src/fork_choice.rs @@ -0,0 +1,40 @@ +use crate::{ + consensus::{ComputedBankState, SwitchForkDecision, Tower}, + progress_map::ProgressMap, + replay_stage::HeaviestForkFailures, +}; +use solana_ledger::bank_forks::BankForks; +use solana_runtime::bank::Bank; +use std::{ + collections::{HashMap, HashSet}, + sync::{Arc, RwLock}, +}; + +pub(crate) struct SelectVoteAndResetForkResult { + pub vote_bank: Option<(Arc, SwitchForkDecision)>, + pub reset_bank: Option>, + pub heaviest_fork_failures: Vec, +} + +pub(crate) trait ForkChoice { + fn compute_bank_stats( + &mut self, + bank: &Bank, + tower: &Tower, + progress: &mut ProgressMap, + computed_bank_stats: &ComputedBankState, + ); + + // Returns: + // 1) The heaviest overall bbank + // 2) The heavest bank on the same fork as the last vote (doesn't require a + // switching proof to vote for) + fn select_forks( + &self, + frozen_banks: &[Arc], + tower: &Tower, + progress: &ProgressMap, + ancestors: &HashMap>, + bank_forks: &RwLock, + ) -> (Arc, Option>); +} diff --git a/core/src/heaviest_subtree_fork_choice/fork_choice.rs b/core/src/heaviest_subtree_fork_choice/fork_choice.rs new file mode 100644 index 0000000000..5fb46cca14 --- /dev/null +++ b/core/src/heaviest_subtree_fork_choice/fork_choice.rs @@ -0,0 +1,72 @@ +use crate::{ + consensus::{ComputedBankState, Tower}, + fork_choice::ForkChoice, + heaviest_subtree_fork_choice::HeaviestSubtreeForkChoice, + progress_map::ProgressMap, +}; +use solana_ledger::bank_forks::BankForks; +use solana_runtime::bank::Bank; +use std::{ + collections::{HashMap, HashSet}, + sync::{Arc, RwLock}, +}; + +impl ForkChoice for HeaviestSubtreeForkChoice { + fn compute_bank_stats( + &mut self, + bank: &Bank, + _tower: &Tower, + _progress: &mut ProgressMap, + computed_bank_stats: &ComputedBankState, + ) { + let ComputedBankState { pubkey_votes, .. } = computed_bank_stats; + + // Update `heaviest_subtree_fork_choice` to find the best fork to build on + let best_overall_slot = self.add_votes( + &pubkey_votes, + bank.epoch_stakes_map(), + bank.epoch_schedule(), + ); + + datapoint_info!( + "best_slot", + ("slot", bank.slot(), i64), + ("best_slot", best_overall_slot, i64), + ); + } + + // Returns: + // 1) The heaviest overall bbank + // 2) The heavest bank on the same fork as the last vote (doesn't require a + // switching proof to vote for) + fn select_forks( + &self, + _frozen_banks: &[Arc], + tower: &Tower, + _progress: &ProgressMap, + _ancestors: &HashMap>, + bank_forks: &RwLock, + ) -> (Arc, Option>) { + let last_vote = tower.last_vote().slots.last().cloned(); + let heaviest_slot_on_same_voted_fork = last_vote.map(|last_vote| { + let heaviest_slot_on_same_voted_fork = + self.best_slot(last_vote).expect("last_vote is a frozen bank so must have been added to heaviest_subtree_fork_choice at time of freezing"); + if heaviest_slot_on_same_voted_fork == last_vote { + None + } else { + Some(heaviest_slot_on_same_voted_fork) + } + }).unwrap_or(None); + let heaviest_slot = self.best_overall_slot(); + let r_bank_forks = bank_forks.read().unwrap(); + ( + r_bank_forks.get(heaviest_slot).unwrap().clone(), + heaviest_slot_on_same_voted_fork.map(|heaviest_slot_on_same_voted_fork| { + r_bank_forks + .get(heaviest_slot_on_same_voted_fork) + .unwrap() + .clone() + }), + ) + } +} diff --git a/core/src/heaviest_subtree_fork_choice/mod.rs b/core/src/heaviest_subtree_fork_choice/mod.rs new file mode 100644 index 0000000000..8a620d7627 --- /dev/null +++ b/core/src/heaviest_subtree_fork_choice/mod.rs @@ -0,0 +1,1040 @@ +#[cfg(test)] +use solana_ledger::bank_forks::BankForks; +use solana_runtime::{bank::Bank, epoch_stakes::EpochStakes}; +use solana_sdk::{ + clock::{Epoch, Slot}, + epoch_schedule::EpochSchedule, + pubkey::Pubkey, +}; +use std::{ + collections::{BTreeMap, HashMap}, + sync::Arc, +}; +#[cfg(test)] +use trees::{Tree, TreeWalk}; + +mod fork_choice; +pub type ForkWeight = u64; + +#[derive(PartialEq, Eq, Clone, Debug, PartialOrd, Ord)] +enum UpdateLabel { + Aggregate, + Add, + Subtract, +} + +#[derive(PartialEq, Eq, Clone, Debug)] +enum UpdateOperation { + Aggregate, + Add(u64), + Subtract(u64), +} + +impl UpdateOperation { + fn update_stake(&mut self, new_stake: u64) { + match self { + Self::Aggregate => panic!("Should not get here"), + Self::Add(stake) => *stake += new_stake, + Self::Subtract(stake) => *stake += new_stake, + } + } +} + +struct ForkInfo { + // Amount of stake that has voted for exactly this slot + stake_voted_at: ForkWeight, + // Amount of stake that has voted for this slot and the subtree + // rooted at this slot + stake_voted_subtree: ForkWeight, + // Best slot in the subtree rooted at this slot, does not + // have to be a direct child in `children` + best_slot: Slot, + parent: Option, + children: Vec, +} + +#[derive(Default)] +pub struct HeaviestSubtreeForkChoice { + fork_infos: HashMap, + latest_votes: HashMap, + root: Slot, +} + +impl HeaviestSubtreeForkChoice { + pub(crate) fn new(root: Slot) -> Self { + let mut heaviest_subtree_fork_choice = Self { + root, + ..HeaviestSubtreeForkChoice::default() + }; + heaviest_subtree_fork_choice.add_new_leaf_slot(root, None); + heaviest_subtree_fork_choice + } + + // Given a root and a list of `frozen_banks` sorted smallest to greatest by slot, + // return a new HeaviestSubtreeForkChoice + pub(crate) fn new_from_frozen_banks(root: Slot, frozen_banks: &[Arc]) -> Self { + let mut heaviest_subtree_fork_choice = HeaviestSubtreeForkChoice::new(root); + let mut prev_slot = root; + for bank in frozen_banks.iter() { + assert!(bank.is_frozen()); + if bank.slot() > root { + // Make sure the list is sorted + assert!(bank.slot() > prev_slot); + prev_slot = bank.slot(); + heaviest_subtree_fork_choice + .add_new_leaf_slot(bank.slot(), Some(bank.parent_slot())); + } + } + + heaviest_subtree_fork_choice + } + + #[cfg(test)] + pub(crate) fn new_from_bank_forks(bank_forks: &BankForks) -> Self { + let mut frozen_banks: Vec<_> = bank_forks.frozen_banks().values().cloned().collect(); + + frozen_banks.sort_by_key(|bank| bank.slot()); + let root = bank_forks.root(); + Self::new_from_frozen_banks(root, &frozen_banks) + } + + #[cfg(test)] + pub(crate) fn new_from_tree(forks: Tree) -> Self { + let root = forks.root().data; + let mut walk = TreeWalk::from(forks); + let mut heaviest_subtree_fork_choice = HeaviestSubtreeForkChoice::new(root); + while let Some(visit) = walk.get() { + let slot = visit.node().data; + if heaviest_subtree_fork_choice.fork_infos.contains_key(&slot) { + walk.forward(); + continue; + } + let parent = walk.get_parent().map(|n| n.data); + heaviest_subtree_fork_choice.add_new_leaf_slot(slot, parent); + walk.forward(); + } + + heaviest_subtree_fork_choice + } + + pub fn best_slot(&self, slot: Slot) -> Option { + self.fork_infos + .get(&slot) + .map(|fork_info| fork_info.best_slot) + } + + pub fn best_overall_slot(&self) -> Slot { + self.best_slot(self.root).unwrap() + } + + pub fn stake_voted_subtree(&self, slot: Slot) -> Option { + self.fork_infos + .get(&slot) + .map(|fork_info| fork_info.stake_voted_subtree) + } + + // Add new votes, returns the best slot + pub fn add_votes( + &mut self, + // newly updated votes on a fork + pubkey_votes: &[(Pubkey, Slot)], + epoch_stakes: &HashMap, + epoch_schedule: &EpochSchedule, + ) -> Slot { + // Generate the set of updates + let update_operations = + self.generate_update_operations(pubkey_votes, epoch_stakes, epoch_schedule); + + // Finalize all updates + self.process_update_operations(update_operations); + self.best_overall_slot() + } + + pub fn set_root(&mut self, root: Slot) { + self.root = root; + let mut pending_slots = vec![root]; + let mut new_fork_infos = HashMap::new(); + while !pending_slots.is_empty() { + let current_slot = pending_slots.pop().unwrap(); + let fork_info = self + .fork_infos + .remove(¤t_slot) + .expect("Anything reachable from root must exist in the map"); + for child in &fork_info.children { + pending_slots.push(*child); + } + new_fork_infos.insert(current_slot, fork_info); + } + + std::mem::swap(&mut self.fork_infos, &mut new_fork_infos); + self.fork_infos + .get_mut(&root) + .expect("new root must exist in fork_infos map") + .parent = None; + } + + pub fn add_new_leaf_slot(&mut self, slot: Slot, parent: Option) { + self.fork_infos + .entry(slot) + .and_modify(|slot_info| slot_info.parent = parent) + .or_insert(ForkInfo { + stake_voted_at: 0, + stake_voted_subtree: 0, + // The `best_slot` of a leaf is itself + best_slot: slot, + children: vec![], + parent, + }); + + if parent.is_none() { + return; + } + + let parent = parent.unwrap(); + + // Parent must already exist by time child is being added + self.fork_infos + .get_mut(&parent) + .unwrap() + .children + .push(slot); + + // Propagate leaf up the tree to any ancestors who considered the previous leaf + // the `best_slot` + self.propagate_new_leaf(slot, parent) + } + + fn propagate_new_leaf(&mut self, slot: Slot, parent: Slot) { + let parent_best_slot = self + .best_slot(parent) + .expect("parent must exist in self.fork_infos after being created above"); + let should_replace_parent_best_slot = parent_best_slot == parent + || (parent_best_slot > slot + // if `stake_voted_subtree(parent).unwrap() - stake_voted_at(parent) == 0` + // then none of the children of `parent` have been voted on. Because + // this new leaf also hasn't been voted on, it must have the same weight + // as any existing child, so this new leaf can replace the previous best child + // if the new leaf is a lower slot. + && self.stake_voted_subtree(parent).unwrap() - self.stake_voted_at(parent).unwrap() == 0); + + if should_replace_parent_best_slot { + let mut ancestor = Some(parent); + loop { + if ancestor.is_none() { + break; + } + let ancestor_fork_info = self.fork_infos.get_mut(&ancestor.unwrap()).unwrap(); + if ancestor_fork_info.best_slot == parent_best_slot { + ancestor_fork_info.best_slot = slot; + } else { + break; + } + ancestor = ancestor_fork_info.parent; + } + } + } + + #[allow(clippy::map_entry)] + fn insert_aggregate_operations( + &self, + update_operations: &mut BTreeMap<(Slot, UpdateLabel), UpdateOperation>, + slot: Slot, + ) { + for parent in self.ancestor_iterator(slot) { + let label = (parent, UpdateLabel::Aggregate); + if update_operations.contains_key(&label) { + break; + } else { + update_operations.insert(label, UpdateOperation::Aggregate); + } + } + } + + fn ancestor_iterator(&self, start_slot: Slot) -> AncestorIterator { + AncestorIterator::new(start_slot, &self.fork_infos) + } + + fn aggregate_slot(&mut self, slot: Slot) { + let mut stake_voted_subtree; + let mut best_slot = slot; + if let Some(fork_info) = self.fork_infos.get(&slot) { + stake_voted_subtree = fork_info.stake_voted_at; + let mut best_child_stake_voted_subtree = 0; + let mut best_child_slot = slot; + let should_print = fork_info.children.len() > 1; + for &child in &fork_info.children { + let child_stake_voted_subtree = self.stake_voted_subtree(child).unwrap(); + if should_print { + info!( + "child: {} of slot: {} has weight: {}", + child, slot, child_stake_voted_subtree + ); + } + stake_voted_subtree += child_stake_voted_subtree; + if best_child_slot == slot || + child_stake_voted_subtree > best_child_stake_voted_subtree || + // tiebreaker by slot height, prioritize earlier slot + (child_stake_voted_subtree == best_child_stake_voted_subtree && child < best_child_slot) + { + best_child_stake_voted_subtree = child_stake_voted_subtree; + best_child_slot = child; + best_slot = self + .best_slot(child) + .expect("`child` must exist in `self.fork_infos`"); + } + } + } else { + return; + } + + let fork_info = self.fork_infos.get_mut(&slot).unwrap(); + fork_info.stake_voted_subtree = stake_voted_subtree; + fork_info.best_slot = best_slot; + } + + fn generate_update_operations( + &mut self, + pubkey_votes: &[(Pubkey, Slot)], + epoch_stakes: &HashMap, + epoch_schedule: &EpochSchedule, + ) -> BTreeMap<(Slot, UpdateLabel), UpdateOperation> { + let mut update_operations: BTreeMap<(Slot, UpdateLabel), UpdateOperation> = BTreeMap::new(); + + // Sort the `pubkey_votes` in a BTreeMap by the slot voted + for &(pubkey, new_vote_slot) in pubkey_votes.iter() { + let pubkey_latest_vote = self.latest_votes.get(&pubkey).unwrap_or(&0); + // Filter out any votes or slots <= any slot this pubkey has + // already voted for, we only care about the latest votes + if new_vote_slot <= *pubkey_latest_vote { + continue; + } + + // Remove this pubkey stake from previous fork + if let Some(old_latest_vote_slot) = self.latest_votes.insert(pubkey, new_vote_slot) { + let epoch = epoch_schedule.get_epoch(old_latest_vote_slot); + let stake_update = epoch_stakes + .get(&epoch) + .map(|epoch_stakes| epoch_stakes.vote_account_stake(&pubkey)) + .unwrap_or(0); + if stake_update > 0 { + update_operations + .entry((old_latest_vote_slot, UpdateLabel::Subtract)) + .and_modify(|update| update.update_stake(stake_update)) + .or_insert(UpdateOperation::Subtract(stake_update)); + self.insert_aggregate_operations(&mut update_operations, old_latest_vote_slot); + } + } + + // Add this pubkey stake to new fork + let epoch = epoch_schedule.get_epoch(new_vote_slot); + let stake_update = epoch_stakes + .get(&epoch) + .map(|epoch_stakes| epoch_stakes.vote_account_stake(&pubkey)) + .unwrap_or(0); + update_operations + .entry((new_vote_slot, UpdateLabel::Add)) + .and_modify(|update| update.update_stake(stake_update)) + .or_insert(UpdateOperation::Add(stake_update)); + self.insert_aggregate_operations(&mut update_operations, new_vote_slot); + } + + update_operations + } + + fn process_update_operations( + &mut self, + update_operations: BTreeMap<(Slot, UpdateLabel), UpdateOperation>, + ) { + // Iterate through the update operations from greatest to smallest slot + for ((slot, _), operation) in update_operations.into_iter().rev() { + match operation { + UpdateOperation::Aggregate => self.aggregate_slot(slot), + UpdateOperation::Add(stake) => self.add_slot_stake(slot, stake), + UpdateOperation::Subtract(stake) => self.subtract_slot_stake(slot, stake), + } + } + } + + fn add_slot_stake(&mut self, slot: Slot, stake: u64) { + if let Some(fork_info) = self.fork_infos.get_mut(&slot) { + fork_info.stake_voted_at += stake; + fork_info.stake_voted_subtree += stake; + } + } + + fn subtract_slot_stake(&mut self, slot: Slot, stake: u64) { + if let Some(fork_info) = self.fork_infos.get_mut(&slot) { + fork_info.stake_voted_at -= stake; + fork_info.stake_voted_subtree -= stake; + } + } + + fn stake_voted_at(&self, slot: Slot) -> Option { + self.fork_infos + .get(&slot) + .map(|fork_info| fork_info.stake_voted_at) + } + + #[cfg(test)] + fn set_stake_voted_at(&mut self, slot: Slot, stake_voted_at: u64) { + self.fork_infos.get_mut(&slot).unwrap().stake_voted_at = stake_voted_at; + } + + #[cfg(test)] + fn is_leaf(&self, slot: Slot) -> bool { + self.fork_infos.get(&slot).unwrap().children.is_empty() + } + + #[cfg(test)] + fn children(&self, slot: Slot) -> Option<&Vec> { + self.fork_infos + .get(&slot) + .map(|fork_info| &fork_info.children) + } + + #[cfg(test)] + fn parent(&self, slot: Slot) -> Option { + self.fork_infos + .get(&slot) + .map(|fork_info| fork_info.parent) + .unwrap_or(None) + } +} + +struct AncestorIterator<'a> { + current_slot: Slot, + fork_infos: &'a HashMap, +} + +impl<'a> AncestorIterator<'a> { + fn new(start_slot: Slot, fork_infos: &'a HashMap) -> Self { + Self { + current_slot: start_slot, + fork_infos, + } + } +} + +impl<'a> Iterator for AncestorIterator<'a> { + type Item = Slot; + fn next(&mut self) -> Option { + let parent = self + .fork_infos + .get(&self.current_slot) + .map(|fork_info| fork_info.parent) + .unwrap_or(None); + + parent + .map(|parent| { + self.current_slot = parent; + Some(self.current_slot) + }) + .unwrap_or(None) + } +} + +#[cfg(test)] +mod test { + use super::*; + use crate::consensus::test::VoteSimulator; + use solana_runtime::{ + bank::Bank, + genesis_utils::{self, GenesisConfigInfo, ValidatorVoteKeypairs}, + }; + use solana_sdk::signature::{Keypair, Signer}; + use std::{collections::HashSet, ops::Range}; + use trees::tr; + + #[test] + fn test_ancestor_iterator() { + let mut heaviest_subtree_fork_choice = setup_forks(); + let parents: Vec<_> = heaviest_subtree_fork_choice.ancestor_iterator(6).collect(); + assert_eq!(parents, vec![5, 3, 1, 0]); + let parents: Vec<_> = heaviest_subtree_fork_choice.ancestor_iterator(4).collect(); + assert_eq!(parents, vec![2, 1, 0]); + let parents: Vec<_> = heaviest_subtree_fork_choice.ancestor_iterator(1).collect(); + assert_eq!(parents, vec![0]); + let parents: Vec<_> = heaviest_subtree_fork_choice.ancestor_iterator(0).collect(); + assert!(parents.is_empty()); + + // Set a root, everything but slots 2, 4 should be removed + heaviest_subtree_fork_choice.set_root(2); + let parents: Vec<_> = heaviest_subtree_fork_choice.ancestor_iterator(4).collect(); + assert_eq!(parents, vec![2]); + } + + #[test] + fn test_new_from_frozen_banks() { + /* + Build fork structure: + slot 0 + | + slot 1 + / \ + slot 2 | + | slot 3 + slot 4 + */ + let forks = tr(0) / (tr(1) / (tr(2) / (tr(4))) / (tr(3))); + let mut vote_simulator = VoteSimulator::new(1); + vote_simulator.fill_bank_forks(forks, &HashMap::new()); + let bank_forks = vote_simulator.bank_forks; + let mut frozen_banks: Vec<_> = bank_forks + .read() + .unwrap() + .frozen_banks() + .values() + .cloned() + .collect(); + frozen_banks.sort_by_key(|bank| bank.slot()); + + let root = bank_forks.read().unwrap().root(); + let heaviest_subtree_fork_choice = + HeaviestSubtreeForkChoice::new_from_frozen_banks(root, &frozen_banks); + assert!(heaviest_subtree_fork_choice.parent(0).is_none()); + assert_eq!(heaviest_subtree_fork_choice.children(0).unwrap(), &[1]); + assert_eq!(heaviest_subtree_fork_choice.parent(1), Some(0)); + assert_eq!(heaviest_subtree_fork_choice.children(1).unwrap(), &[2, 3]); + assert_eq!(heaviest_subtree_fork_choice.parent(2), Some(1)); + assert_eq!(heaviest_subtree_fork_choice.children(2).unwrap(), &[4]); + assert_eq!(heaviest_subtree_fork_choice.parent(3), Some(1)); + assert!(heaviest_subtree_fork_choice.children(3).unwrap().is_empty()); + assert_eq!(heaviest_subtree_fork_choice.parent(4), Some(2)); + assert!(heaviest_subtree_fork_choice.children(4).unwrap().is_empty()); + } + + #[test] + fn test_set_root() { + let mut heaviest_subtree_fork_choice = setup_forks(); + + // Set root to 1, should only purge 0 + heaviest_subtree_fork_choice.set_root(1); + for i in 0..=6 { + let exists = i != 0; + assert_eq!( + heaviest_subtree_fork_choice.fork_infos.contains_key(&i), + exists + ); + } + + // Set root to 5, should purge everything except 5, 6 + heaviest_subtree_fork_choice.set_root(5); + for i in 0..=6 { + let exists = i == 5 || i == 6; + assert_eq!( + heaviest_subtree_fork_choice.fork_infos.contains_key(&i), + exists + ); + } + } + + #[test] + fn test_propagate_new_leaf() { + let mut heaviest_subtree_fork_choice = setup_forks(); + let ancestors = heaviest_subtree_fork_choice + .ancestor_iterator(6) + .collect::>(); + for a in ancestors.into_iter().chain(std::iter::once(6)) { + assert_eq!(heaviest_subtree_fork_choice.best_slot(a).unwrap(), 6); + } + + // Add a leaf 10, it should be the best choice + heaviest_subtree_fork_choice.add_new_leaf_slot(10, Some(6)); + let ancestors = heaviest_subtree_fork_choice + .ancestor_iterator(10) + .collect::>(); + for a in ancestors.into_iter().chain(std::iter::once(10)) { + assert_eq!(heaviest_subtree_fork_choice.best_slot(a).unwrap(), 10); + } + + // Add a smaller leaf 9, it should be the best choice + heaviest_subtree_fork_choice.add_new_leaf_slot(9, Some(6)); + let ancestors = heaviest_subtree_fork_choice + .ancestor_iterator(9) + .collect::>(); + for a in ancestors.into_iter().chain(std::iter::once(9)) { + assert_eq!(heaviest_subtree_fork_choice.best_slot(a).unwrap(), 9); + } + + // Add a higher leaf 11, should not change the best choice + heaviest_subtree_fork_choice.add_new_leaf_slot(11, Some(6)); + let ancestors = heaviest_subtree_fork_choice + .ancestor_iterator(11) + .collect::>(); + for a in ancestors.into_iter().chain(std::iter::once(9)) { + assert_eq!(heaviest_subtree_fork_choice.best_slot(a).unwrap(), 9); + } + + // If an earlier ancestor than the direct parent already has another `best_slot`, + // it shouldn't change. + let stake = 100; + let (bank, vote_pubkeys) = setup_bank_and_vote_pubkeys(3, stake); + let other_best_leaf = 4; + // Leaf slot 9 stops being the `best_slot` at slot 1 b/c there are votes + // for slot 2 + for pubkey in &vote_pubkeys[0..1] { + heaviest_subtree_fork_choice.add_votes( + &[(*pubkey, other_best_leaf)], + bank.epoch_stakes_map(), + bank.epoch_schedule(), + ); + } + heaviest_subtree_fork_choice.add_new_leaf_slot(8, Some(6)); + let ancestors = heaviest_subtree_fork_choice + .ancestor_iterator(8) + .collect::>(); + for a in ancestors.into_iter().chain(std::iter::once(8)) { + let best_slot = if a > 1 { 8 } else { other_best_leaf }; + assert_eq!( + heaviest_subtree_fork_choice.best_slot(a).unwrap(), + best_slot + ); + } + + // If the direct parent's `best_slot` has non-zero stake voting + // for it, then the `best_slot` should not change, even with a lower + // leaf being added + assert_eq!(heaviest_subtree_fork_choice.best_slot(6).unwrap(), 8); + // Add a vote for slot 8 + heaviest_subtree_fork_choice.add_votes( + &[(vote_pubkeys[2], 8)], + bank.epoch_stakes_map(), + bank.epoch_schedule(), + ); + heaviest_subtree_fork_choice.add_new_leaf_slot(7, Some(6)); + let ancestors = heaviest_subtree_fork_choice + .ancestor_iterator(7) + .collect::>(); + // best_slot's remain unchanged + for a in ancestors.into_iter().chain(std::iter::once(8)) { + let best_slot = if a > 1 { 8 } else { other_best_leaf }; + assert_eq!( + heaviest_subtree_fork_choice.best_slot(a).unwrap(), + best_slot + ); + } + + // All the leaves should think they are their own best choice + for leaf in [8, 9, 10, 11].iter() { + assert_eq!( + heaviest_subtree_fork_choice.best_slot(*leaf).unwrap(), + *leaf + ); + } + } + + #[test] + fn test_propagate_new_leaf_2() { + /* + Build fork structure: + slot 0 + | + slot 4 + | + slot 5 + */ + let forks = tr(0) / (tr(4) / (tr(5))); + let mut heaviest_subtree_fork_choice = HeaviestSubtreeForkChoice::new_from_tree(forks); + + let stake = 100; + let (bank, vote_pubkeys) = setup_bank_and_vote_pubkeys(1, stake); + + // slot 5 should be the best because it's the only leaef + assert_eq!(heaviest_subtree_fork_choice.best_overall_slot(), 5); + + // Add a leaf slot 2 on a different fork than slot 5. Slot 2 should + // be the new best because it's for a lesser slot + heaviest_subtree_fork_choice.add_new_leaf_slot(2, Some(0)); + assert_eq!(heaviest_subtree_fork_choice.best_overall_slot(), 2); + + // Add a vote for slot 4 + heaviest_subtree_fork_choice.add_votes( + &[(vote_pubkeys[0], 4)], + bank.epoch_stakes_map(), + bank.epoch_schedule(), + ); + + // slot 5 should be the best again + assert_eq!(heaviest_subtree_fork_choice.best_overall_slot(), 5); + + // Adding a slot 1 that is less than the current best slot 5 should not change the best + // slot because the fork slot 5 is on has a higher weight + heaviest_subtree_fork_choice.add_new_leaf_slot(1, Some(0)); + assert_eq!(heaviest_subtree_fork_choice.best_overall_slot(), 5); + } + + #[test] + fn test_aggregate_slot() { + let mut heaviest_subtree_fork_choice = setup_forks(); + + // No weights are present, weights should be zero + heaviest_subtree_fork_choice.aggregate_slot(1); + assert_eq!(heaviest_subtree_fork_choice.stake_voted_at(1).unwrap(), 0); + assert_eq!( + heaviest_subtree_fork_choice.stake_voted_subtree(1).unwrap(), + 0 + ); + // The best leaf when weights are equal should prioritize the lower leaf + assert_eq!(heaviest_subtree_fork_choice.best_slot(3).unwrap(), 6); + assert_eq!(heaviest_subtree_fork_choice.best_slot(2).unwrap(), 4); + assert_eq!(heaviest_subtree_fork_choice.best_slot(1).unwrap(), 4); + + // Update the weights that have voted *exactly* at each slot + heaviest_subtree_fork_choice.set_stake_voted_at(6, 3); + heaviest_subtree_fork_choice.set_stake_voted_at(5, 3); + heaviest_subtree_fork_choice.set_stake_voted_at(2, 4); + heaviest_subtree_fork_choice.set_stake_voted_at(4, 2); + let total_stake = 12; + + // Aggregate up each off the two forks (order matters, has to be + // reverse order for each fork, and aggregating a slot multiple times + // is fine) + let slots_to_aggregate: Vec<_> = std::iter::once(6) + .chain(heaviest_subtree_fork_choice.ancestor_iterator(6)) + .chain(std::iter::once(4)) + .chain(heaviest_subtree_fork_choice.ancestor_iterator(4)) + .collect(); + + for slot in slots_to_aggregate { + heaviest_subtree_fork_choice.aggregate_slot(slot); + } + + for slot in 0..=1 { + // No stake has voted exactly at slots 0 or 1 + assert_eq!( + heaviest_subtree_fork_choice.stake_voted_at(slot).unwrap(), + 0 + ); + // Subtree stake is sum of thee `stake_voted_at` across + // all slots in the subtree + assert_eq!( + heaviest_subtree_fork_choice + .stake_voted_subtree(slot) + .unwrap(), + total_stake + ); + // The best path is 0 -> 1 -> 2 -> 4, so slot 4 should be the best choice + assert_eq!(heaviest_subtree_fork_choice.best_slot(slot).unwrap(), 4); + } + } + + #[test] + fn test_process_update_operations() { + let mut heaviest_subtree_fork_choice = setup_forks(); + let stake = 100; + let (bank, vote_pubkeys) = setup_bank_and_vote_pubkeys(3, stake); + + let pubkey_votes: Vec<(Pubkey, Slot)> = vec![ + (vote_pubkeys[0], 3), + (vote_pubkeys[1], 2), + (vote_pubkeys[2], 1), + ]; + let expected_best_slot = + |slot, heaviest_subtree_fork_choice: &HeaviestSubtreeForkChoice| -> Slot { + if !heaviest_subtree_fork_choice.is_leaf(slot) { + // Both branches have equal weight, so should pick the lesser leaf + if heaviest_subtree_fork_choice + .ancestor_iterator(4) + .collect::>() + .contains(&slot) + { + 4 + } else { + 6 + } + } else { + slot + } + }; + + check_process_update_correctness( + &mut heaviest_subtree_fork_choice, + &pubkey_votes, + 0..7, + &bank, + stake, + expected_best_slot, + ); + + // Everyone makes newer votes + let pubkey_votes: Vec<(Pubkey, Slot)> = vec![ + (vote_pubkeys[0], 4), + (vote_pubkeys[1], 3), + (vote_pubkeys[2], 3), + ]; + + let expected_best_slot = + |slot, heaviest_subtree_fork_choice: &HeaviestSubtreeForkChoice| -> Slot { + if !heaviest_subtree_fork_choice.is_leaf(slot) { + // The branch with leaf 6 now has two votes, so should pick that one + if heaviest_subtree_fork_choice + .ancestor_iterator(6) + .collect::>() + .contains(&slot) + { + 6 + } else { + 4 + } + } else { + slot + } + }; + + check_process_update_correctness( + &mut heaviest_subtree_fork_choice, + &pubkey_votes, + 0..7, + &bank, + stake, + expected_best_slot, + ); + } + + #[test] + fn test_generate_update_operations() { + let mut heaviest_subtree_fork_choice = setup_forks(); + let stake = 100; + let (bank, vote_pubkeys) = setup_bank_and_vote_pubkeys(3, stake); + let pubkey_votes: Vec<(Pubkey, Slot)> = vec![ + (vote_pubkeys[0], 3), + (vote_pubkeys[1], 4), + (vote_pubkeys[2], 1), + ]; + + let expected_update_operations: BTreeMap<(Slot, UpdateLabel), UpdateOperation> = vec![ + // Add/remove from new/old forks + ((1, UpdateLabel::Add), UpdateOperation::Add(stake)), + ((3, UpdateLabel::Add), UpdateOperation::Add(stake)), + ((4, UpdateLabel::Add), UpdateOperation::Add(stake)), + // Aggregate all ancestors of changed slots + ((0, UpdateLabel::Aggregate), UpdateOperation::Aggregate), + ((1, UpdateLabel::Aggregate), UpdateOperation::Aggregate), + ((2, UpdateLabel::Aggregate), UpdateOperation::Aggregate), + ] + .into_iter() + .collect(); + + let generated_update_operations = heaviest_subtree_fork_choice.generate_update_operations( + &pubkey_votes, + bank.epoch_stakes_map(), + bank.epoch_schedule(), + ); + assert_eq!(expected_update_operations, generated_update_operations); + + // Everyone makes older/same votes, should be ignored + let pubkey_votes: Vec<(Pubkey, Slot)> = vec![ + (vote_pubkeys[0], 3), + (vote_pubkeys[1], 2), + (vote_pubkeys[2], 1), + ]; + let generated_update_operations = heaviest_subtree_fork_choice.generate_update_operations( + &pubkey_votes, + bank.epoch_stakes_map(), + bank.epoch_schedule(), + ); + assert!(generated_update_operations.is_empty()); + + // Some people make newer votes + let pubkey_votes: Vec<(Pubkey, Slot)> = vec![ + // old, ignored + (vote_pubkeys[0], 3), + // new, switched forks + (vote_pubkeys[1], 5), + // new, same fork + (vote_pubkeys[2], 3), + ]; + + let expected_update_operations: BTreeMap<(Slot, UpdateLabel), UpdateOperation> = vec![ + // Add/remove to/from new/old forks + ((3, UpdateLabel::Add), UpdateOperation::Add(stake)), + ((5, UpdateLabel::Add), UpdateOperation::Add(stake)), + ((1, UpdateLabel::Subtract), UpdateOperation::Subtract(stake)), + ((4, UpdateLabel::Subtract), UpdateOperation::Subtract(stake)), + // Aggregate all ancestors of changed slots + ((0, UpdateLabel::Aggregate), UpdateOperation::Aggregate), + ((1, UpdateLabel::Aggregate), UpdateOperation::Aggregate), + ((2, UpdateLabel::Aggregate), UpdateOperation::Aggregate), + ((3, UpdateLabel::Aggregate), UpdateOperation::Aggregate), + ] + .into_iter() + .collect(); + + let generated_update_operations = heaviest_subtree_fork_choice.generate_update_operations( + &pubkey_votes, + bank.epoch_stakes_map(), + bank.epoch_schedule(), + ); + assert_eq!(expected_update_operations, generated_update_operations); + + // People make new votes + let pubkey_votes: Vec<(Pubkey, Slot)> = vec![ + // new, switch forks + (vote_pubkeys[0], 4), + // new, same fork + (vote_pubkeys[1], 6), + // new, same fork + (vote_pubkeys[2], 6), + ]; + + let expected_update_operations: BTreeMap<(Slot, UpdateLabel), UpdateOperation> = vec![ + // Add/remove from new/old forks + ((4, UpdateLabel::Add), UpdateOperation::Add(stake)), + ((6, UpdateLabel::Add), UpdateOperation::Add(2 * stake)), + ( + (3, UpdateLabel::Subtract), + UpdateOperation::Subtract(2 * stake), + ), + ((5, UpdateLabel::Subtract), UpdateOperation::Subtract(stake)), + // Aggregate all ancestors of changed slots + ((0, UpdateLabel::Aggregate), UpdateOperation::Aggregate), + ((1, UpdateLabel::Aggregate), UpdateOperation::Aggregate), + ((2, UpdateLabel::Aggregate), UpdateOperation::Aggregate), + ((3, UpdateLabel::Aggregate), UpdateOperation::Aggregate), + ((5, UpdateLabel::Aggregate), UpdateOperation::Aggregate), + ] + .into_iter() + .collect(); + + let generated_update_operations = heaviest_subtree_fork_choice.generate_update_operations( + &pubkey_votes, + bank.epoch_stakes_map(), + bank.epoch_schedule(), + ); + assert_eq!(expected_update_operations, generated_update_operations); + } + + #[test] + fn test_add_votes() { + let mut heaviest_subtree_fork_choice = setup_forks(); + let stake = 100; + let (bank, vote_pubkeys) = setup_bank_and_vote_pubkeys(3, stake); + + let pubkey_votes: Vec<(Pubkey, Slot)> = vec![ + (vote_pubkeys[0], 3), + (vote_pubkeys[1], 2), + (vote_pubkeys[2], 1), + ]; + assert_eq!( + heaviest_subtree_fork_choice.add_votes( + &pubkey_votes, + bank.epoch_stakes_map(), + bank.epoch_schedule() + ), + 4 + ); + + assert_eq!(heaviest_subtree_fork_choice.best_overall_slot(), 4) + } + + fn setup_forks() -> HeaviestSubtreeForkChoice { + /* + Build fork structure: + slot 0 + | + slot 1 + / \ + slot 2 | + | slot 3 + slot 4 | + slot 5 + | + slot 6 + */ + let forks = tr(0) / (tr(1) / (tr(2) / (tr(4))) / (tr(3) / (tr(5) / (tr(6))))); + HeaviestSubtreeForkChoice::new_from_tree(forks) + } + + fn setup_bank_and_vote_pubkeys(num_vote_accounts: usize, stake: u64) -> (Bank, Vec) { + // Create some voters at genesis + let validator_voting_keypairs: Vec<_> = (0..num_vote_accounts) + .map(|_| ValidatorVoteKeypairs::new(Keypair::new(), Keypair::new(), Keypair::new())) + .collect(); + + let vote_pubkeys: Vec<_> = validator_voting_keypairs + .iter() + .map(|k| k.vote_keypair.pubkey()) + .collect(); + let GenesisConfigInfo { genesis_config, .. } = + genesis_utils::create_genesis_config_with_vote_accounts( + 10_000, + &validator_voting_keypairs, + stake, + ); + let bank = Bank::new(&genesis_config); + (bank, vote_pubkeys) + } + + fn check_process_update_correctness( + heaviest_subtree_fork_choice: &mut HeaviestSubtreeForkChoice, + pubkey_votes: &[(Pubkey, Slot)], + slots_range: Range, + bank: &Bank, + stake: u64, + mut expected_best_slot: F, + ) where + F: FnMut(Slot, &HeaviestSubtreeForkChoice) -> Slot, + { + let unique_votes: HashSet = pubkey_votes + .iter() + .map(|(_, vote_slot)| *vote_slot) + .collect(); + let vote_ancestors: HashMap> = unique_votes + .iter() + .map(|v| { + ( + *v, + heaviest_subtree_fork_choice.ancestor_iterator(*v).collect(), + ) + }) + .collect(); + let mut vote_count: HashMap = HashMap::new(); + for (_, vote) in pubkey_votes { + vote_count.entry(*vote).and_modify(|c| *c += 1).or_insert(1); + } + + // Maps a slot to the number of descendants of that slot + // that have been voted on + let num_voted_descendants: HashMap = slots_range + .clone() + .map(|slot| { + let num_voted_descendants = vote_ancestors + .iter() + .map(|(vote_slot, ancestors)| { + (ancestors.contains(&slot) || *vote_slot == slot) as usize + * vote_count.get(vote_slot).unwrap() + }) + .sum(); + (slot, num_voted_descendants) + }) + .collect(); + + let update_operations = heaviest_subtree_fork_choice.generate_update_operations( + &pubkey_votes, + bank.epoch_stakes_map(), + bank.epoch_schedule(), + ); + + heaviest_subtree_fork_choice.process_update_operations(update_operations); + for slot in slots_range { + let expected_stake_voted_at = + vote_count.get(&slot).cloned().unwrap_or(0) as u64 * stake; + let expected_stake_voted_subtree = + *num_voted_descendants.get(&slot).unwrap() as u64 * stake; + assert_eq!( + expected_stake_voted_at, + heaviest_subtree_fork_choice.stake_voted_at(slot).unwrap() + ); + assert_eq!( + expected_stake_voted_subtree, + heaviest_subtree_fork_choice + .stake_voted_subtree(slot) + .unwrap() + ); + assert_eq!( + expected_best_slot(slot, heaviest_subtree_fork_choice), + heaviest_subtree_fork_choice.best_slot(slot).unwrap() + ); + } + } +} diff --git a/core/src/lib.rs b/core/src/lib.rs index 9f1016ac16..e67e33be88 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -15,6 +15,7 @@ mod deprecated; pub mod shred_fetch_stage; #[macro_use] pub mod contact_info; +pub mod bank_weight_fork_choice; pub mod cluster_info; pub mod cluster_slots; pub mod consensus; @@ -26,8 +27,10 @@ pub mod crds_gossip_push; pub mod crds_value; pub mod epoch_slots; pub mod fetch_stage; +pub mod fork_choice; pub mod gen_keys; pub mod gossip_service; +pub mod heaviest_subtree_fork_choice; pub mod ledger_cleanup_service; pub mod local_vote_signer_service; pub mod non_circulating_supply; diff --git a/core/src/replay_stage.rs b/core/src/replay_stage.rs index c4063d6b7a..27f2c14b5d 100644 --- a/core/src/replay_stage.rs +++ b/core/src/replay_stage.rs @@ -1,14 +1,17 @@ //! The `replay_stage` replays transactions broadcast by the leader. use crate::{ + bank_weight_fork_choice::BankWeightForkChoice, broadcast_stage::RetransmitSlotsSender, cluster_info::ClusterInfo, cluster_info_vote_listener::VoteTracker, cluster_slots::ClusterSlots, commitment::{AggregateCommitmentService, BlockCommitmentCache, CommitmentAggregationData}, - consensus::{StakeLockout, SwitchForkDecision, Tower}, + consensus::{ComputedBankState, StakeLockout, SwitchForkDecision, Tower}, + fork_choice::{ForkChoice, SelectVoteAndResetForkResult}, + heaviest_subtree_fork_choice::HeaviestSubtreeForkChoice, poh_recorder::{PohRecorder, GRACE_TICKS_FACTOR, MAX_GRACE_SLOTS}, - progress_map::{ForkProgress, ForkStats, ProgressMap, PropagatedStats}, + progress_map::{ForkProgress, ProgressMap, PropagatedStats}, pubkey_references::PubkeyReferences, repair_service::DuplicateSlotsResetReceiver, result::Result, @@ -33,7 +36,7 @@ use solana_sdk::{ hash::Hash, pubkey::Pubkey, signature::{Keypair, Signer}, - timing::{self, duration_as_ms}, + timing::duration_as_ms, transaction::Transaction, }; use solana_vote_program::{ @@ -140,12 +143,6 @@ impl ReplayTiming { } } -pub(crate) struct SelectVoteAndResetForkResult { - pub vote_bank: Option<(Arc, SwitchForkDecision)>, - pub reset_bank: Option>, - pub heaviest_fork_failures: Vec, -} - pub struct ReplayStage { t_replay: JoinHandle>, commitment_service: AggregateCommitmentService, @@ -180,8 +177,6 @@ impl ReplayStage { } = config; trace!("replay stage"); - let mut tower = Tower::new(&my_pubkey, &vote_account, &bank_forks.read().unwrap()); - // Start the replay stage loop let (lockouts_sender, commitment_service) = AggregateCommitmentService::new( &exit, @@ -208,12 +203,12 @@ impl ReplayStage { frozen_banks.sort_by_key(|bank| bank.slot()); // Initialize progress map with any root banks - for bank in frozen_banks { - let prev_leader_slot = progress.get_bank_prev_leader_slot(&bank); + for bank in &frozen_banks { + let prev_leader_slot = progress.get_bank_prev_leader_slot(bank); progress.insert( bank.slot(), ForkProgress::new_from_bank( - &bank, + bank, &my_pubkey, &vote_account, prev_leader_slot, @@ -222,6 +217,27 @@ impl ReplayStage { ), ); } + let root_bank = bank_forks.read().unwrap().root_bank().clone(); + let root = root_bank.slot(); + let unlock_heaviest_subtree_fork_choice_slot = + Self::get_unlock_heaviest_subtree_fork_choice(root_bank.operating_mode()); + let mut heaviest_subtree_fork_choice = + HeaviestSubtreeForkChoice::new_from_frozen_banks(root, &frozen_banks); + let mut bank_weight_fork_choice = BankWeightForkChoice::default(); + let heaviest_bank = if root > unlock_heaviest_subtree_fork_choice_slot { + bank_forks + .read() + .unwrap() + .get(heaviest_subtree_fork_choice.best_overall_slot()) + .expect( + "The best overall slot must be one of `frozen_banks` which all + exist in bank_forks", + ) + .clone() + } else { + Tower::find_heaviest_bank(&bank_forks, &my_pubkey).unwrap_or(root_bank) + }; + let mut tower = Tower::new(&my_pubkey, &vote_account, root, &heaviest_bank); let mut current_leader = None; let mut last_reset = Hash::default(); let mut partition = false; @@ -259,6 +275,7 @@ impl ReplayStage { &mut progress, transaction_status_sender.clone(), &verify_recyclers, + &mut heaviest_subtree_fork_choice, ); Self::report_memory(&allocated, "replay_active_banks", start); @@ -296,6 +313,8 @@ impl ReplayStage { &cluster_slots, &bank_forks, &mut all_pubkeys, + &mut heaviest_subtree_fork_choice, + &mut bank_weight_fork_choice, ); let compute_bank_stats_elapsed = now.elapsed().as_micros(); for slot in newly_computed_slot_stats { @@ -317,8 +336,14 @@ impl ReplayStage { } } - let (heaviest_bank, heaviest_bank_on_same_fork) = - Self::select_forks(&frozen_banks, &tower, &progress, &ancestors); + let fork_choice: &mut dyn ForkChoice = + if forks_root > unlock_heaviest_subtree_fork_choice_slot { + &mut heaviest_subtree_fork_choice + } else { + &mut bank_weight_fork_choice + }; + let (heaviest_bank, heaviest_bank_on_same_voted_fork) = fork_choice + .select_forks(&frozen_banks, &tower, &progress, &ancestors, &bank_forks); Self::report_memory(&allocated, "select_fork", start); @@ -329,7 +354,7 @@ impl ReplayStage { heaviest_fork_failures, } = Self::select_vote_and_reset_forks( &heaviest_bank, - &heaviest_bank_on_same_fork, + &heaviest_bank_on_same_voted_fork, &ancestors, &descendants, &progress, @@ -341,13 +366,10 @@ impl ReplayStage { select_vote_and_reset_forks_elapsed as u64, ); - if heaviest_bank.is_some() - && tower.is_recent(heaviest_bank.as_ref().unwrap().slot()) - && !heaviest_fork_failures.is_empty() - { + if tower.is_recent(heaviest_bank.slot()) && !heaviest_fork_failures.is_empty() { info!( "Couldn't vote on heaviest fork: {:?}, heaviest_fork_failures: {:?}", - heaviest_bank.as_ref().map(|b| b.slot()), + heaviest_bank.slot(), heaviest_fork_failures ); @@ -394,6 +416,7 @@ impl ReplayStage { &mut all_pubkeys, &subscriptions, &block_commitment_cache, + &mut heaviest_subtree_fork_choice, )?; }; @@ -872,6 +895,7 @@ impl ReplayStage { all_pubkeys: &mut PubkeyReferences, subscriptions: &Arc, block_commitment_cache: &Arc>, + heaviest_subtree_fork_choice: &mut HeaviestSubtreeForkChoice, ) -> Result<()> { if bank.is_empty() { inc_new_counter_info!("replay_stage-voted_empty_bank", 1); @@ -910,6 +934,7 @@ impl ReplayStage { accounts_hash_sender, all_pubkeys, largest_confirmed_root, + heaviest_subtree_fork_choice, ); subscriptions.notify_roots(rooted_slots); latest_root_senders.iter().for_each(|s| { @@ -1076,6 +1101,7 @@ impl ReplayStage { progress: &mut ProgressMap, transaction_status_sender: Option, verify_recyclers: &VerifyRecyclers, + heaviest_subtree_fork_choice: &mut HeaviestSubtreeForkChoice, ) -> bool { let mut did_complete_bank = false; let mut tx_count = 0; @@ -1143,6 +1169,8 @@ impl ReplayStage { did_complete_bank = true; info!("bank frozen: {}", bank.slot()); bank.freeze(); + heaviest_subtree_fork_choice + .add_new_leaf_slot(bank.slot(), Some(bank.parent_slot())); } else { trace!( "bank {} not completed tick_height: {}, max_tick_height: {}", @@ -1156,6 +1184,7 @@ impl ReplayStage { did_complete_bank } + #[allow(clippy::too_many_arguments)] pub(crate) fn compute_bank_stats( my_pubkey: &Pubkey, ancestors: &HashMap>, @@ -1166,37 +1195,56 @@ impl ReplayStage { cluster_slots: &ClusterSlots, bank_forks: &RwLock, all_pubkeys: &mut PubkeyReferences, + heaviest_subtree_fork_choice: &mut dyn ForkChoice, + bank_weight_fork_choice: &mut dyn ForkChoice, ) -> Vec { frozen_banks.sort_by_key(|bank| bank.slot()); let mut new_stats = vec![]; for bank in frozen_banks { let bank_slot = bank.slot(); - // Only time progress map should be missing a bank slot // is if this node was the leader for this slot as those banks // are not replayed in replay_active_banks() - let parent_weight = bank - .parent() - .and_then(|b| progress.get(&b.slot())) - .map(|x| x.fork_stats.fork_weight) - .unwrap_or(0); { - let stats = progress + let is_computed = progress .get_fork_stats_mut(bank_slot) - .expect("All frozen banks must exist in the Progress map"); - - if !stats.computed { - let (stake_lockouts, total_staked, bank_weight, lockout_intervals) = tower - .collect_vote_lockouts( - bank_slot, - bank.vote_accounts().into_iter(), - &ancestors, - all_pubkeys, - ); + .expect("All frozen banks must exist in the Progress map") + .computed; + if !is_computed { + let computed_bank_state = Tower::collect_vote_lockouts( + my_pubkey, + bank_slot, + bank.vote_accounts().into_iter(), + &ancestors, + all_pubkeys, + ); + heaviest_subtree_fork_choice.compute_bank_stats( + &bank, + tower, + progress, + &computed_bank_state, + ); + bank_weight_fork_choice.compute_bank_stats( + &bank, + tower, + progress, + &computed_bank_state, + ); + let ComputedBankState { + stake_lockouts, + total_staked, + lockout_intervals, + .. + } = computed_bank_state; + let stats = progress + .get_fork_stats_mut(bank_slot) + .expect("All frozen banks must exist in the Progress map"); stats.total_staked = total_staked; - stats.weight = bank_weight; - stats.fork_weight = stats.weight + parent_weight; - + stats.stake_lockouts = stake_lockouts; + stats.lockout_intervals = lockout_intervals; + stats.block_height = bank.block_height(); + stats.computed = true; + new_stats.push(bank_slot); datapoint_info!( "bank_weight", ("slot", bank_slot, i64), @@ -1211,11 +1259,6 @@ impl ReplayStage { stats.fork_weight, bank.parent().map(|b| b.slot()).unwrap_or(0) ); - stats.stake_lockouts = stake_lockouts; - stats.lockout_intervals = lockout_intervals; - stats.block_height = bank.block_height(); - stats.computed = true; - new_stats.push(bank_slot); } } @@ -1289,11 +1332,11 @@ impl ReplayStage { let newly_voted_pubkeys = slot_vote_tracker .as_ref() .and_then(|slot_vote_tracker| slot_vote_tracker.write().unwrap().get_updates()) - .unwrap_or_else(Vec::new); + .unwrap_or_default(); let cluster_slot_pubkeys = cluster_slot_pubkeys .map(|v| v.read().unwrap().keys().cloned().collect()) - .unwrap_or_else(Vec::new); + .unwrap_or_default(); Self::update_fork_propagated_threshold_from_votes( progress, @@ -1305,115 +1348,12 @@ impl ReplayStage { ); } - // Returns: - // 1) The heaviest bank - // 2) The latest votable bank on the same fork as the last vote - pub(crate) fn select_forks( - frozen_banks: &[Arc], - tower: &Tower, - progress: &ProgressMap, - ancestors: &HashMap>, - ) -> (Option>, Option>) { - let tower_start = Instant::now(); - let num_frozen_banks = frozen_banks.len(); - - trace!("frozen_banks {}", frozen_banks.len()); - let num_old_banks = frozen_banks - .iter() - .filter(|b| b.slot() < tower.root().unwrap_or(0)) - .count(); - - let last_vote = tower.last_vote().slots.last().cloned(); - let mut heaviest_bank_on_same_fork = None; - let mut heaviest_same_fork_weight = 0; - let stats: Vec<&ForkStats> = frozen_banks - .iter() - .map(|bank| { - // Only time progress map should be missing a bank slot - // is if this node was the leader for this slot as those banks - // are not replayed in replay_active_banks() - let stats = progress - .get_fork_stats(bank.slot()) - .expect("All frozen banks must exist in the Progress map"); - - if let Some(last_vote) = last_vote { - if ancestors - .get(&bank.slot()) - .expect("Entry in frozen banks must exist in ancestors") - .contains(&last_vote) - { - // Descendant of last vote cannot be locked out - assert!(!stats.is_locked_out); - - // ancestors(slot) should not contain the slot itself, - // so we should never get the same bank as the last vote - assert_ne!(bank.slot(), last_vote); - // highest weight, lowest slot first. frozen_banks is sorted - // from least slot to greatest slot, so if two banks have - // the same fork weight, the lower slot will be picked - if stats.fork_weight > heaviest_same_fork_weight { - heaviest_bank_on_same_fork = Some(bank.clone()); - heaviest_same_fork_weight = stats.fork_weight; - } - } - } - - stats - }) - .collect(); - let num_not_recent = stats.iter().filter(|s| !s.is_recent).count(); - let num_has_voted = stats.iter().filter(|s| s.has_voted).count(); - let num_empty = stats.iter().filter(|s| s.is_empty).count(); - let num_threshold_failure = stats.iter().filter(|s| !s.vote_threshold).count(); - let num_votable_threshold_failure = stats - .iter() - .filter(|s| s.is_recent && !s.has_voted && !s.vote_threshold) - .count(); - - let mut candidates: Vec<_> = frozen_banks.iter().zip(stats.iter()).collect(); - - //highest weight, lowest slot first - candidates.sort_by_key(|b| (b.1.fork_weight, 0i64 - b.0.slot() as i64)); - let rv = candidates.last(); - let ms = timing::duration_as_ms(&tower_start.elapsed()); - let weights: Vec<(u128, u64, u64)> = candidates - .iter() - .map(|x| (x.1.weight, x.0.slot(), x.1.block_height)) - .collect(); - debug!( - "@{:?} tower duration: {:?} len: {}/{} weights: {:?} voting: {}", - timing::timestamp(), - ms, - candidates.len(), - stats.iter().filter(|s| !s.has_voted).count(), - weights, - rv.is_some() - ); - datapoint_debug!( - "replay_stage-select_forks", - ("frozen_banks", num_frozen_banks as i64, i64), - ("not_recent", num_not_recent as i64, i64), - ("has_voted", num_has_voted as i64, i64), - ("old_banks", num_old_banks as i64, i64), - ("empty_banks", num_empty as i64, i64), - ("threshold_failure", num_threshold_failure as i64, i64), - ( - "votable_threshold_failure", - num_votable_threshold_failure as i64, - i64 - ), - ("tower_duration", ms as i64, i64), - ); - - (rv.map(|x| x.0.clone()), heaviest_bank_on_same_fork) - } - // Given a heaviest bank, `heaviest_bank` and the next votable bank - // `heaviest_bank_on_same_fork` as the validator's last vote, return + // `heaviest_bank_on_same_voted_fork` as the validator's last vote, return // a bank to vote on, a bank to reset to, pub(crate) fn select_vote_and_reset_forks( - heaviest_bank: &Option>, - heaviest_bank_on_same_fork: &Option>, + heaviest_bank: &Arc, + heaviest_bank_on_same_voted_fork: &Option>, ancestors: &HashMap>, descendants: &HashMap>, progress: &ProgressMap, @@ -1434,37 +1374,35 @@ impl ReplayStage { // switch_threshold succceeds let mut failure_reasons = vec![]; let selected_fork = { - if let Some(bank) = heaviest_bank { - let switch_fork_decision = tower.check_switch_threshold( - bank.slot(), - &ancestors, - &descendants, - &progress, - bank.total_epoch_stake(), - bank.epoch_vote_accounts(bank.epoch()).expect( - "Bank epoch vote accounts must contain entry for the bank's own epoch", - ), + let switch_fork_decision = tower.check_switch_threshold( + heaviest_bank.slot(), + &ancestors, + &descendants, + &progress, + heaviest_bank.total_epoch_stake(), + heaviest_bank + .epoch_vote_accounts(heaviest_bank.epoch()) + .expect("Bank epoch vote accounts must contain entry for the bank's own epoch"), + ); + if switch_fork_decision == SwitchForkDecision::FailedSwitchThreshold { + // If we can't switch, then reset to the the next votable + // bank on the same fork as our last vote, but don't vote + info!( + "Waiting to switch vote to {}, resetting to slot {:?} on same fork for now", + heaviest_bank.slot(), + heaviest_bank_on_same_voted_fork.as_ref().map(|b| b.slot()) ); - if switch_fork_decision == SwitchForkDecision::FailedSwitchThreshold { - // If we can't switch, then reset to the the next votable - // bank on the same fork as our last vote, but don't vote - info!( - "Waiting to switch vote to {}, resetting to slot {:?} on same fork for now", - bank.slot(), - heaviest_bank_on_same_fork.as_ref().map(|b| b.slot()) - ); - failure_reasons.push(HeaviestForkFailures::FailedSwitchThreshold(bank.slot())); - heaviest_bank_on_same_fork - .as_ref() - .map(|b| (b, switch_fork_decision)) - } else { - // If the switch threshold is observed, halt voting on - // the current fork and attempt to vote/reset Poh to - // the heaviest bank - heaviest_bank.as_ref().map(|b| (b, switch_fork_decision)) - } + failure_reasons.push(HeaviestForkFailures::FailedSwitchThreshold( + heaviest_bank.slot(), + )); + heaviest_bank_on_same_voted_fork + .as_ref() + .map(|b| (b, switch_fork_decision)) } else { - None + // If the switch threshold is observed, halt voting on + // the current fork and attempt to vote/reset Poh to + // the heaviest bank + Some((heaviest_bank, switch_fork_decision)) } }; @@ -1690,6 +1628,7 @@ impl ReplayStage { accounts_hash_sender: &Option, all_pubkeys: &mut PubkeyReferences, largest_confirmed_root: Option, + heaviest_subtree_fork_choice: &mut HeaviestSubtreeForkChoice, ) { let old_epoch = bank_forks.read().unwrap().root_bank().epoch(); bank_forks.write().unwrap().set_root( @@ -1703,6 +1642,7 @@ impl ReplayStage { all_pubkeys.purge(); } progress.handle_new_root(&r_bank_forks); + heaviest_subtree_fork_choice.set_root(new_root); } fn generate_new_bank_forks( @@ -1812,6 +1752,14 @@ impl ReplayStage { } } + pub fn get_unlock_heaviest_subtree_fork_choice(operating_mode: OperatingMode) -> Slot { + match operating_mode { + OperatingMode::Development => std::u64::MAX / 2, + OperatingMode::Stable => std::u64::MAX / 2, + OperatingMode::Preview => std::u64::MAX / 2, + } + } + pub fn join(self) -> thread::Result<()> { self.commitment_service.join()?; self.t_replay.join().map(|_| ()) @@ -1844,21 +1792,18 @@ pub(crate) mod tests { }; use solana_runtime::genesis_utils::{self, GenesisConfigInfo, ValidatorVoteKeypairs}; use solana_sdk::{ - account::Account, clock::NUM_CONSECUTIVE_LEADER_SLOTS, genesis_config, hash::{hash, Hash}, instruction::InstructionError, packet::PACKET_DATA_SIZE, - rent::Rent, signature::{Keypair, Signature, Signer}, system_transaction, transaction::TransactionError, }; - use solana_stake_program::stake_state; use solana_transaction_status::{EncodedTransaction, TransactionWithStatusMeta}; use solana_vote_program::{ - vote_state::{self, Vote, VoteState, VoteStateVersions}, + vote_state::{VoteState, VoteStateVersions}, vote_transaction, }; use std::{ @@ -1869,292 +1814,6 @@ pub(crate) mod tests { }; use trees::tr; - struct ForkInfo { - leader: usize, - fork: Vec, - voters: Vec, - } - - struct ValidatorInfo { - stake: u64, - keypair: Keypair, - authorized_voter_keypair: Keypair, - staking_keypair: Keypair, - } - - struct ForkSelectionResponse { - slot: u64, - is_locked_out: bool, - } - - fn simulate_fork_selection( - neutral_fork: &ForkInfo, - forks: &[ForkInfo], - validators: &[ValidatorInfo], - ) -> Vec> { - fn vote(bank: &Arc, pubkey: &Pubkey, slot: Slot) { - let mut vote_account = bank.get_account(&pubkey).unwrap(); - let mut vote_state = VoteState::from(&vote_account).unwrap(); - vote_state.process_slot_vote_unchecked(slot); - let versioned = VoteStateVersions::Current(Box::new(vote_state)); - VoteState::to(&versioned, &mut vote_account).unwrap(); - bank.store_account(&pubkey, &vote_account); - } - - let vote_tracker = VoteTracker::default(); - let cluster_slots = ClusterSlots::default(); - let mut towers: Vec = iter::repeat_with(|| Tower::new_for_tests(8, 0.67)) - .take(validators.len()) - .collect(); - - for slot in &neutral_fork.fork { - for tower in towers.iter_mut() { - tower.record_bank_vote(Vote { - hash: Hash::default(), - slots: vec![*slot], - timestamp: None, - }); - } - } - - for fork_info in forks.iter() { - for slot in fork_info.fork.iter() { - for voter_index in fork_info.voters.iter() { - towers[*voter_index].record_bank_vote(Vote { - hash: Hash::default(), - slots: vec![*slot], - timestamp: None, - }); - } - } - } - - let genesis_vote_accounts: Vec = validators - .iter() - .map(|validator| { - vote_state::create_account( - &validator.authorized_voter_keypair.pubkey(), - &validator.keypair.pubkey(), - 0, - validator.stake, - ) - }) - .collect(); - - let genesis_stake_accounts: Vec = validators - .iter() - .enumerate() - .map(|(i, validator)| { - stake_state::create_account( - &validator.staking_keypair.pubkey(), - &validator.authorized_voter_keypair.pubkey(), - &genesis_vote_accounts[i], - &Rent::default(), - validator.stake, - ) - }) - .collect(); - - let mut genesis_config = create_genesis_config(10_000).genesis_config; - genesis_config.accounts.clear(); - - for i in 0..validators.len() { - genesis_config.accounts.insert( - validators[i].authorized_voter_keypair.pubkey(), - genesis_vote_accounts[i].clone(), - ); - genesis_config.accounts.insert( - validators[i].staking_keypair.pubkey(), - genesis_stake_accounts[i].clone(), - ); - } - - let mut bank_forks = BankForks::new(neutral_fork.fork[0], Bank::new(&genesis_config)); - - let mut fork_progresses: Vec = iter::repeat_with(ProgressMap::default) - .take(validators.len()) - .collect(); - - for fork_progress in fork_progresses.iter_mut() { - let bank = &bank_forks.banks[&0]; - fork_progress - .entry(neutral_fork.fork[0]) - .or_insert_with(|| ForkProgress::new(bank.last_blockhash(), None, None, 0, 0)); - } - - for index in 1..neutral_fork.fork.len() { - let bank = Bank::new_from_parent( - &bank_forks.banks[&neutral_fork.fork[index - 1]].clone(), - &validators[neutral_fork.leader].keypair.pubkey(), - neutral_fork.fork[index], - ); - - bank_forks.insert(bank); - - for validator in validators.iter() { - vote( - &bank_forks.banks[&neutral_fork.fork[index]].clone(), - &validator.authorized_voter_keypair.pubkey(), - neutral_fork.fork[index - 1], - ); - } - - bank_forks.banks[&neutral_fork.fork[index]].freeze(); - - for fork_progress in fork_progresses.iter_mut() { - let bank = &bank_forks.banks[&neutral_fork.fork[index]]; - fork_progress - .entry(bank_forks.banks[&neutral_fork.fork[index]].slot()) - .or_insert_with(|| ForkProgress::new(bank.last_blockhash(), None, None, 0, 0)); - } - } - - let last_neutral_bank = &bank_forks.banks[neutral_fork.fork.last().unwrap()].clone(); - - for fork_info in forks.iter() { - for index in 0..fork_info.fork.len() { - let last_bank: &Arc; - let last_bank_in_fork: Arc; - - if index == 0 { - last_bank = &last_neutral_bank; - } else { - last_bank_in_fork = bank_forks.banks[&fork_info.fork[index - 1]].clone(); - last_bank = &last_bank_in_fork; - } - - let bank = Bank::new_from_parent( - last_bank, - &validators[fork_info.leader].keypair.pubkey(), - fork_info.fork[index], - ); - - bank_forks.insert(bank); - - for voter_index in fork_info.voters.iter() { - vote( - &bank_forks.banks[&fork_info.fork[index]].clone(), - &validators[*voter_index].authorized_voter_keypair.pubkey(), - last_bank.slot(), - ); - } - - bank_forks.banks[&fork_info.fork[index]].freeze(); - - for fork_progress in fork_progresses.iter_mut() { - let bank = &bank_forks.banks[&fork_info.fork[index]]; - fork_progress - .entry(bank_forks.banks[&fork_info.fork[index]].slot()) - .or_insert_with(|| { - ForkProgress::new(bank.last_blockhash(), None, None, 0, 0) - }); - } - } - } - - let bank_fork_ancestors = bank_forks.ancestors(); - let wrapped_bank_fork = Arc::new(RwLock::new(bank_forks)); - let mut all_pubkeys = PubkeyReferences::default(); - (0..validators.len()) - .map(|i| { - let mut frozen_banks: Vec<_> = wrapped_bank_fork - .read() - .unwrap() - .frozen_banks() - .values() - .cloned() - .collect(); - ReplayStage::compute_bank_stats( - &validators[i].keypair.pubkey(), - &bank_fork_ancestors, - &mut frozen_banks, - &towers[i], - &mut fork_progresses[i], - &vote_tracker, - &cluster_slots, - &wrapped_bank_fork, - &mut all_pubkeys, - ); - let (heaviest_bank, _) = ReplayStage::select_forks( - &frozen_banks, - &towers[i], - &fork_progresses[i], - &bank_fork_ancestors, - ); - - if let Some(bank) = heaviest_bank { - let stats = &fork_progresses[i].get_fork_stats(bank.slot()).unwrap(); - Some(ForkSelectionResponse { - slot: bank.slot(), - is_locked_out: stats.is_locked_out, - }) - } else { - None - } - }) - .collect() - } - - #[test] - fn test_minority_fork_overcommit_attack() { - let neutral_fork = ForkInfo { - leader: 0, - fork: vec![0, 1, 2], - voters: vec![], - }; - - let forks: Vec = vec![ - // Minority fork - ForkInfo { - leader: 2, - fork: (3..=3 + 8).collect(), - voters: vec![2], - }, - ForkInfo { - leader: 1, - fork: (12..12 + 8).collect(), - voters: vec![0, 1], - }, - ]; - - let validators: Vec = vec![ - ValidatorInfo { - stake: 34_000_000, - keypair: Keypair::new(), - authorized_voter_keypair: Keypair::new(), - staking_keypair: Keypair::new(), - }, - ValidatorInfo { - stake: 33_000_000, - keypair: Keypair::new(), - authorized_voter_keypair: Keypair::new(), - staking_keypair: Keypair::new(), - }, - // Malicious Node - ValidatorInfo { - stake: 33_000_000, - keypair: Keypair::new(), - authorized_voter_keypair: Keypair::new(), - staking_keypair: Keypair::new(), - }, - ]; - - let resp = simulate_fork_selection(&neutral_fork, &forks, &validators); - // Both honest nodes are now want to switch to minority fork and are locked out - assert!(resp[0].is_some()); - assert_eq!(resp[0].as_ref().unwrap().is_locked_out, true); - assert_eq!( - resp[0].as_ref().unwrap().slot, - forks[0].fork.last().unwrap().clone() - ); - assert!(resp[1].is_some()); - assert_eq!(resp[1].as_ref().unwrap().is_locked_out, true); - assert_eq!( - resp[1].as_ref().unwrap().slot, - forks[0].fork.last().unwrap().clone() - ); - } - #[test] fn test_child_slots_of_same_parent() { let ledger_path = get_tmp_ledger_path!(); @@ -2304,6 +1963,7 @@ pub(crate) mod tests { let bank0 = Bank::new(&genesis_config); let bank_forks = Arc::new(RwLock::new(BankForks::new(0, bank0))); let root = 3; + let mut heaviest_subtree_fork_choice = HeaviestSubtreeForkChoice::new(root); let root_bank = Bank::new_from_parent( bank_forks.read().unwrap().get(0).unwrap(), &Pubkey::default(), @@ -2321,6 +1981,7 @@ pub(crate) mod tests { &None, &mut PubkeyReferences::default(), None, + &mut heaviest_subtree_fork_choice, ); assert_eq!(bank_forks.read().unwrap().root(), root); assert_eq!(progress.len(), 1); @@ -2353,6 +2014,7 @@ pub(crate) mod tests { root, ); bank_forks.write().unwrap().insert(root_bank); + let mut heaviest_subtree_fork_choice = HeaviestSubtreeForkChoice::new(root); let mut progress = ProgressMap::default(); for i in 0..=root { progress.insert(i, ForkProgress::new(Hash::default(), None, None, 0, 0)); @@ -2364,6 +2026,7 @@ pub(crate) mod tests { &None, &mut PubkeyReferences::default(), Some(confirmed_root), + &mut heaviest_subtree_fork_choice, ); assert_eq!(bank_forks.read().unwrap().root(), root); assert!(bank_forks.read().unwrap().get(confirmed_root).is_some()); @@ -2880,7 +2543,8 @@ pub(crate) mod tests { ValidatorVoteKeypairs::new(node_keypair, vote_keypair, stake_keypair), ); - let (bank_forks, mut progress) = initialize_state(&keypairs, 10_000); + let (bank_forks, mut progress, mut heaviest_subtree_fork_choice) = + initialize_state(&keypairs, 10_000); let bank0 = bank_forks.get(0).unwrap().clone(); let my_keypairs = keypairs.get(&node_pubkey).unwrap(); let vote_tx = vote_transaction::new_vote_transaction( @@ -2917,6 +2581,8 @@ pub(crate) mod tests { &ClusterSlots::default(), &bank_forks, &mut PubkeyReferences::default(), + &mut heaviest_subtree_fork_choice, + &mut BankWeightForkChoice::default(), ); assert_eq!(newly_computed, vec![0]); // The only vote is in bank 1, and bank_forks does not currently contain @@ -2958,6 +2624,8 @@ pub(crate) mod tests { &ClusterSlots::default(), &bank_forks, &mut PubkeyReferences::default(), + &mut heaviest_subtree_fork_choice, + &mut BankWeightForkChoice::default(), ); assert_eq!(newly_computed, vec![1]); @@ -2991,6 +2659,8 @@ pub(crate) mod tests { &ClusterSlots::default(), &bank_forks, &mut PubkeyReferences::default(), + &mut heaviest_subtree_fork_choice, + &mut BankWeightForkChoice::default(), ); // No new stats should have been computed assert!(newly_computed.is_empty()); @@ -3005,7 +2675,8 @@ pub(crate) mod tests { // Create the tree of banks in a BankForks object let forks = tr(0) / (tr(1)) / (tr(2)); - vote_simulator.fill_bank_forks(forks, &HashMap::new()); + vote_simulator.fill_bank_forks(forks.clone(), &HashMap::new()); + let mut heaviest_subtree_fork_choice = HeaviestSubtreeForkChoice::new_from_tree(forks); let mut frozen_banks: Vec<_> = vote_simulator .bank_forks .read() @@ -3017,7 +2688,7 @@ pub(crate) mod tests { let ancestors = vote_simulator.bank_forks.read().unwrap().ancestors(); ReplayStage::compute_bank_stats( - &Pubkey::default(), + &node_pubkey, &ancestors, &mut frozen_banks, &tower, @@ -3026,25 +2697,36 @@ pub(crate) mod tests { &ClusterSlots::default(), &vote_simulator.bank_forks, &mut PubkeyReferences::default(), + &mut heaviest_subtree_fork_choice, + &mut BankWeightForkChoice::default(), ); assert_eq!( - vote_simulator - .progress - .get_fork_stats(1) - .unwrap() - .fork_weight, - vote_simulator - .progress - .get_fork_stats(2) - .unwrap() - .fork_weight, + heaviest_subtree_fork_choice.stake_voted_subtree(1).unwrap(), + heaviest_subtree_fork_choice.stake_voted_subtree(2).unwrap() + ); + + let (heaviest_bank, _) = heaviest_subtree_fork_choice.select_forks( + &frozen_banks, + &tower, + &vote_simulator.progress, + &ancestors, + &vote_simulator.bank_forks, ); - let (heaviest_bank, _) = - ReplayStage::select_forks(&frozen_banks, &tower, &vote_simulator.progress, &ancestors); // Should pick the lower of the two equally weighted banks - assert_eq!(heaviest_bank.unwrap().slot(), 1); + assert_eq!(heaviest_bank.slot(), 1); + + let (heaviest_bank, _) = BankWeightForkChoice::default().select_forks( + &frozen_banks, + &tower, + &vote_simulator.progress, + &ancestors, + &vote_simulator.bank_forks, + ); + + // Should pick the lower of the two equally weighted banks + assert_eq!(heaviest_bank.slot(), 1); } #[test] @@ -3062,6 +2744,8 @@ pub(crate) mod tests { let votes = vec![0, 2]; cluster_votes.insert(node_pubkey, votes.clone()); vote_simulator.fill_bank_forks(forks, &cluster_votes); + + // Fill banks with votes for vote in votes { assert!(vote_simulator .simulate_vote(vote, &node_pubkey, &mut tower,) @@ -3078,7 +2762,7 @@ pub(crate) mod tests { .collect(); ReplayStage::compute_bank_stats( - &Pubkey::default(), + &node_pubkey, &vote_simulator.bank_forks.read().unwrap().ancestors(), &mut frozen_banks, &tower, @@ -3087,6 +2771,8 @@ pub(crate) mod tests { &ClusterSlots::default(), &vote_simulator.bank_forks, &mut PubkeyReferences::default(), + &mut vote_simulator.heaviest_subtree_fork_choice, + &mut BankWeightForkChoice::default(), ); frozen_banks.sort_by_key(|bank| bank.slot()); @@ -3103,6 +2789,16 @@ pub(crate) mod tests { .fork_weight; assert!(second >= first); } + for bank in frozen_banks { + // The only leaf should always be chosen over parents + assert_eq!( + vote_simulator + .heaviest_subtree_fork_choice + .best_slot(bank.slot()) + .unwrap(), + 3 + ); + } } #[test] @@ -3200,7 +2896,7 @@ pub(crate) mod tests { success_index: usize, ) { let stake = 10_000; - let (bank_forks, _) = initialize_state(&all_keypairs, stake); + let (bank_forks, _, _) = initialize_state(&all_keypairs, stake); let root_bank = bank_forks.root_bank().clone(); let mut propagated_stats = PropagatedStats { total_epoch_stake: stake * all_keypairs.len() as u64, @@ -3337,7 +3033,7 @@ pub(crate) mod tests { ); let stake = 10_000; - let (mut bank_forks, mut progress_map) = initialize_state(&keypairs, stake); + let (mut bank_forks, mut progress_map, _) = initialize_state(&keypairs, stake); let bank0 = bank_forks.get(0).unwrap().clone(); bank_forks.insert(Bank::new_from_parent(&bank0, &Pubkey::default(), 9)); @@ -3435,7 +3131,8 @@ pub(crate) mod tests { .collect(); let stake_per_validator = 10_000; - let (mut bank_forks, mut progress_map) = initialize_state(&keypairs, stake_per_validator); + let (mut bank_forks, mut progress_map, _) = + initialize_state(&keypairs, stake_per_validator); bank_forks.set_root(0, &None, None); let total_epoch_stake = bank_forks.root_bank().total_epoch_stake(); @@ -3517,7 +3214,8 @@ pub(crate) mod tests { .collect(); let stake_per_validator = 10_000; - let (mut bank_forks, mut progress_map) = initialize_state(&keypairs, stake_per_validator); + let (mut bank_forks, mut progress_map, _) = + initialize_state(&keypairs, stake_per_validator); bank_forks.set_root(0, &None, None); let total_epoch_stake = num_validators as u64 * stake_per_validator; diff --git a/local-cluster/tests/local_cluster.rs b/local-cluster/tests/local_cluster.rs index 92620f913b..53b1bca97c 100644 --- a/local-cluster/tests/local_cluster.rs +++ b/local-cluster/tests/local_cluster.rs @@ -357,20 +357,18 @@ fn test_cluster_partition_1_1_1() { #[test] #[serial] -fn test_kill_partition() { +fn test_kill_heaviest_partition() { // This test: - // 1) Spins up three partitions - // 2) Forces more slots in the leader schedule for the first partition so - // that this partition will be the heaviiest - // 3) Schedules the other validators for sufficient slots in the schedule + // 1) Spins up four partitions, the heaviest being the first with more stake + // 2) Schedules the other validators for sufficient slots in the schedule // so that they will still be locked out of voting for the major partitoin // when the partition resolves - // 4) Kills the major partition. Validators are locked out, but should be - // able to reset to the major partition - // 5) Check for recovery + // 3) Kills the most staked partition. Validators are locked out, but should all + // eventually choose the major partition + // 4) Check for recovery let mut leader_schedule = vec![]; let num_slots_per_validator = 8; - let partitions: [&[usize]; 3] = [&[9], &[10], &[10]]; + let partitions: [&[usize]; 4] = [&[11], &[10], &[10], &[10]]; let validator_keys: Vec<_> = iter::repeat_with(|| Arc::new(Keypair::new())) .take(partitions.len()) .collect(); diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index b0f4067152..7a01388512 100644 --- a/runtime/src/bank.rs +++ b/runtime/src/bank.rs @@ -2411,6 +2411,10 @@ impl Bank { self.epoch_stakes.get(&epoch) } + pub fn epoch_stakes_map(&self) -> &HashMap { + &self.epoch_stakes + } + /// vote accounts for the specific epoch along with the stake /// attributed to each account pub fn epoch_vote_accounts(&self, epoch: Epoch) -> Option<&HashMap> { @@ -2448,11 +2452,11 @@ impl Bank { } /// Get the fixed stake of the given vote account for the current epoch - pub fn epoch_vote_account_stake(&self, voting_pubkey: &Pubkey) -> u64 { + pub fn epoch_vote_account_stake(&self, vote_account: &Pubkey) -> u64 { *self .epoch_vote_accounts(self.epoch()) .expect("Bank epoch vote accounts must contain entry for the bank's own epoch") - .get(voting_pubkey) + .get(vote_account) .map(|(stake, _)| stake) .unwrap_or(&0) } diff --git a/runtime/src/epoch_stakes.rs b/runtime/src/epoch_stakes.rs index 7b95ecabe1..b7885c02d4 100644 --- a/runtime/src/epoch_stakes.rs +++ b/runtime/src/epoch_stakes.rs @@ -50,6 +50,13 @@ impl EpochStakes { &self.epoch_authorized_voters } + pub fn vote_account_stake(&self, vote_account: &Pubkey) -> u64 { + Stakes::vote_accounts(&self.stakes) + .get(vote_account) + .map(|(stake, _)| *stake) + .unwrap_or(0) + } + fn parse_epoch_vote_accounts( epoch_vote_accounts: &HashMap, leader_schedule_epoch: Epoch,