expose stake directly from bank (#4281)
This commit is contained in:
@ -54,13 +54,13 @@ impl EpochStakes {
|
||||
&Pubkey::default(),
|
||||
)
|
||||
}
|
||||
pub fn new_from_stake_accounts(epoch: u64, accounts: &[(Pubkey, Account)]) -> Self {
|
||||
let stakes = accounts.iter().map(|(k, v)| (*k, v.lamports)).collect();
|
||||
pub fn new_from_stakes(epoch: u64, accounts: &[(Pubkey, (u64, Account))]) -> Self {
|
||||
let stakes = accounts.iter().map(|(k, (v, _))| (*k, *v)).collect();
|
||||
Self::new(epoch, stakes, &accounts[0].0)
|
||||
}
|
||||
pub fn new_from_bank(bank: &Bank, my_id: &Pubkey) -> Self {
|
||||
let bank_epoch = bank.get_epoch_and_slot_index(bank.slot()).0;
|
||||
let stakes = staking_utils::vote_account_balances_at_epoch(bank, bank_epoch)
|
||||
let stakes = staking_utils::vote_account_stakes_at_epoch(bank, bank_epoch)
|
||||
.expect("voting require a bank with stakes");
|
||||
Self::new(bank_epoch, stakes, my_id)
|
||||
}
|
||||
@ -105,10 +105,10 @@ impl Locktower {
|
||||
ancestors: &HashMap<u64, HashSet<u64>>,
|
||||
) -> HashMap<u64, StakeLockout>
|
||||
where
|
||||
F: Iterator<Item = (Pubkey, Account)>,
|
||||
F: Iterator<Item = (Pubkey, (u64, Account))>,
|
||||
{
|
||||
let mut stake_lockouts = HashMap::new();
|
||||
for (key, account) in vote_accounts {
|
||||
for (key, (_, account)) in vote_accounts {
|
||||
let lamports: u64 = *self.epoch_stakes.stakes.get(&key).unwrap_or(&0);
|
||||
if lamports == 0 {
|
||||
continue;
|
||||
@ -361,7 +361,7 @@ impl Locktower {
|
||||
fn initialize_lockouts_from_bank(bank: &Bank, current_epoch: u64) -> VoteState {
|
||||
let mut lockouts = VoteState::default();
|
||||
if let Some(iter) = staking_utils::node_staked_accounts_at_epoch(&bank, current_epoch) {
|
||||
for (delegate_id, _, account) in iter {
|
||||
for (delegate_id, (_, account)) in iter {
|
||||
if *delegate_id == bank.collector_id() {
|
||||
let state = VoteState::deserialize(&account.data).expect("votes");
|
||||
if lockouts.votes.len() < state.votes.len() {
|
||||
@ -378,8 +378,8 @@ impl Locktower {
|
||||
mod test {
|
||||
use super::*;
|
||||
|
||||
fn gen_accounts(stake_votes: &[(u64, &[u64])]) -> Vec<(Pubkey, Account)> {
|
||||
let mut accounts = vec![];
|
||||
fn gen_stakes(stake_votes: &[(u64, &[u64])]) -> Vec<(Pubkey, (u64, Account))> {
|
||||
let mut stakes = vec![];
|
||||
for (lamports, votes) in stake_votes {
|
||||
let mut account = Account::default();
|
||||
account.data = vec![0; 1024];
|
||||
@ -391,14 +391,14 @@ mod test {
|
||||
vote_state
|
||||
.serialize(&mut account.data)
|
||||
.expect("serialize state");
|
||||
accounts.push((Pubkey::new_rand(), account));
|
||||
stakes.push((Pubkey::new_rand(), (*lamports, account)));
|
||||
}
|
||||
accounts
|
||||
stakes
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_collect_vote_lockouts_no_epoch_stakes() {
|
||||
let accounts = gen_accounts(&[(1, &[0])]);
|
||||
let accounts = gen_stakes(&[(1, &[0])]);
|
||||
let epoch_stakes = EpochStakes::new_for_tests(2);
|
||||
let locktower = Locktower::new(epoch_stakes, 0, 0.67);
|
||||
let ancestors = vec![(1, vec![0].into_iter().collect()), (0, HashSet::new())]
|
||||
@ -411,8 +411,8 @@ mod test {
|
||||
#[test]
|
||||
fn test_collect_vote_lockouts_sums() {
|
||||
//two accounts voting for slot 0 with 1 token staked
|
||||
let accounts = gen_accounts(&[(1, &[0]), (1, &[0])]);
|
||||
let epoch_stakes = EpochStakes::new_from_stake_accounts(0, &accounts);
|
||||
let accounts = gen_stakes(&[(1, &[0]), (1, &[0])]);
|
||||
let epoch_stakes = EpochStakes::new_from_stakes(0, &accounts);
|
||||
let locktower = Locktower::new(epoch_stakes, 0, 0.67);
|
||||
let ancestors = vec![(1, vec![0].into_iter().collect()), (0, HashSet::new())]
|
||||
.into_iter()
|
||||
@ -426,8 +426,8 @@ mod test {
|
||||
fn test_collect_vote_lockouts_root() {
|
||||
let votes: Vec<u64> = (0..MAX_LOCKOUT_HISTORY as u64).into_iter().collect();
|
||||
//two accounts voting for slot 0 with 1 token staked
|
||||
let accounts = gen_accounts(&[(1, &votes), (1, &votes)]);
|
||||
let epoch_stakes = EpochStakes::new_from_stake_accounts(0, &accounts);
|
||||
let accounts = gen_stakes(&[(1, &votes), (1, &votes)]);
|
||||
let epoch_stakes = EpochStakes::new_from_stakes(0, &accounts);
|
||||
let mut locktower = Locktower::new(epoch_stakes, 0, 0.67);
|
||||
let mut ancestors = HashMap::new();
|
||||
for i in 0..(MAX_LOCKOUT_HISTORY + 1) {
|
||||
@ -754,13 +754,13 @@ mod test {
|
||||
let threshold_size = 0.67;
|
||||
let threshold_stake = (f64::ceil(total_stake as f64 * threshold_size)) as u64;
|
||||
let locktower_votes: Vec<u64> = (0..VOTE_THRESHOLD_DEPTH as u64).collect();
|
||||
let accounts = gen_accounts(&[
|
||||
let accounts = gen_stakes(&[
|
||||
(threshold_stake, &[(VOTE_THRESHOLD_DEPTH - 2) as u64]),
|
||||
(total_stake - threshold_stake, &locktower_votes[..]),
|
||||
]);
|
||||
|
||||
// Initialize locktower
|
||||
let stakes: HashMap<_, _> = accounts.iter().map(|(pk, a)| (*pk, a.lamports)).collect();
|
||||
let stakes: HashMap<_, _> = accounts.iter().map(|(pk, (s, _))| (*pk, *s)).collect();
|
||||
let epoch_stakes = EpochStakes::new(0, stakes, &Pubkey::default());
|
||||
let mut locktower = Locktower::new(epoch_stakes, VOTE_THRESHOLD_DEPTH, threshold_size);
|
||||
|
||||
|
@ -17,10 +17,10 @@ pub fn get_supermajority_slot(bank: &Bank, epoch_height: u64) -> Option<u64> {
|
||||
find_supermajority_slot(supermajority_stake, stakes_and_lockouts.iter())
|
||||
}
|
||||
|
||||
pub fn vote_account_balances(bank: &Bank) -> HashMap<Pubkey, u64> {
|
||||
pub fn vote_account_stakes(bank: &Bank) -> HashMap<Pubkey, u64> {
|
||||
let node_staked_accounts = node_staked_accounts(bank);
|
||||
node_staked_accounts
|
||||
.map(|(id, stake, _)| (id, stake))
|
||||
.map(|(id, (stake, _))| (id, stake))
|
||||
.collect()
|
||||
}
|
||||
|
||||
@ -34,12 +34,13 @@ pub fn delegated_stakes(bank: &Bank) -> HashMap<Pubkey, u64> {
|
||||
|
||||
/// At the specified epoch, collect the node account balance and vote states for nodes that
|
||||
/// have non-zero balance in their corresponding staking accounts
|
||||
pub fn vote_account_balances_at_epoch(
|
||||
pub fn vote_account_stakes_at_epoch(
|
||||
bank: &Bank,
|
||||
epoch_height: u64,
|
||||
) -> Option<HashMap<Pubkey, u64>> {
|
||||
let node_staked_accounts = node_staked_accounts_at_epoch(bank, epoch_height);
|
||||
node_staked_accounts.map(|epoch_state| epoch_state.map(|(id, stake, _)| (*id, stake)).collect())
|
||||
node_staked_accounts
|
||||
.map(|epoch_state| epoch_state.map(|(id, (stake, _))| (*id, *stake)).collect())
|
||||
}
|
||||
|
||||
/// At the specified epoch, collect the delegate account balance and vote states for delegates
|
||||
@ -52,21 +53,18 @@ pub fn delegated_stakes_at_epoch(bank: &Bank, epoch_height: u64) -> Option<HashM
|
||||
|
||||
/// Collect the node account balance and vote states for nodes have non-zero balance in
|
||||
/// their corresponding staking accounts
|
||||
fn node_staked_accounts(bank: &Bank) -> impl Iterator<Item = (Pubkey, u64, Account)> {
|
||||
bank.vote_accounts()
|
||||
.into_iter()
|
||||
.map(|(account_id, account)| (account_id, Bank::read_balance(&account), account))
|
||||
fn node_staked_accounts(bank: &Bank) -> impl Iterator<Item = (Pubkey, (u64, Account))> {
|
||||
bank.vote_accounts().into_iter()
|
||||
}
|
||||
|
||||
pub fn node_staked_accounts_at_epoch(
|
||||
bank: &Bank,
|
||||
epoch_height: u64,
|
||||
) -> Option<impl Iterator<Item = (&Pubkey, u64, &Account)>> {
|
||||
) -> Option<impl Iterator<Item = (&Pubkey, &(u64, Account))>> {
|
||||
bank.epoch_vote_accounts(epoch_height).map(|vote_accounts| {
|
||||
vote_accounts
|
||||
.into_iter()
|
||||
.filter(|(account_id, account)| filter_no_delegate(account_id, account))
|
||||
.map(|(account_id, account)| (account_id, Bank::read_balance(&account), account))
|
||||
.filter(|(account_id, (_, account))| filter_no_delegate(account_id, account))
|
||||
})
|
||||
}
|
||||
|
||||
@ -77,12 +75,12 @@ fn filter_no_delegate(account_id: &Pubkey, account: &Account) -> bool {
|
||||
}
|
||||
|
||||
fn to_vote_state(
|
||||
node_staked_accounts: impl Iterator<Item = (impl Borrow<Pubkey>, u64, impl Borrow<Account>)>,
|
||||
node_staked_accounts: impl Iterator<Item = (impl Borrow<Pubkey>, impl Borrow<(u64, Account)>)>,
|
||||
) -> impl Iterator<Item = (u64, VoteState)> {
|
||||
node_staked_accounts.filter_map(|(_, stake, account)| {
|
||||
VoteState::deserialize(&account.borrow().data)
|
||||
node_staked_accounts.filter_map(|(_, stake_account)| {
|
||||
VoteState::deserialize(&stake_account.borrow().1.data)
|
||||
.ok()
|
||||
.map(|vote_state| (stake, vote_state))
|
||||
.map(|vote_state| (stake_account.borrow().0, vote_state))
|
||||
})
|
||||
}
|
||||
|
||||
@ -158,17 +156,17 @@ pub mod tests {
|
||||
|
||||
// Epoch doesn't exist
|
||||
let mut expected = HashMap::new();
|
||||
assert_eq!(vote_account_balances_at_epoch(&bank, 10), None);
|
||||
assert_eq!(vote_account_stakes_at_epoch(&bank, 10), None);
|
||||
|
||||
// First epoch has the bootstrap leader
|
||||
expected.insert(voting_keypair.pubkey(), BOOTSTRAP_LEADER_LAMPORTS);
|
||||
let expected = Some(expected);
|
||||
assert_eq!(vote_account_balances_at_epoch(&bank, 0), expected);
|
||||
assert_eq!(vote_account_stakes_at_epoch(&bank, 0), expected);
|
||||
|
||||
// Second epoch carries same information
|
||||
let bank = new_from_parent(&Arc::new(bank), 1);
|
||||
assert_eq!(vote_account_balances_at_epoch(&bank, 0), expected);
|
||||
assert_eq!(vote_account_balances_at_epoch(&bank, 1), expected);
|
||||
assert_eq!(vote_account_stakes_at_epoch(&bank, 0), expected);
|
||||
assert_eq!(vote_account_stakes_at_epoch(&bank, 1), expected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
Reference in New Issue
Block a user