Store and compute node/stake state in EpochStakes struct (#8958)

* Store and compute needed bank state in EpochStakes struct

ci/test-checks.sh

@@ -22,7 +22,7 @@ _ cargo +"$rust_stable" clippy --all --exclude solana-sdk-c -- --deny=warnings
 _ cargo +"$rust_stable" clippy --manifest-path sdk-c/Cargo.toml -- --deny=warnings
 
 _ cargo +"$rust_stable" audit --version
-_ cargo +"$rust_stable" audit --ignore RUSTSEC-2020-0002
+_ cargo +"$rust_stable" audit --ignore RUSTSEC-2020-0002 --ignore RUSTSEC-2020-0006
 _ ci/nits.sh
 _ ci/order-crates-for-publishing.py
 _ docs/build.sh

ledger/src/snapshot_utils.rs

@@ -31,7 +31,7 @@ pub const TAR_SNAPSHOTS_DIR: &str = "snapshots";
 pub const TAR_ACCOUNTS_DIR: &str = "accounts";
 pub const TAR_VERSION_FILE: &str = "version";
 
-pub const SNAPSHOT_VERSION: &str = "1.0.0";
+pub const SNAPSHOT_VERSION: &str = "1.1.0";
 
 #[derive(PartialEq, Ord, Eq, Debug)]
 pub struct SlotSnapshotPaths {

runtime/src/bank.rs

@@ -6,6 +6,7 @@ use crate::{
     accounts::{Accounts, TransactionAccounts, TransactionLoadResult, TransactionLoaders},
     accounts_db::{AccountsDBSerialize, ErrorCounters, SnapshotStorage, SnapshotStorages},
     blockhash_queue::BlockhashQueue,
+    epoch_stakes::{EpochStakes, NodeVoteAccounts},
     message_processor::{MessageProcessor, ProcessInstruction},
     nonce_utils,
     rent_collector::RentCollector,
@@ -323,7 +324,7 @@ pub struct Bank {
 
     /// staked nodes on epoch boundaries, saved off when a bank.slot() is at
     ///   a leader schedule calculation boundary
-    epoch_stakes: HashMap<Epoch, Stakes>,
+    epoch_stakes: HashMap<Epoch, EpochStakes>,
 
     /// A boolean reflecting whether any entries were recorded into the PoH
     /// stream for the slot == self.slot
@@ -380,7 +381,8 @@ impl Bank {
         {
             let stakes = bank.stakes.read().unwrap();
             for epoch in 0..=bank.get_leader_schedule_epoch(bank.slot) {
-                bank.epoch_stakes.insert(epoch, stakes.clone());
+                bank.epoch_stakes
+                    .insert(epoch, EpochStakes::new(&stakes, epoch));
             }
             bank.update_stake_history(None);
         }
@@ -592,8 +594,24 @@ impl Bank {
                 epoch >= leader_schedule_epoch.saturating_sub(MAX_LEADER_SCHEDULE_STAKES)
             });
 
+            let vote_stakes: HashMap<_, _> = self
+                .stakes
+                .read()
+                .unwrap()
+                .vote_accounts()
+                .iter()
+                .map(|(epoch, (stake, _))| (*epoch, *stake))
+                .collect();
+            let new_epoch_stakes =
+                EpochStakes::new(&self.stakes.read().unwrap(), leader_schedule_epoch);
+            info!(
+                "new epoch stakes, epoch: {}, stakes: {:#?}, total_stake: {}",
+                leader_schedule_epoch,
+                vote_stakes,
+                new_epoch_stakes.total_stake(),
+            );
             self.epoch_stakes
-                .insert(leader_schedule_epoch, self.stakes.read().unwrap().clone());
+                .insert(leader_schedule_epoch, new_epoch_stakes);
         }
     }
 
@@ -2062,10 +2080,55 @@ impl Bank {
         self.stakes.read().unwrap().vote_accounts().clone()
     }
 
+    /// Get the EpochStakes for a given epoch
+    pub fn epoch_stakes(&self, epoch: Epoch) -> Option<&EpochStakes> {
+        self.epoch_stakes.get(&epoch)
+    }
+
     /// vote accounts for the specific epoch along with the stake
     ///   attributed to each account
     pub fn epoch_vote_accounts(&self, epoch: Epoch) -> Option<&HashMap<Pubkey, (u64, Account)>> {
-        self.epoch_stakes.get(&epoch).map(Stakes::vote_accounts)
+        self.epoch_stakes
+            .get(&epoch)
+            .map(|epoch_stakes| Stakes::vote_accounts(epoch_stakes.stakes()))
+    }
+
+    /// Get the fixed authorized voter for the given vote account for the
+    /// current epoch
+    pub fn epoch_authorized_voter(&self, vote_account: &Pubkey) -> Option<&Pubkey> {
+        self.epoch_stakes
+            .get(&self.epoch)
+            .expect("Epoch stakes for bank's own epoch must exist")
+            .epoch_authorized_voters()
+            .get(vote_account)
+    }
+
+    /// Get the fixed set of vote accounts for the given node id for the
+    /// current epoch
+    pub fn epoch_vote_accounts_for_node_id(&self, node_id: &Pubkey) -> Option<&NodeVoteAccounts> {
+        self.epoch_stakes
+            .get(&self.epoch)
+            .expect("Epoch stakes for bank's own epoch must exist")
+            .node_id_to_vote_accounts()
+            .get(node_id)
+    }
+
+    /// Get the fixed total stake of all vote accounts for current epoch
+    pub fn total_epoch_stake(&self) -> u64 {
+        self.epoch_stakes
+            .get(&self.epoch)
+            .expect("Epoch stakes for bank's own epoch must exist")
+            .total_stake()
+    }
+
+    /// Get the fixed stake of the given vote account for the current epoch
+    pub fn epoch_vote_account_stake(&self, voting_pubkey: &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)
+            .map(|(stake, _)| stake)
+            .unwrap_or(&0)
     }
 
     /// given a slot, return the epoch and offset into the epoch this slot falls

runtime/src/epoch_stakes.rs

@@ -0,0 +1,208 @@
+use crate::stakes::Stakes;
+use serde::{Deserialize, Serialize};
+use solana_sdk::{account::Account, clock::Epoch, pubkey::Pubkey};
+use solana_vote_program::vote_state::VoteState;
+use std::{collections::HashMap, sync::Arc};
+
+pub type NodeIdToVoteAccounts = HashMap<Pubkey, NodeVoteAccounts>;
+pub type EpochAuthorizedVoters = HashMap<Pubkey, Pubkey>;
+
+#[derive(Clone, Serialize, Deserialize, Default, PartialEq)]
+pub struct NodeVoteAccounts {
+    pub vote_accounts: Vec<Pubkey>,
+    pub total_stake: u64,
+}
+
+#[derive(Clone, Serialize, Deserialize)]
+pub struct EpochStakes {
+    stakes: Arc<Stakes>,
+    total_stake: u64,
+    node_id_to_vote_accounts: Arc<NodeIdToVoteAccounts>,
+    epoch_authorized_voters: Arc<EpochAuthorizedVoters>,
+}
+
+impl EpochStakes {
+    pub fn new(stakes: &Stakes, leader_schedule_epoch: Epoch) -> Self {
+        let epoch_vote_accounts = Stakes::vote_accounts(stakes);
+        let (total_stake, node_id_to_vote_accounts, epoch_authorized_voters) =
+            Self::parse_epoch_vote_accounts(&epoch_vote_accounts, leader_schedule_epoch);
+        Self {
+            stakes: Arc::new(stakes.clone()),
+            total_stake,
+            node_id_to_vote_accounts: Arc::new(node_id_to_vote_accounts),
+            epoch_authorized_voters: Arc::new(epoch_authorized_voters),
+        }
+    }
+
+    pub fn stakes(&self) -> &Stakes {
+        &self.stakes
+    }
+
+    pub fn total_stake(&self) -> u64 {
+        self.total_stake
+    }
+
+    pub fn node_id_to_vote_accounts(&self) -> &Arc<NodeIdToVoteAccounts> {
+        &self.node_id_to_vote_accounts
+    }
+
+    pub fn epoch_authorized_voters(&self) -> &Arc<EpochAuthorizedVoters> {
+        &self.epoch_authorized_voters
+    }
+
+    fn parse_epoch_vote_accounts(
+        epoch_vote_accounts: &HashMap<Pubkey, (u64, Account)>,
+        leader_schedule_epoch: Epoch,
+    ) -> (u64, NodeIdToVoteAccounts, EpochAuthorizedVoters) {
+        let mut node_id_to_vote_accounts: NodeIdToVoteAccounts = HashMap::new();
+        let total_stake = epoch_vote_accounts
+            .iter()
+            .map(|(_, (stake, _))| stake)
+            .sum();
+        let epoch_authorized_voters = epoch_vote_accounts
+            .iter()
+            .filter_map(|(key, (stake, account))| {
+                let vote_state = VoteState::from(&account);
+                if vote_state.is_none() {
+                    datapoint_warn!(
+                        "parse_epoch_vote_accounts",
+                        (
+                            "warn",
+                            format!("Unable to get vote_state from account {}", key),
+                            String
+                        ),
+                    );
+                    return None;
+                }
+                let vote_state = vote_state.unwrap();
+                if *stake > 0 {
+                    // Read out the authorized voters
+                    let authorized_voter = vote_state
+                        .authorized_voters()
+                        .get_authorized_voter(leader_schedule_epoch)
+                        .expect("Authorized voter for current epoch must be known");
+
+                    let node_vote_accounts = node_id_to_vote_accounts
+                        .entry(vote_state.node_pubkey)
+                        .or_default();
+
+                    node_vote_accounts.total_stake += stake;
+                    node_vote_accounts.vote_accounts.push(*key);
+
+                    Some((*key, authorized_voter))
+                } else {
+                    None
+                }
+            })
+            .collect();
+        (
+            total_stake,
+            node_id_to_vote_accounts,
+            epoch_authorized_voters,
+        )
+    }
+}
+
+#[cfg(test)]
+pub(crate) mod tests {
+    use super::*;
+    use solana_vote_program::vote_state::create_account_with_authorized;
+    use std::iter;
+
+    struct VoteAccountInfo {
+        vote_account: Pubkey,
+        account: Account,
+        authorized_voter: Pubkey,
+    }
+
+    #[test]
+    fn test_parse_epoch_vote_accounts() {
+        let stake_per_account = 100;
+        let num_vote_accounts_per_node = 2;
+        // Create some vote accounts for each pubkey
+        let vote_accounts_map: HashMap<Pubkey, Vec<VoteAccountInfo>> = (0..10)
+            .map(|_| {
+                let node_id = Pubkey::new_rand();
+                (
+                    node_id,
+                    iter::repeat_with(|| {
+                        let authorized_voter = Pubkey::new_rand();
+                        VoteAccountInfo {
+                            vote_account: Pubkey::new_rand(),
+                            account: create_account_with_authorized(
+                                &node_id,
+                                &authorized_voter,
+                                &node_id,
+                                0,
+                                100,
+                            ),
+                            authorized_voter,
+                        }
+                    })
+                    .take(num_vote_accounts_per_node)
+                    .collect(),
+                )
+            })
+            .collect();
+
+        let expected_authorized_voters: HashMap<_, _> = vote_accounts_map
+            .iter()
+            .flat_map(|(_, vote_accounts)| {
+                vote_accounts
+                    .iter()
+                    .map(|v| (v.vote_account, v.authorized_voter))
+            })
+            .collect();
+
+        let expected_node_id_to_vote_accounts: HashMap<_, _> = vote_accounts_map
+            .iter()
+            .map(|(node_pubkey, vote_accounts)| {
+                let mut vote_accounts = vote_accounts
+                    .iter()
+                    .map(|v| (v.vote_account))
+                    .collect::<Vec<_>>();
+                vote_accounts.sort();
+                let node_vote_accounts = NodeVoteAccounts {
+                    vote_accounts,
+                    total_stake: stake_per_account * num_vote_accounts_per_node as u64,
+                };
+                (*node_pubkey, node_vote_accounts)
+            })
+            .collect();
+
+        // Create and process the vote accounts
+        let epoch_vote_accounts: HashMap<_, _> = vote_accounts_map
+            .iter()
+            .flat_map(|(_, vote_accounts)| {
+                vote_accounts
+                    .iter()
+                    .map(|v| (v.vote_account, (stake_per_account, v.account.clone())))
+            })
+            .collect();
+
+        let (total_stake, mut node_id_to_vote_accounts, epoch_authorized_voters) =
+            EpochStakes::parse_epoch_vote_accounts(&epoch_vote_accounts, 0);
+
+        // Verify the results
+        node_id_to_vote_accounts
+            .iter_mut()
+            .for_each(|(_, node_vote_accounts)| node_vote_accounts.vote_accounts.sort());
+
+        assert!(
+            node_id_to_vote_accounts.len() == expected_node_id_to_vote_accounts.len()
+                && node_id_to_vote_accounts
+                    .iter()
+                    .all(|(k, v)| expected_node_id_to_vote_accounts.get(k).unwrap() == v)
+        );
+        assert!(
+            epoch_authorized_voters.len() == expected_authorized_voters.len()
+                && epoch_authorized_voters
+                    .iter()
+                    .all(|(k, v)| expected_authorized_voters.get(k).unwrap() == v)
+        );
+        assert_eq!(
+            total_stake,
+            vote_accounts_map.len() as u64 * num_vote_accounts_per_node as u64 * 100
+        );
+    }
+}

runtime/src/lib.rs

@@ -6,6 +6,7 @@ pub mod bank;
 pub mod bank_client;
 mod blockhash_queue;
 pub mod bloom;
+pub mod epoch_stakes;
 pub mod genesis_utils;
 pub mod loader_utils;
 pub mod message_processor;