Add get_supermajority_slot() function (#2976)

* Moved supermajority functions into new module, staking_utils

* Move staking functions out of bank, and into staking_utils, change get_supermajority_slot to only use state from epoch boundary

* Move bank slot height in staked_nodes_at_slot() to be bank id

runtime/src/bank.rs

@@ -8,7 +8,6 @@ use crate::last_id_queue::LastIdQueue;
 use crate::runtime::{self, RuntimeError};
 use crate::status_cache::StatusCache;
 use bincode::serialize;
-use hashbrown::HashMap;
 use log::*;
 use solana_metrics::counter::Counter;
 use solana_sdk::account::Account;
@@ -609,7 +608,7 @@ impl Bank {
     }
 
     /// Compute all the parents of the bank in order
-    fn parents(&self) -> Vec<Arc<Bank>> {
+    pub fn parents(&self) -> Vec<Arc<Bank>> {
         let mut parents = vec![];
         let mut bank = self.parent();
         while let Some(parent) = bank {
@@ -687,48 +686,10 @@ impl Bank {
         extend_and_hash(&self.parent_hash, &serialize(&accounts_delta_hash).unwrap())
     }
 
-    pub fn vote_states<F>(&self, cond: F) -> Vec<VoteState>
+    /// Return the number of slots in advance of an epoch that a leader scheduler
-    where
+    /// should be generated.
-        F: Fn(&VoteState) -> bool,
+    pub fn stakers_slot_offset(&self) -> u64 {
-    {
+        self.stakers_slot_offset
-        self.accounts()
-            .accounts_db
-            .get_vote_accounts(self.id)
-            .iter()
-            .filter_map(|account| {
-                if let Ok(vote_state) = VoteState::deserialize(&account.userdata) {
-                    if cond(&vote_state) {
-                        return Some(vote_state);
-                    }
-                }
-                None
-            })
-            .collect()
-    }
-
-    /// Collect the node Pubkey and staker account balance for nodes
-    /// that have non-zero balance in their corresponding staker accounts
-    pub fn staked_nodes(&self) -> HashMap<Pubkey, u64> {
-        self.vote_states(|state| self.get_balance(&state.staker_id) > 0)
-            .iter()
-            .map(|state| (state.node_id, self.get_balance(&state.staker_id)))
-            .collect()
-    }
-
-    /// Return the checkpointed stakes that should be used to generate a leader schedule.
-    fn staked_nodes_at_slot(&self, current_slot_height: u64) -> HashMap<Pubkey, u64> {
-        let slot_height = current_slot_height.saturating_sub(self.stakers_slot_offset);
-
-        let parents = self.parents();
-        let mut banks = vec![self];
-        banks.extend(parents.iter().map(|x| x.as_ref()));
-
-        let bank = banks
-            .iter()
-            .find(|bank| bank.slot_height() <= slot_height)
-            .unwrap_or_else(|| banks.last().unwrap());
-
-        bank.staked_nodes()
     }
 
     /// Return the number of ticks per slot that should be used calls to slot_height().
@@ -756,10 +717,23 @@ impl Bank {
         self.slots_per_epoch
     }
 
-    /// Return the checkpointed stakes that should be used to generate a leader schedule.
+    pub fn vote_states<F>(&self, cond: F) -> Vec<VoteState>
-    pub fn staked_nodes_at_epoch(&self, epoch_height: u64) -> HashMap<Pubkey, u64> {
+    where
-        let epoch_slot_height = epoch_height * self.slots_per_epoch();
+        F: Fn(&VoteState) -> bool,
-        self.staked_nodes_at_slot(epoch_slot_height)
+    {
+        self.accounts()
+            .accounts_db
+            .get_vote_accounts(self.id)
+            .iter()
+            .filter_map(|account| {
+                if let Ok(vote_state) = VoteState::deserialize(&account.userdata) {
+                    if cond(&vote_state) {
+                        return Some(vote_state);
+                    }
+                }
+                None
+            })
+            .collect()
     }
 
     /// Return the number of slots since the last epoch boundary.
@@ -1173,23 +1147,6 @@ mod tests {
         assert_eq!(register_ticks(&bank, ticks_per_epoch), (0, 1, 1));
     }
 
-    #[test]
-    fn test_bank_staked_nodes_at_epoch() {
-        let pubkey = Keypair::new().pubkey();
-        let bootstrap_tokens = 2;
-        let (genesis_block, _) = GenesisBlock::new_with_leader(2, pubkey, bootstrap_tokens);
-        let bank = Bank::new(&genesis_block);
-        let bank = Bank::new_from_parent(&Arc::new(bank));
-        let ticks_per_offset = bank.stakers_slot_offset * bank.ticks_per_slot();
-        register_ticks(&bank, ticks_per_offset);
-        assert_eq!(bank.slot_height(), bank.stakers_slot_offset);
-
-        let mut expected = HashMap::new();
-        expected.insert(pubkey, bootstrap_tokens - 1);
-        let bank = Bank::new_from_parent(&Arc::new(bank));
-        assert_eq!(bank.staked_nodes_at_slot(bank.slot_height()), expected);
-    }
-
     #[test]
     fn test_interleaving_locks() {
         let (genesis_block, mint_keypair) = GenesisBlock::new(3);

src/broadcast_service.rs

@@ -9,6 +9,7 @@ use crate::erasure::CodingGenerator;
 use crate::packet::index_blobs;
 use crate::result::{Error, Result};
 use crate::service::Service;
+use crate::staking_utils;
 use rayon::prelude::*;
 use solana_metrics::counter::Counter;
 use solana_metrics::{influxdb, submit};
@@ -189,7 +190,7 @@ impl BroadcastService {
             let mut broadcast_table = cluster_info
                 .read()
                 .unwrap()
-                .sorted_tvu_peers(&bank.staked_nodes());
+                .sorted_tvu_peers(&staking_utils::staked_nodes(&bank));
             // Layer 1, leader nodes are limited to the fanout size.
             broadcast_table.truncate(DATA_PLANE_FANOUT);
             inc_new_counter_info!("broadcast_service-num_peers", broadcast_table.len() + 1);

src/cluster_info.rs

@@ -22,6 +22,7 @@ use crate::crds_value::{CrdsValue, CrdsValueLabel, LeaderId, Vote};
 use crate::packet::{to_shared_blob, Blob, SharedBlob, BLOB_SIZE};
 use crate::result::Result;
 use crate::rpc_service::RPC_PORT;
+use crate::staking_utils;
 use crate::streamer::{BlobReceiver, BlobSender};
 use bincode::{deserialize, serialize};
 use core::cmp;
@@ -877,7 +878,7 @@ impl ClusterInfo {
                     let start = timestamp();
                     let stakes: HashMap<_, _> = match bank_forks {
                         Some(ref bank_forks) => {
-                            bank_forks.read().unwrap().working_bank().staked_nodes()
+                            staking_utils::staked_nodes(&bank_forks.read().unwrap().working_bank())
                         }
                         None => HashMap::new(),
                     };

src/leader_schedule_utils.rs

@@ -1,10 +1,11 @@
 use crate::leader_schedule::LeaderSchedule;
+use crate::staking_utils;
 use solana_runtime::bank::Bank;
 use solana_sdk::pubkey::Pubkey;
 
 /// Return the leader schedule for the given epoch.
 fn leader_schedule(epoch_height: u64, bank: &Bank) -> LeaderSchedule {
-    let stakes = bank.staked_nodes_at_epoch(epoch_height);
+    let stakes = staking_utils::staked_nodes_at_epoch(bank, epoch_height);
     let mut seed = [0u8; 32];
     seed[0..8].copy_from_slice(&epoch_height.to_le_bytes());
     let stakes: Vec<_> = stakes.into_iter().collect();
@@ -80,6 +81,7 @@ pub fn num_ticks_left_in_slot(bank: &Bank, tick_height: u64) -> u64 {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::staking_utils;
     use solana_sdk::genesis_block::GenesisBlock;
     use solana_sdk::signature::{Keypair, KeypairUtil};
 
@@ -89,7 +91,7 @@ mod tests {
         let (genesis_block, _mint_keypair) = GenesisBlock::new_with_leader(2, pubkey, 2);
         let bank = Bank::new(&genesis_block);
 
-        let ids_and_stakes: Vec<_> = bank.staked_nodes().into_iter().collect();
+        let ids_and_stakes: Vec<_> = staking_utils::staked_nodes(&bank).into_iter().collect();
         let seed = [0u8; 32];
         let leader_schedule =
             LeaderSchedule::new(&ids_and_stakes, seed, genesis_block.slots_per_epoch);

src/lib.rs

@@ -64,6 +64,7 @@ pub mod rpc_subscriptions;
 pub mod service;
 pub mod sigverify;
 pub mod sigverify_stage;
+pub mod staking_utils;
 pub mod storage_stage;
 pub mod streamer;
 pub mod test_tx;

src/retransmit_stage.rs

@@ -9,6 +9,7 @@ use crate::cluster_info::{
 use crate::packet::SharedBlob;
 use crate::result::{Error, Result};
 use crate::service::Service;
+use crate::staking_utils;
 use crate::streamer::BlobReceiver;
 use crate::window_service::WindowService;
 use solana_metrics::counter::Counter;
@@ -39,7 +40,7 @@ fn retransmit(
             .to_owned(),
     );
     let (neighbors, children) = compute_retransmit_peers(
-        &bank_forks.read().unwrap().working_bank().staked_nodes(),
+        &staking_utils::staked_nodes(&bank_forks.read().unwrap().working_bank()),
         cluster_info,
         DATA_PLANE_FANOUT,
         NEIGHBORHOOD_SIZE,

src/staking_utils.rs

@@ -0,0 +1,236 @@
+use hashbrown::HashMap;
+use solana_runtime::bank::Bank;
+use solana_sdk::pubkey::Pubkey;
+use solana_sdk::vote_program::VoteState;
+
+/// Looks through vote accounts, and finds the latest slot that has achieved
+/// supermajority lockout
+pub fn get_supermajority_slot(bank: &Bank, epoch_height: u64) -> Option<u64> {
+    // Find the amount of stake needed for supermajority
+    let stakes_and_lockouts = epoch_stakes_and_lockouts(bank, epoch_height);
+    let total_stake: u64 = stakes_and_lockouts.values().map(|s| s.0).sum();
+    let supermajority_stake = total_stake * 2 / 3;
+
+    // Filter out the states that don't have a max lockout
+    find_supermajority_slot(supermajority_stake, stakes_and_lockouts.values())
+}
+
+pub fn staked_nodes(bank: &Bank) -> HashMap<Pubkey, u64> {
+    staked_nodes_extractor(bank, |stake, _| stake)
+}
+
+/// Return the checkpointed stakes that should be used to generate a leader schedule.
+pub fn staked_nodes_at_epoch(bank: &Bank, epoch_height: u64) -> HashMap<Pubkey, u64> {
+    staked_nodes_at_epoch_extractor(bank, epoch_height, |stake, _| stake)
+}
+
+/// Return the checkpointed stakes that should be used to generate a leader schedule.
+/// state_extractor takes (stake, vote_state) and maps to an output.
+fn staked_nodes_at_epoch_extractor<F, T>(
+    bank: &Bank,
+    epoch_height: u64,
+    state_extractor: F,
+) -> HashMap<Pubkey, T>
+where
+    F: Fn(u64, &VoteState) -> T,
+{
+    let epoch_slot_height = epoch_height * bank.slots_per_epoch();
+    staked_nodes_at_slot_extractor(bank, epoch_slot_height, state_extractor)
+}
+
+/// Return the checkpointed stakes that should be used to generate a leader schedule.
+/// state_extractor takes (stake, vote_state) and maps to an output
+fn staked_nodes_at_slot_extractor<F, T>(
+    bank: &Bank,
+    current_slot_height: u64,
+    state_extractor: F,
+) -> HashMap<Pubkey, T>
+where
+    F: Fn(u64, &VoteState) -> T,
+{
+    let slot_height = current_slot_height.saturating_sub(bank.stakers_slot_offset());
+
+    let parents = bank.parents();
+    let mut banks = vec![bank];
+    banks.extend(parents.iter().map(|x| x.as_ref()));
+
+    let bank = banks
+        .iter()
+        .find(|bank| bank.id() <= slot_height)
+        .unwrap_or_else(|| banks.last().unwrap());
+
+    staked_nodes_extractor(bank, state_extractor)
+}
+
+/// Collect the node Pubkey and staker account balance for nodes
+/// that have non-zero balance in their corresponding staker accounts.
+/// state_extractor takes (stake, vote_state) and maps to an output
+fn staked_nodes_extractor<F, T>(bank: &Bank, state_extractor: F) -> HashMap<Pubkey, T>
+where
+    F: Fn(u64, &VoteState) -> T,
+{
+    bank.vote_states(|_| true)
+        .iter()
+        .filter_map(|state| {
+            let balance = bank.get_balance(&state.staker_id);
+            if balance > 0 {
+                Some((state.node_id, state_extractor(balance, &state)))
+            } else {
+                None
+            }
+        })
+        .collect()
+}
+
+fn epoch_stakes_and_lockouts(
+    bank: &Bank,
+    epoch_height: u64,
+) -> HashMap<Pubkey, (u64, Option<u64>)> {
+    staked_nodes_at_epoch_extractor(bank, epoch_height, |stake, state| (stake, state.root_slot))
+}
+
+fn find_supermajority_slot<'a, I>(supermajority_stake: u64, stakes_and_lockouts: I) -> Option<u64>
+where
+    I: Iterator<Item = &'a (u64, Option<u64>)>,
+{
+    // Filter out the states that don't have a max lockout
+    let mut stakes_and_lockouts: Vec<_> = stakes_and_lockouts
+        .filter_map(|(stake, slot)| slot.map(|s| (stake, s)))
+        .collect();
+
+    // Sort by the root slot, in descending order
+    stakes_and_lockouts.sort_unstable_by(|s1, s2| s1.1.cmp(&s2.1).reverse());
+
+    // Find if any slot has achieved sufficient votes for supermajority lockout
+    let mut total = 0;
+    for (stake, slot) in stakes_and_lockouts {
+        total += stake;
+        if total > supermajority_stake {
+            return Some(slot);
+        }
+    }
+
+    None
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::voting_keypair::tests as voting_keypair_tests;
+    use hashbrown::HashSet;
+    use solana_sdk::genesis_block::GenesisBlock;
+    use solana_sdk::hash::Hash;
+    use solana_sdk::signature::{Keypair, KeypairUtil};
+    use std::iter::FromIterator;
+    use std::sync::Arc;
+
+    fn register_ticks(bank: &Bank, n: u64) -> (u64, u64, u64) {
+        for _ in 0..n {
+            bank.register_tick(&Hash::default());
+        }
+        (bank.tick_index(), bank.slot_index(), bank.epoch_height())
+    }
+
+    #[test]
+    fn test_bank_staked_nodes_at_epoch() {
+        let pubkey = Keypair::new().pubkey();
+        let bootstrap_tokens = 2;
+        let (genesis_block, _) = GenesisBlock::new_with_leader(2, pubkey, bootstrap_tokens);
+        let bank = Bank::new(&genesis_block);
+        let bank = Bank::new_from_parent(&Arc::new(bank));
+        let ticks_per_offset = bank.stakers_slot_offset() * bank.ticks_per_slot();
+        register_ticks(&bank, ticks_per_offset);
+        assert_eq!(bank.slot_height(), bank.stakers_slot_offset());
+
+        let mut expected = HashMap::new();
+        expected.insert(pubkey, bootstrap_tokens - 1);
+        let bank = Bank::new_from_parent(&Arc::new(bank));
+        assert_eq!(
+            staked_nodes_at_slot_extractor(&bank, bank.slot_height(), |s, _| s),
+            expected
+        );
+    }
+
+    #[test]
+    fn test_epoch_stakes_and_lockouts() {
+        let validator = Keypair::new();
+        let voter = Keypair::new();
+
+        let (genesis_block, mint_keypair) = GenesisBlock::new(500);
+        let bank = Bank::new(&genesis_block);
+        let bank_voter = Keypair::new();
+
+        // Give the validator some stake
+        bank.transfer(
+            1,
+            &mint_keypair,
+            validator.pubkey(),
+            genesis_block.last_id(),
+        )
+        .unwrap();
+
+        voting_keypair_tests::new_vote_account_with_vote(&validator, &voter, &bank, 1, 0);
+        assert_eq!(bank.get_balance(&validator.pubkey()), 0);
+        // Validator has zero balance, so they get filtered out. Only the bootstrap leader
+        // created by the genesis block will get included
+        let expected: Vec<_> = epoch_stakes_and_lockouts(&bank, 0)
+            .values()
+            .cloned()
+            .collect();
+        assert_eq!(expected, vec![(1, None)]);
+
+        voting_keypair_tests::new_vote_account_with_vote(&mint_keypair, &bank_voter, &bank, 1, 0);
+
+        let result: HashSet<_> =
+            HashSet::from_iter(epoch_stakes_and_lockouts(&bank, 0).values().cloned());
+        let expected: HashSet<_> = HashSet::from_iter(vec![(1, None), (498, None)]);
+        assert_eq!(result, expected);
+    }
+
+    #[test]
+    fn test_find_supermajority_slot() {
+        let supermajority = 10;
+
+        let stakes_and_slots = vec![];
+        assert_eq!(
+            find_supermajority_slot(supermajority, stakes_and_slots.iter()),
+            None
+        );
+
+        let stakes_and_slots = vec![(5, None), (5, None)];
+        assert_eq!(
+            find_supermajority_slot(supermajority, stakes_and_slots.iter()),
+            None
+        );
+
+        let stakes_and_slots = vec![(5, None), (5, None), (9, Some(2))];
+        assert_eq!(
+            find_supermajority_slot(supermajority, stakes_and_slots.iter()),
+            None
+        );
+
+        let stakes_and_slots = vec![(5, None), (5, None), (9, Some(2)), (1, Some(3))];
+        assert_eq!(
+            find_supermajority_slot(supermajority, stakes_and_slots.iter()),
+            None
+        );
+
+        let stakes_and_slots = vec![(5, None), (5, None), (9, Some(2)), (2, Some(3))];
+        assert_eq!(
+            find_supermajority_slot(supermajority, stakes_and_slots.iter()),
+            Some(2)
+        );
+
+        let stakes_and_slots = vec![(9, Some(2)), (2, Some(3)), (9, None)];
+        assert_eq!(
+            find_supermajority_slot(supermajority, stakes_and_slots.iter()),
+            Some(2)
+        );
+
+        let stakes_and_slots = vec![(9, Some(2)), (2, Some(3)), (9, Some(3))];
+        assert_eq!(
+            find_supermajority_slot(supermajority, stakes_and_slots.iter()),
+            Some(3)
+        );
+    }
+}