anza-xyz-agave/core/src/vote_simulator.rs

#![cfg(feature = "dev-context-only-utils")]
use {
    crate::{
        cluster_info_vote_listener::VoteTracker,
        cluster_slots_service::cluster_slots::ClusterSlots,
        consensus::{
            fork_choice::{select_vote_and_reset_forks, SelectVoteAndResetForkResult},
            heaviest_subtree_fork_choice::HeaviestSubtreeForkChoice,
            latest_validator_votes_for_frozen_banks::LatestValidatorVotesForFrozenBanks,
            progress_map::{ForkProgress, ProgressMap},
            tower_vote_state::TowerVoteState,
            Tower,
        },
        repair::cluster_slot_state_verifier::{
            DuplicateConfirmedSlots, DuplicateSlotsTracker, EpochSlotsFrozenSlots,
        },
        replay_stage::{HeaviestForkFailures, ReplayStage, TowerBFTStructures},
        unfrozen_gossip_verified_vote_hashes::UnfrozenGossipVerifiedVoteHashes,
    },
    crossbeam_channel::unbounded,
    solana_clock::Slot,
    solana_hash::Hash,
    solana_pubkey::Pubkey,
    solana_runtime::{
        bank::Bank,
        bank_forks::BankForks,
        genesis_utils::{
            create_genesis_config_with_vote_accounts, GenesisConfigInfo, ValidatorVoteKeypairs,
        },
    },
    solana_signer::Signer,
    solana_vote::vote_transaction,
    solana_vote_program::vote_state::{Lockout, TowerSync},
    std::{
        collections::{HashMap, HashSet, VecDeque},
        sync::{Arc, RwLock},
    },
    trees::{tr, Tree, TreeWalk},
};

pub struct VoteSimulator {
    pub validator_keypairs: HashMap<Pubkey, ValidatorVoteKeypairs>,
    pub node_pubkeys: Vec<Pubkey>,
    pub vote_pubkeys: Vec<Pubkey>,
    pub bank_forks: Arc<RwLock<BankForks>>,
    pub progress: ProgressMap,
    pub latest_validator_votes_for_frozen_banks: LatestValidatorVotesForFrozenBanks,
    pub tbft_structs: TowerBFTStructures,
}

impl VoteSimulator {
    pub fn new(num_keypairs: usize) -> Self {
        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,
            progress,
            latest_validator_votes_for_frozen_banks: LatestValidatorVotesForFrozenBanks::default(),
            tbft_structs: TowerBFTStructures {
                heaviest_subtree_fork_choice,
                duplicate_slots_tracker: DuplicateSlotsTracker::default(),
                duplicate_confirmed_slots: DuplicateConfirmedSlots::default(),
                unfrozen_gossip_verified_vote_hashes: UnfrozenGossipVerifiedVoteHashes::default(),
                epoch_slots_frozen_slots: EpochSlotsFrozenSlots::default(),
            },
        }
    }

    pub fn fill_bank_forks(
        &mut self,
        forks: Tree<u64>,
        cluster_votes: &HashMap<Pubkey, Vec<u64>>,
        is_frozen: bool,
    ) {
        let root = *forks.root().data();
        assert!(self.bank_forks.read().unwrap().get(root).is_some());

        let mut walk = TreeWalk::from(forks);

        while let Some(visit) = walk.get() {
            let slot = *visit.node().data();
            if self.bank_forks.read().unwrap().get(slot).is_some() {
                walk.forward();
                continue;
            }
            let parent = *walk.get_parent().unwrap().data();
            let parent_bank = self.bank_forks.read().unwrap().get(parent).unwrap();
            let new_bank = Bank::new_from_parent(parent_bank.clone(), &Pubkey::default(), slot);
            let new_bank = self
                .bank_forks
                .write()
                .unwrap()
                .insert(new_bank)
                .clone_without_scheduler();
            self.progress
                .entry(slot)
                .or_insert_with(|| ForkProgress::new(Hash::default(), None, None, 0, 0));
            for (pubkey, vote) in cluster_votes.iter() {
                if vote.contains(&parent) {
                    let keypairs = self.validator_keypairs.get(pubkey).unwrap();
                    let latest_blockhash = parent_bank.last_blockhash();
                    let tower_sync = if let Some(vote_account) =
                        parent_bank.get_vote_account(&keypairs.vote_keypair.pubkey())
                    {
                        let mut vote_state = TowerVoteState::from(
                            vote_account.vote_state_view().expect("must be TowerBFT"),
                        );
                        vote_state.process_next_vote_slot(parent);
                        TowerSync::new(
                            vote_state.votes,
                            vote_state.root_slot,
                            parent_bank.hash(),
                            Hash::default(),
                        )
                    } else {
                        TowerSync::new(
                            VecDeque::from([Lockout::new(parent)]),
                            Some(root),
                            parent_bank.hash(),
                            Hash::default(),
                        )
                    };
                    let vote_tx = vote_transaction::new_tower_sync_transaction(
                        tower_sync,
                        latest_blockhash,
                        &keypairs.node_keypair,
                        &keypairs.vote_keypair,
                        &keypairs.vote_keypair,
                        None,
                    );
                    info!("voting {} {}", parent_bank.slot(), parent_bank.hash());
                    new_bank.process_transaction(&vote_tx).unwrap();

                    // Check the vote landed
                    let vote_account = new_bank
                        .get_vote_account(&keypairs.vote_keypair.pubkey())
                        .unwrap();
                    let vote_state_view = vote_account.vote_state_view().unwrap();
                    assert!(vote_state_view
                        .votes_iter()
                        .any(|lockout| lockout.slot() == parent));
                }
            }

            new_bank.fill_bank_with_ticks_for_tests();
            if !visit.node().has_no_child() || is_frozen {
                new_bank.set_block_id(Some(Hash::new_unique()));
                new_bank.freeze();
                self.progress
                    .get_fork_stats_mut(new_bank.slot())
                    .expect("All frozen banks must exist in the Progress map")
                    .bank_hash = Some(new_bank.hash());
                self.tbft_structs
                    .heaviest_subtree_fork_choice
                    .add_new_leaf_slot(
                        (new_bank.slot(), new_bank.hash()),
                        Some((new_bank.parent_slot(), new_bank.parent_hash())),
                    );
            }

            walk.forward();
        }
    }

    pub fn simulate_vote(
        &mut self,
        vote_slot: Slot,
        my_pubkey: &Pubkey,
        tower: &mut Tower,
    ) -> Vec<HeaviestForkFailures> {
        // Try to simulate the vote
        let ancestors = self.bank_forks.read().unwrap().ancestors();
        let mut frozen_banks: Vec<_> = self
            .bank_forks
            .read()
            .unwrap()
            .frozen_banks()
            .map(|(_slot, bank)| bank)
            .collect();

        let _ = ReplayStage::compute_bank_stats(
            my_pubkey,
            &ancestors,
            &mut frozen_banks,
            tower,
            &mut self.progress,
            &VoteTracker::default(),
            &ClusterSlots::default(),
            &self.bank_forks,
            &mut self.tbft_structs.heaviest_subtree_fork_choice,
            &mut self.latest_validator_votes_for_frozen_banks,
        );

        let vote_bank = self
            .bank_forks
            .read()
            .unwrap()
            .get(vote_slot)
            .expect("Bank must have been created before vote simulation");

        // Try to vote on the given slot
        let descendants = self.bank_forks.read().unwrap().descendants();
        let SelectVoteAndResetForkResult {
            heaviest_fork_failures,
            ..
        } = select_vote_and_reset_forks(
            &vote_bank,
            None,
            &ancestors,
            &descendants,
            &self.progress,
            tower,
            &self.latest_validator_votes_for_frozen_banks,
            &self.tbft_structs.heaviest_subtree_fork_choice,
        );

        // Make sure this slot isn't locked out or failing threshold
        info!("Checking vote: {}", vote_bank.slot());
        if !heaviest_fork_failures.is_empty() {
            return heaviest_fork_failures;
        }

        let new_root = tower.record_bank_vote(&vote_bank);
        if let Some(new_root) = new_root {
            self.set_root(new_root);
        }

        vec![]
    }

    pub fn set_root(&mut self, new_root: Slot) {
        let (drop_bank_sender, _drop_bank_receiver) = unbounded();
        ReplayStage::handle_new_root(
            new_root,
            &self.bank_forks,
            &mut self.progress,
            None, // snapshot_controller
            None,
            &mut true,
            &drop_bank_sender,
            &mut self.tbft_structs,
        )
        .unwrap();
    }

    pub fn create_and_vote_new_branch(
        &mut self,
        start_slot: Slot,
        end_slot: Slot,
        cluster_votes: &HashMap<Pubkey, Vec<u64>>,
        votes_to_simulate: &HashSet<Slot>,
        my_pubkey: &Pubkey,
        tower: &mut Tower,
    ) -> HashMap<Slot, Vec<HeaviestForkFailures>> {
        (start_slot + 1..=end_slot)
            .filter_map(|slot| {
                let mut fork_tip_parent = tr(slot - 1);
                fork_tip_parent.push_front(tr(slot));
                self.fill_bank_forks(fork_tip_parent, cluster_votes, true);
                if votes_to_simulate.contains(&slot) {
                    Some((slot, self.simulate_vote(slot, my_pubkey, tower)))
                } else {
                    None
                }
            })
            .collect()
    }

    pub fn simulate_lockout_interval(
        &mut self,
        slot: Slot,
        lockout_interval: (u64, u64),
        vote_account_pubkey: &Pubkey,
    ) {
        self.progress
            .entry(slot)
            .or_insert_with(|| ForkProgress::new(Hash::default(), None, None, 0, 0))
            .fork_stats
            .lockout_intervals
            .entry(lockout_interval.1)
            .or_default()
            .push((lockout_interval.0, *vote_account_pubkey));
    }

    pub fn clear_lockout_intervals(&mut self, slot: Slot) {
        self.progress
            .entry(slot)
            .or_insert_with(|| ForkProgress::new(Hash::default(), None, None, 0, 0))
            .fork_stats
            .lockout_intervals
            .clear()
    }

    pub fn can_progress_on_fork(
        &mut self,
        my_pubkey: &Pubkey,
        tower: &mut Tower,
        start_slot: u64,
        num_slots: u64,
        cluster_votes: &mut HashMap<Pubkey, Vec<u64>>,
    ) -> bool {
        // Check that within some reasonable time, validator can make a new
        // root on this fork
        let old_root = tower.root();

        for i in 1..num_slots {
            // The parent of the tip of the fork
            let mut fork_tip_parent = tr(start_slot + i - 1);
            // The tip of the fork
            fork_tip_parent.push_front(tr(start_slot + i));
            self.fill_bank_forks(fork_tip_parent, cluster_votes, true);
            if self
                .simulate_vote(i + start_slot, my_pubkey, tower)
                .is_empty()
            {
                cluster_votes
                    .entry(*my_pubkey)
                    .or_default()
                    .push(start_slot + i);
            }
            if old_root != tower.root() {
                return true;
            }
        }

        false
    }

    #[allow(clippy::type_complexity)]
    fn init_state(
        num_keypairs: usize,
    ) -> (
        HashMap<Pubkey, ValidatorVoteKeypairs>,
        Vec<Pubkey>,
        Vec<Pubkey>,
        Arc<RwLock<BankForks>>,
        ProgressMap,
        HeaviestSubtreeForkChoice,
    ) {
        let keypairs: HashMap<_, _> = std::iter::repeat_with(|| {
            let vote_keypairs = ValidatorVoteKeypairs::new_rand();
            (vote_keypairs.node_keypair.pubkey(), vote_keypairs)
        })
        .take(num_keypairs)
        .collect();
        let node_pubkeys: Vec<_> = keypairs
            .values()
            .map(|keys| keys.node_keypair.pubkey())
            .collect();
        let vote_pubkeys: Vec<_> = keypairs
            .values()
            .map(|keys| keys.vote_keypair.pubkey())
            .collect();

        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,
        )
    }
}

// Setup BankForks with bank 0 and all the validator accounts
pub fn initialize_state(
    validator_keypairs_map: &HashMap<Pubkey, ValidatorVoteKeypairs>,
    stake: u64,
) -> (
    Arc<RwLock<BankForks>>,
    ProgressMap,
    HeaviestSubtreeForkChoice,
) {
    let validator_keypairs: Vec<_> = validator_keypairs_map.values().collect();
    let GenesisConfigInfo {
        mut genesis_config,
        mint_keypair,
        ..
    } = create_genesis_config_with_vote_accounts(
        1_000_000_000,
        &validator_keypairs,
        vec![stake; validator_keypairs.len()],
    );

    genesis_config.poh_config.hashes_per_tick = Some(2);
    let (bank0, bank_forks) = Bank::new_with_bank_forks_for_tests(&genesis_config);
    bank0.set_block_id(Some(Hash::new_unique()));

    for pubkey in validator_keypairs_map.keys() {
        bank0.transfer(10_000, &mint_keypair, pubkey).unwrap();
    }

    bank0.fill_bank_with_ticks_for_tests();
    bank0.freeze();
    let mut progress = ProgressMap::default();
    progress.insert(
        0,
        ForkProgress::new_from_bank(&bank0, bank0.collector_id(), &Pubkey::default(), None, 0, 0),
    );
    let heaviest_subtree_fork_choice =
        HeaviestSubtreeForkChoice::new_from_bank_forks(bank_forks.clone());
    (bank_forks, progress, heaviest_subtree_fork_choice)
}