From a7c3aacf2e0ab18655153bee01f2a285c3b2863e Mon Sep 17 00:00:00 2001 From: Ashwin Sekar Date: Tue, 13 Feb 2024 08:38:49 +0000 Subject: [PATCH] pr feedback: add unit tests --- core/src/consensus.rs | 22 ++++++++++ core/src/replay_stage.rs | 73 ++++++++++++++++++++++++++++--- sdk/program/src/vote/state/mod.rs | 18 ++++++++ 3 files changed, 106 insertions(+), 7 deletions(-) diff --git a/core/src/consensus.rs b/core/src/consensus.rs index 4f129b18282218..3e24f33233863e 100644 --- a/core/src/consensus.rs +++ b/core/src/consensus.rs @@ -292,6 +292,28 @@ impl Tower { } } + #[cfg(test)] + pub fn new_random(node_pubkey: Pubkey) -> Self { + use rand::Rng; + + let mut rng = rand::thread_rng(); + let root_slot = rng.gen(); + let vote_state = VoteState::new_rand_for_tests(node_pubkey, root_slot); + let last_vote = VoteStateUpdate::from( + vote_state + .votes + .iter() + .map(|lv| (lv.slot(), lv.confirmation_count())) + .collect::>(), + ); + Self { + node_pubkey, + vote_state, + last_vote: VoteTransaction::CompactVoteStateUpdate(last_vote), + ..Tower::default() + } + } + pub fn new_from_bankforks( bank_forks: &BankForks, node_pubkey: &Pubkey, diff --git a/core/src/replay_stage.rs b/core/src/replay_stage.rs index 90eb8857f8142d..c09bc92b404ddb 100644 --- a/core/src/replay_stage.rs +++ b/core/src/replay_stage.rs @@ -584,7 +584,7 @@ impl ReplayStage { // set-identity was called during the startup procedure, ensure the tower is consistent // before starting the loop. further calls to set-identity will reload the tower in the loop tower = Self::reload_tower_with_new_identity( - &tower_storage, + tower_storage.as_ref(), &my_pubkey, &vote_account, &bank_forks, @@ -1000,7 +1000,7 @@ impl ReplayStage { // Load the new identity's tower tower = Self::reload_tower_with_new_identity( - &tower_storage, + tower_storage.as_ref(), &my_pubkey, &vote_account, &bank_forks, @@ -1155,12 +1155,12 @@ impl ReplayStage { } fn reload_tower_with_new_identity( - tower_storage: &Arc, + tower_storage: &dyn TowerStorage, new_identity: &Pubkey, vote_account: &Pubkey, bank_forks: &Arc>, ) -> Tower { - Tower::restore(tower_storage.as_ref(), new_identity) + Tower::restore(tower_storage, new_identity) .and_then(|restored_tower| { let root_bank = bank_forks.read().unwrap().root_bank(); let slot_history = root_bank.get_slot_history(); @@ -4256,9 +4256,9 @@ pub(crate) mod tests { crate::{ consensus::{ progress_map::{ValidatorStakeInfo, RETRANSMIT_BASE_DELAY_MS}, - tower_storage::NullTowerStorage, + tower_storage::{FileTowerStorage, NullTowerStorage}, tree_diff::TreeDiff, - Tower, + Tower, VOTE_THRESHOLD_DEPTH, }, replay_stage::ReplayStage, vote_simulator::{self, VoteSimulator}, @@ -4280,7 +4280,7 @@ pub(crate) mod tests { }, solana_runtime::{ accounts_background_service::AbsRequestSender, - commitment::BlockCommitment, + commitment::{BlockCommitment, VOTE_THRESHOLD_SIZE}, genesis_utils::{GenesisConfigInfo, ValidatorVoteKeypairs}, }, solana_sdk::{ @@ -4304,6 +4304,7 @@ pub(crate) mod tests { iter, sync::{atomic::AtomicU64, Arc, RwLock}, }, + tempfile::tempdir, trees::{tr, Tree}, }; @@ -8624,4 +8625,62 @@ pub(crate) mod tests { assert_eq!(reset_fork, Some(4)); assert_eq!(failures, vec![HeaviestForkFailures::LockedOut(4),]); } + + #[test] + fn test_tower_reload_missing() { + let mut tower_file = tempdir().unwrap().into_path(); + tower_file = tower_file.with_file_name("tower.txt"); + let tower_storage = FileTowerStorage::new(tower_file); + let identity = Pubkey::new_unique(); + let vote_account = Pubkey::new_unique(); + let tree = tr(0) / (tr(1) / (tr(3) / (tr(4))) / (tr(2) / (tr(5) / (tr(6))))); + let generate_votes = |pubkeys: Vec| { + pubkeys + .into_iter() + .zip(iter::once(vec![0, 1, 2, 5, 6]).chain(iter::repeat(vec![0, 1, 3, 4]).take(2))) + .collect() + }; + let (vote_simulator, _blockstore) = + setup_forks_from_tree(tree, 3, Some(Box::new(generate_votes))); + let bank_forks = vote_simulator.bank_forks; + + let tower = ReplayStage::reload_tower_with_new_identity( + &tower_storage, + &identity, + &vote_account, + &bank_forks, + ); + let expected_tower = Tower::new_for_tests(VOTE_THRESHOLD_DEPTH, VOTE_THRESHOLD_SIZE); + assert_eq!(tower.vote_state, expected_tower.vote_state); + } + + #[test] + fn test_tower_reload() { + let mut tower_file = tempdir().unwrap().into_path(); + tower_file = tower_file.with_file_name("tower.txt"); + let tower_storage = FileTowerStorage::new(tower_file); + let node_keypair = Keypair::new(); + let identity = node_keypair.pubkey(); + let vote_account = Pubkey::new_unique(); + let tree = tr(0) / (tr(1) / (tr(3) / (tr(4))) / (tr(2) / (tr(5) / (tr(6))))); + let generate_votes = |pubkeys: Vec| { + pubkeys + .into_iter() + .zip(iter::once(vec![0, 1, 2, 5, 6]).chain(iter::repeat(vec![0, 1, 3, 4]).take(2))) + .collect() + }; + let (vote_simulator, _blockstore) = + setup_forks_from_tree(tree, 3, Some(Box::new(generate_votes))); + let bank_forks = vote_simulator.bank_forks; + let expected_tower = Tower::new_random(identity); + expected_tower.save(&tower_storage, &node_keypair).unwrap(); + + let tower = ReplayStage::reload_tower_with_new_identity( + &tower_storage, + &identity, + &vote_account, + &bank_forks, + ); + assert_eq!(tower.vote_state, expected_tower.vote_state); + } } diff --git a/sdk/program/src/vote/state/mod.rs b/sdk/program/src/vote/state/mod.rs index 8cfcd0ef19d9e8..a6e765472750c6 100644 --- a/sdk/program/src/vote/state/mod.rs +++ b/sdk/program/src/vote/state/mod.rs @@ -352,6 +352,24 @@ impl VoteState { } } + pub fn new_rand_for_tests(node_pubkey: Pubkey, root_slot: Slot) -> Self { + let votes = (1..32) + .map(|x| LandedVote { + latency: 0, + lockout: Lockout::new_with_confirmation_count( + u64::from(x).saturating_add(root_slot), + 32_u32.saturating_sub(x), + ), + }) + .collect(); + Self { + node_pubkey, + root_slot: Some(root_slot), + votes, + ..VoteState::default() + } + } + pub fn get_authorized_voter(&self, epoch: Epoch) -> Option { self.authorized_voters.get_authorized_voter(epoch) }