caches StakeAccount instead of Delegation in Stakes

The commit makes values in stake_delegations map in Stakes struct
generic. Stakes<Delegation> is equivalent to the old code and is used
for backward compatibility in BankFieldsTo{Serialize,Deserialize}.

But banks cache Stakes<StakeAccount> which includes the entire stake
account and StakeState deserialized from account. Doing so, will remove
the need to load stake account from accounts-db when working with
stake-delegations.

core/tests/snapshots.rs

@@ -65,7 +65,7 @@ mod tests {
             accounts_index::AccountSecondaryIndexes,
             bank::{Bank, BankSlotDelta},
             bank_forks::BankForks,
-            genesis_utils::{create_genesis_config, GenesisConfigInfo},
+            genesis_utils::{create_genesis_config_with_leader, GenesisConfigInfo},
             snapshot_archive_info::FullSnapshotArchiveInfo,
             snapshot_config::SnapshotConfig,
             snapshot_package::{
@@ -124,7 +124,15 @@ mod tests {
             let accounts_dir = TempDir::new().unwrap();
             let bank_snapshots_dir = TempDir::new().unwrap();
             let snapshot_archives_dir = TempDir::new().unwrap();
-            let mut genesis_config_info = create_genesis_config(10_000);
+            // validator_stake_lamports should be non-zero otherwise stake
+            // account will not be stored in accounts-db but still cached in
+            // bank stakes which results in mismatch when banks are loaded from
+            // snapshots.
+            let mut genesis_config_info = create_genesis_config_with_leader(
+                10_000,                          // mint_lamports
+                &solana_sdk::pubkey::new_rand(), // validator_pubkey
+                1,                               // validator_stake_lamports
+            );
             genesis_config_info.genesis_config.cluster_type = cluster_type;
             let bank0 = Bank::new_with_paths_for_tests(
                 &genesis_config_info.genesis_config,

runtime/src/bank.rs

@@ -128,6 +128,7 @@ use {
         signature::{Keypair, Signature},
         slot_hashes::SlotHashes,
         slot_history::SlotHistory,
+        stake::state::Delegation,
         system_transaction,
         sysvar::{self, Sysvar, SysvarId},
         timing::years_as_slots,
@@ -938,7 +939,7 @@ pub(crate) struct BankFieldsToDeserialize {
     pub(crate) rent_collector: RentCollector,
     pub(crate) epoch_schedule: EpochSchedule,
     pub(crate) inflation: Inflation,
-    pub(crate) stakes: Stakes,
+    pub(crate) stakes: Stakes<Delegation>,
     pub(crate) epoch_stakes: HashMap<Epoch, EpochStakes>,
     pub(crate) is_delta: bool,
     pub(crate) accounts_data_len: u64,
@@ -1497,10 +1498,12 @@ impl Bank {
         // genesis needs stakes for all epochs up to the epoch implied by
         //  slot = 0 and genesis configuration
         {
-            let stakes = bank.stakes_cache.stakes();
+            let stakes = bank.stakes_cache.stakes().clone();
+            let stakes = Stakes::<Delegation>::try_from(stakes).unwrap();
+            let stakes = Arc::new(stakes);
             for epoch in 0..=bank.get_leader_schedule_epoch(bank.slot) {
                 bank.epoch_stakes
-                    .insert(epoch, EpochStakes::new(&stakes, epoch));
+                    .insert(epoch, EpochStakes::new(stakes.clone(), epoch));
             }
             bank.update_stake_history(None);
         }
@@ -1983,6 +1986,23 @@ impl Bank {
         debug_do_not_add_builtins: bool,
         accounts_data_len: u64,
     ) -> Self {
+        let now = Instant::now();
+        let ancestors = Ancestors::from(&fields.ancestors);
+        // For backward compatibility, we can only serialize and deserialize
+        // Stakes<Delegation> in BankFieldsTo{Serialize,Deserialize}. But Bank
+        // caches Stakes<StakeAccount>. Below Stakes<StakeAccount> is obtained
+        // from Stakes<Delegation> by reading the full account state from
+        // accounts-db. Note that it is crucial that these accounts are loaded
+        // at the right slot and match precisely with serialized Delegations.
+        let stakes = Stakes::<StakeAccount>::new(&fields.stakes, |pubkey| {
+            let (account, _slot) = bank_rc.accounts.load_with_fixed_root(&ancestors, pubkey)?;
+            Some(account)
+        })
+        .expect(
+            "Stakes cache is inconsistent with accounts-db. This can indicate \
+            a corrupted snapshot or bugs in cached accounts or accounts-db.",
+        );
+        let stakes_accounts_load_duration = now.elapsed();
         fn new<T: Default>() -> T {
             T::default()
         }
@@ -1992,7 +2012,7 @@ impl Bank {
             rc: bank_rc,
             src: new(),
             blockhash_queue: RwLock::new(fields.blockhash_queue),
-            ancestors: Ancestors::from(&fields.ancestors),
+            ancestors,
             hash: RwLock::new(fields.hash),
             parent_hash: fields.parent_hash,
             parent_slot: fields.parent_slot,
@@ -2023,7 +2043,7 @@ impl Bank {
             rent_collector: Self::get_rent_collector_from(&fields.rent_collector, fields.epoch),
             epoch_schedule: fields.epoch_schedule,
             inflation: Arc::new(RwLock::new(fields.inflation)),
-            stakes_cache: StakesCache::new(fields.stakes),
+            stakes_cache: StakesCache::new(stakes),
             epoch_stakes: fields.epoch_stakes,
             is_delta: AtomicBool::new(fields.is_delta),
             builtin_programs: new(),
@@ -2106,6 +2126,11 @@ impl Bank {
                 accounts_data_len as i64,
                 i64
             ),
+            (
+                "stakes_accounts_load_duration_us",
+                stakes_accounts_load_duration.as_micros(),
+                i64
+            ),
         );
         bank
     }
@@ -2378,9 +2403,9 @@ impl Bank {
             self.epoch_stakes.retain(|&epoch, _| {
                 epoch >= leader_schedule_epoch.saturating_sub(MAX_LEADER_SCHEDULE_STAKES)
             });
-
+            let stakes = self.stakes_cache.stakes().clone();
-            let new_epoch_stakes =
+            let stakes = Stakes::<Delegation>::try_from(stakes).unwrap();
-                EpochStakes::new(&self.stakes_cache.stakes(), leader_schedule_epoch);
+            let new_epoch_stakes = EpochStakes::new(Arc::new(stakes), leader_schedule_epoch);
             {
                 let vote_stakes: HashMap<_, _> = self
                     .stakes_cache
@@ -2612,7 +2637,8 @@ impl Bank {
         thread_pool.install(|| {
             stake_delegations
                 .into_par_iter()
-                .for_each(|(stake_pubkey, delegation)| {
+                .for_each(|(stake_pubkey, stake_account)| {
+                    let delegation = stake_account.delegation().unwrap();
                     let vote_pubkey = &delegation.voter_pubkey;
                     if invalid_vote_keys.contains_key(vote_pubkey) {
                         return;
@@ -2683,7 +2709,7 @@ impl Bank {
                         reward_calc_tracer(&RewardCalculationEvent::Staking(
                             stake_pubkey,
                             &InflationPointCalculationEvent::Delegation(
-                                *delegation,
+                                delegation,
                                 solana_vote_program::id(),
                             ),
                         ));

runtime/src/epoch_stakes.rs

@@ -1,7 +1,7 @@
 use {
     crate::{stakes::Stakes, vote_account::VoteAccountsHashMap},
     serde::{Deserialize, Serialize},
-    solana_sdk::{clock::Epoch, pubkey::Pubkey},
+    solana_sdk::{clock::Epoch, pubkey::Pubkey, stake::state::Delegation},
     std::{collections::HashMap, sync::Arc},
 };
 
@@ -16,26 +16,26 @@ pub struct NodeVoteAccounts {
 
 #[derive(Clone, Debug, Serialize, Deserialize, AbiExample, PartialEq)]
 pub struct EpochStakes {
-    stakes: Arc<Stakes>,
+    stakes: Arc<Stakes<Delegation>>,
     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 {
+    pub(crate) fn new(stakes: Arc<Stakes<Delegation>>, leader_schedule_epoch: Epoch) -> Self {
         let epoch_vote_accounts = stakes.vote_accounts();
         let (total_stake, node_id_to_vote_accounts, epoch_authorized_voters) =
             Self::parse_epoch_vote_accounts(epoch_vote_accounts.as_ref(), leader_schedule_epoch);
         Self {
-            stakes: Arc::new(stakes.clone()),
+            stakes,
             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 {
+    pub fn stakes(&self) -> &Stakes<Delegation> {
         &self.stakes
     }
 

runtime/src/genesis_utils.rs

@@ -75,7 +75,15 @@ pub struct GenesisConfigInfo {
 }
 
 pub fn create_genesis_config(mint_lamports: u64) -> GenesisConfigInfo {
-    create_genesis_config_with_leader(mint_lamports, &solana_sdk::pubkey::new_rand(), 0)
+    // Note that zero lamports for validator stake will result in stake account
+    // not being stored in accounts-db but still cached in bank stakes. This
+    // causes discrepancy between cached stakes accounts in bank and
+    // accounts-db which in particular will break snapshots test.
+    create_genesis_config_with_leader(
+        mint_lamports,
+        &solana_sdk::pubkey::new_rand(), // validator_pubkey
+        0,                               // validator_stake_lamports
+    )
 }
 
 pub fn create_genesis_config_with_vote_accounts(

runtime/src/serde_snapshot/newer.rs

@@ -4,8 +4,9 @@ use {
         utils::{serialize_iter_as_map, serialize_iter_as_seq},
         *,
     },
-    crate::{ancestors::AncestorsForSerialization, stakes::StakesCache},
+    crate::ancestors::AncestorsForSerialization,
     solana_measure::measure::Measure,
+    solana_sdk::stake::state::Delegation,
     std::{cell::RefCell, collections::HashSet, sync::RwLock},
 };
 
@@ -51,7 +52,7 @@ struct DeserializableVersionedBank {
     rent_collector: RentCollector,
     epoch_schedule: EpochSchedule,
     inflation: Inflation,
-    stakes: Stakes,
+    stakes: Stakes<Delegation>,
     #[allow(dead_code)]
     unused_accounts: UnusedAccounts,
     epoch_stakes: HashMap<Epoch, EpochStakes>,
@@ -128,7 +129,7 @@ struct SerializableVersionedBank<'a> {
     rent_collector: RentCollector,
     epoch_schedule: EpochSchedule,
     inflation: Inflation,
-    stakes: &'a StakesCache,
+    stakes: Stakes<Delegation>,
     unused_accounts: UnusedAccounts,
     epoch_stakes: &'a HashMap<Epoch, EpochStakes>,
     is_delta: bool,
@@ -165,7 +166,10 @@ impl<'a> From<crate::bank::BankFieldsToSerialize<'a>> for SerializableVersionedB
             rent_collector: rhs.rent_collector,
             epoch_schedule: rhs.epoch_schedule,
             inflation: rhs.inflation,
-            stakes: rhs.stakes,
+            stakes: {
+                let stakes = rhs.stakes.stakes().clone();
+                Stakes::<Delegation>::try_from(stakes).unwrap()
+            },
             unused_accounts: UnusedAccounts::default(),
             epoch_stakes: rhs.epoch_stakes,
             is_delta: rhs.is_delta,
@@ -306,7 +310,6 @@ impl<'a> TypeContext<'a> for Context {
         let rhs = bank_fields;
         let blockhash_queue = RwLock::new(rhs.blockhash_queue.clone());
         let hard_forks = RwLock::new(rhs.hard_forks.clone());
-        let stakes_cache = StakesCache::new(rhs.stakes.clone());
         let bank = SerializableVersionedBank {
             blockhash_queue: &blockhash_queue,
             ancestors: &rhs.ancestors,
@@ -336,7 +339,7 @@ impl<'a> TypeContext<'a> for Context {
             rent_collector: rhs.rent_collector,
             epoch_schedule: rhs.epoch_schedule,
             inflation: rhs.inflation,
-            stakes: &stakes_cache,
+            stakes: rhs.stakes,
             unused_accounts: UnusedAccounts::default(),
             epoch_stakes: &rhs.epoch_stakes,
             is_delta: rhs.is_delta,

runtime/src/serde_snapshot/tests.rs

@@ -367,7 +367,7 @@ mod test_bank_serialize {
 
     // This some what long test harness is required to freeze the ABI of
     // Bank's serialization due to versioned nature
-    #[frozen_abi(digest = "DnT7iwf29YvGxLTt2iLYPmixwP13LXvusVVKpXGgizhc")]
+    #[frozen_abi(digest = "2RqbikYfG4rXV8gtqsKCNSn6g4g6nfkegXDRe1G6RGFQ")]
     #[derive(Serialize, AbiExample)]
     pub struct BankAbiTestWrapperNewer {
         #[serde(serialize_with = "wrapper_newer")]

runtime/src/stake_account.rs

@@ -1,10 +1,12 @@
+#[cfg(RUSTC_WITH_SPECIALIZATION)]
+use solana_frozen_abi::abi_example::AbiExample;
 use {
     solana_sdk::{
         account::{AccountSharedData, ReadableAccount},
         account_utils::StateMut,
         instruction::InstructionError,
         pubkey::Pubkey,
-        stake::state::StakeState,
+        stake::state::{Delegation, StakeState},
     },
     thiserror::Error,
 };
@@ -22,10 +24,20 @@ pub(crate) enum Error {
 }
 
 impl StakeAccount {
+    #[inline]
+    pub(crate) fn lamports(&self) -> u64 {
+        self.0.lamports()
+    }
+
     #[inline]
     pub(crate) fn stake_state(&self) -> &StakeState {
         &self.1
     }
+
+    #[inline]
+    pub(crate) fn delegation(&self) -> Option<Delegation> {
+        self.1.delegation()
+    }
 }
 
 impl TryFrom<AccountSharedData> for StakeAccount {
@@ -46,3 +58,19 @@ impl From<StakeAccount> for (AccountSharedData, StakeState) {
         (stake_account.0, stake_account.1)
     }
 }
+
+#[cfg(RUSTC_WITH_SPECIALIZATION)]
+impl AbiExample for StakeAccount {
+    fn example() -> Self {
+        use solana_sdk::{
+            account::Account,
+            stake::state::{Meta, Stake},
+        };
+        let stake_state = StakeState::Stake(Meta::example(), Stake::example());
+        let mut account = Account::example();
+        account.data.resize(196, 0u8);
+        account.owner = solana_stake_program::id();
+        let _ = account.set_state(&stake_state).unwrap();
+        Self::try_from(AccountSharedData::from(account)).unwrap()
+    }
+}

runtime/src/stakes.rs

@@ -2,6 +2,7 @@
 //! node stakes
 use {
     crate::{
+        stake_account::{self, StakeAccount},
         stake_history::StakeHistory,
         vote_account::{VoteAccount, VoteAccounts},
     },
@@ -14,20 +15,27 @@ use {
         account::{AccountSharedData, ReadableAccount},
         clock::{Epoch, Slot},
         pubkey::Pubkey,
-        stake::{
+        stake::state::{Delegation, StakeActivationStatus},
-            self,
-            state::{Delegation, StakeActivationStatus, StakeState},
     },
-    },
-    solana_stake_program::stake_state,
     solana_vote_program::vote_state::VoteState,
     std::{
         collections::HashMap,
         ops::Add,
         sync::{Arc, RwLock, RwLockReadGuard},
     },
+    thiserror::Error,
 };
 
+#[derive(Debug, Error)]
+pub(crate) enum Error {
+    #[error("Invalid delegation: {0}")]
+    InvalidDelegation(Pubkey),
+    #[error(transparent)]
+    InvalidStakeAccount(#[from] stake_account::Error),
+    #[error("Stake account not found: {0}")]
+    StakeAccountNotFound(Pubkey),
+}
+
 #[derive(Debug, Clone, PartialEq, ToPrimitive)]
 pub enum InvalidCacheEntryReason {
     Missing,
@@ -35,24 +43,18 @@ pub enum InvalidCacheEntryReason {
     WrongOwner,
 }
 
-#[derive(Default, Debug, Deserialize, Serialize, AbiExample)]
+#[derive(Default, Debug, AbiExample)]
-pub struct StakesCache(RwLock<Stakes>);
+pub(crate) struct StakesCache(RwLock<Stakes<StakeAccount>>);
 
 impl StakesCache {
-    pub fn new(stakes: Stakes) -> Self {
+    pub(crate) fn new(stakes: Stakes<StakeAccount>) -> Self {
         Self(RwLock::new(stakes))
     }
 
-    pub fn stakes(&self) -> RwLockReadGuard<Stakes> {
+    pub(crate) fn stakes(&self) -> RwLockReadGuard<Stakes<StakeAccount>> {
         self.0.read().unwrap()
     }
 
-    pub fn is_stake(account: &AccountSharedData) -> bool {
-        solana_vote_program::check_id(account.owner())
-            || stake::program::check_id(account.owner())
-                && account.data().len() >= StakeState::size_of()
-    }
-
     pub fn check_and_store(&self, pubkey: &Pubkey, account: &AccountSharedData) {
         if solana_vote_program::check_id(account.owner()) {
             let new_vote_account = if account.lamports() != 0
@@ -73,23 +75,9 @@ impl StakesCache {
                 .unwrap()
                 .update_vote_account(pubkey, new_vote_account);
         } else if solana_stake_program::check_id(account.owner()) {
-            let new_delegation = stake_state::delegation_from(account).map(|delegation| {
+            let stake_account = StakeAccount::try_from(account.clone()).ok();
-                let stakes = self.stakes();
+            let mut stakes = self.0.write().unwrap();
-                let stake = if account.lamports() != 0 {
+            stakes.update_stake_delegation(pubkey, stake_account);
-                    delegation.stake(stakes.epoch, Some(&stakes.stake_history))
-                } else {
-                    // when account is removed (lamports == 0), this special `else` clause ensures
-                    // resetting cached stake value below, even if the account happens to be
-                    // still staked for some (odd) reason
-                    0
-                };
-                (stake, delegation)
-            });
-
-            self.0
-                .write()
-                .unwrap()
-                .update_stake_delegation(pubkey, new_delegation);
         }
     }
 
@@ -143,13 +131,20 @@ impl StakesCache {
     }
 }
 
+/// The generic type T is either Delegation or StakeAccount.
+/// Stake<Delegation> is equivalent to the old code and is used for backward
+/// compatibility in BankFieldsToDeserialize.
+/// But banks cache Stakes<StakeAccount> which includes the entire stake
+/// account and StakeState deserialized from the account. Doing so, will remove
+/// the need to load the stake account from accounts-db when working with
+/// stake-delegations.
 #[derive(Default, Clone, PartialEq, Debug, Deserialize, Serialize, AbiExample)]
-pub struct Stakes {
+pub struct Stakes<T: Clone> {
     /// vote accounts
     vote_accounts: VoteAccounts,
 
     /// stake_delegations
-    stake_delegations: ImHashMap<Pubkey, Delegation>,
+    stake_delegations: ImHashMap<Pubkey, T>,
 
     /// unused
     unused: u64,
@@ -161,7 +156,48 @@ pub struct Stakes {
     stake_history: StakeHistory,
 }
 
-impl Stakes {
+impl<T: Clone> Stakes<T> {
+    pub fn vote_accounts(&self) -> &VoteAccounts {
+        &self.vote_accounts
+    }
+
+    pub(crate) fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
+        self.vote_accounts.staked_nodes()
+    }
+}
+
+impl Stakes<StakeAccount> {
+    /// Creates a Stake<StakeAccount> from Stake<Delegation> by loading the
+    /// full account state for respective stake pubkeys. get_account function
+    /// should return the account at the respective slot where stakes where
+    /// cached.
+    pub(crate) fn new<F>(stakes: &Stakes<Delegation>, get_account: F) -> Result<Self, Error>
+    where
+        F: Fn(&Pubkey) -> Option<AccountSharedData>,
+    {
+        let stake_delegations = stakes.stake_delegations.iter().map(|(pubkey, delegation)| {
+            let stake_account = match get_account(pubkey) {
+                None => return Err(Error::StakeAccountNotFound(*pubkey)),
+                Some(account) => account,
+            };
+            let stake_account = StakeAccount::try_from(stake_account)?;
+            // Sanity check that the delegation is consistent with what is
+            // stored in the account.
+            if stake_account.delegation() == Some(*delegation) {
+                Ok((*pubkey, stake_account))
+            } else {
+                Err(Error::InvalidDelegation(*pubkey))
+            }
+        });
+        Ok(Self {
+            vote_accounts: stakes.vote_accounts.clone(),
+            stake_delegations: stake_delegations.collect::<Result<_, _>>()?,
+            unused: stakes.unused,
+            epoch: stakes.epoch,
+            stake_history: stakes.stake_history.clone(),
+        })
+    }
+
     pub fn history(&self) -> &StakeHistory {
         &self.stake_history
     }
@@ -183,7 +219,8 @@ impl Stakes {
         let stake_history_entry = thread_pool.install(|| {
             stake_delegations
                 .par_iter()
-                .fold(StakeActivationStatus::default, |acc, delegation| {
+                .fold(StakeActivationStatus::default, |acc, stake_account| {
+                    let delegation = stake_account.delegation().unwrap();
                     acc + delegation
                         .stake_activating_and_deactivating(self.epoch, Some(&self.stake_history))
                 })
@@ -196,7 +233,8 @@ impl Stakes {
         let delegated_stakes = thread_pool.install(|| {
             stake_delegations
                 .par_iter()
-                .fold(HashMap::default, |mut delegated_stakes, delegation| {
+                .fold(HashMap::default, |mut delegated_stakes, stake_account| {
+                    let delegation = stake_account.delegation().unwrap();
                     let entry = delegated_stakes.entry(delegation.voter_pubkey).or_default();
                     *entry += delegation.stake(self.epoch, Some(&self.stake_history));
                     delegated_stakes
@@ -223,26 +261,20 @@ impl Stakes {
         epoch: Epoch,
         stake_history: &StakeHistory,
     ) -> u64 {
-        let matches_voter_pubkey = |(_, stake_delegation): &(&_, &Delegation)| {
-            &stake_delegation.voter_pubkey == voter_pubkey
-        };
-        let get_stake = |(_, stake_delegation): (_, &Delegation)| {
-            stake_delegation.stake(epoch, Some(stake_history))
-        };
-
         self.stake_delegations
-            .iter()
+            .values()
-            .filter(matches_voter_pubkey)
+            .map(|stake_account| stake_account.delegation().unwrap())
-            .map(get_stake)
+            .filter(|delegation| &delegation.voter_pubkey == voter_pubkey)
+            .map(|delegation| delegation.stake(epoch, Some(stake_history)))
             .sum()
     }
 
     /// Sum the lamports of the vote accounts and the delegated stake
     pub fn vote_balance_and_staked(&self) -> u64 {
-        let get_stake = |(_, stake_delegation): (_, &Delegation)| stake_delegation.stake;
+        let get_stake = |stake_account: &StakeAccount| stake_account.delegation().unwrap().stake;
         let get_lamports = |(_, (_, vote_account)): (_, &(_, VoteAccount))| vote_account.lamports();
 
-        self.stake_delegations.iter().map(get_stake).sum::<u64>()
+        self.stake_delegations.values().map(get_stake).sum::<u64>()
             + self.vote_accounts.iter().map(get_lamports).sum::<u64>()
     }
 
@@ -251,7 +283,8 @@ impl Stakes {
     }
 
     pub fn remove_stake_delegation(&mut self, stake_pubkey: &Pubkey) {
-        if let Some(removed_delegation) = self.stake_delegations.remove(stake_pubkey) {
+        if let Some(stake_account) = self.stake_delegations.remove(stake_pubkey) {
+            let removed_delegation = stake_account.delegation().unwrap();
             let removed_stake = removed_delegation.stake(self.epoch, Some(&self.stake_history));
             self.vote_accounts
                 .sub_stake(&removed_delegation.voter_pubkey, removed_stake);
@@ -282,18 +315,34 @@ impl Stakes {
     pub fn update_stake_delegation(
         &mut self,
         stake_pubkey: &Pubkey,
-        new_delegation: Option<(u64, Delegation)>,
+        new_stake_account: Option<StakeAccount>,
     ) {
         //  old_stake is stake lamports and voter_pubkey from the pre-store() version
-        let old_stake = self.stake_delegations.get(stake_pubkey).map(|delegation| {
+        let old_stake = self
+            .stake_delegations
+            .get(stake_pubkey)
+            .map(|stake_account| {
+                let delegation = stake_account.delegation().unwrap();
                 (
                     delegation.voter_pubkey,
                     delegation.stake(self.epoch, Some(&self.stake_history)),
                 )
             });
-
+        let new_delegation = new_stake_account
-        let new_stake = new_delegation.map(|(stake, delegation)| (delegation.voter_pubkey, stake));
+            .as_ref()
-
+            .and_then(StakeAccount::delegation);
+        let new_stake = new_stake_account.as_ref().and_then(|new_stake_account| {
+            let new_delegation = new_delegation?;
+            // When account is removed (lamports == 0), this check ensures
+            // resetting cached stake value below, even if the account happens
+            // to be still staked for some (odd) reason.
+            let stake = if new_stake_account.lamports() == 0 {
+                0
+            } else {
+                new_delegation.stake(self.epoch, Some(&self.stake_history))
+            };
+            Some((new_delegation.voter_pubkey, stake))
+        });
         // check if adjustments need to be made...
         if new_stake != old_stake {
             if let Some((voter_pubkey, stake)) = old_stake {
@@ -303,9 +352,10 @@ impl Stakes {
                 self.vote_accounts.add_stake(&voter_pubkey, stake);
             }
         }
-
+        // TODO: should this remove stake account if lamports == 0?
-        if let Some((_stake, delegation)) = new_delegation {
+        if new_delegation.is_some() {
-            self.stake_delegations.insert(*stake_pubkey, delegation);
+            self.stake_delegations
+                .insert(*stake_pubkey, new_stake_account.unwrap());
         } else {
             // when stake is no longer delegated, remove it from Stakes so that
             // given `pubkey` can be used for any owner in the future, while not
@@ -314,25 +364,35 @@ impl Stakes {
         }
     }
 
-    pub fn vote_accounts(&self) -> &VoteAccounts {
+    pub(crate) fn stake_delegations(&self) -> &ImHashMap<Pubkey, StakeAccount> {
-        &self.vote_accounts
-    }
-
-    pub(crate) fn stake_delegations(&self) -> &ImHashMap<Pubkey, Delegation> {
         &self.stake_delegations
     }
 
-    pub fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
+    pub(crate) fn highest_staked_node(&self) -> Option<Pubkey> {
-        self.vote_accounts.staked_nodes()
+        let key = |(_pubkey, (stake, _vote_account)): &(_, &(u64, _))| *stake;
+        let (_pubkey, (_stake, vote_account)) = self.vote_accounts.iter().max_by_key(key)?;
+        Some(vote_account.vote_state().as_ref().ok()?.node_pubkey)
+    }
 }
 
-    pub fn highest_staked_node(&self) -> Option<Pubkey> {
+impl TryFrom<Stakes<StakeAccount>> for Stakes<Delegation> {
-        let (_pubkey, (_stake, vote_account)) = self
+    type Error = Error;
-            .vote_accounts
+    fn try_from(stakes: Stakes<StakeAccount>) -> Result<Self, Self::Error> {
-            .iter()
+        let stake_delegations = stakes
-            .max_by(|(_ak, av), (_bk, bv)| av.0.cmp(&bv.0))?;
+            .stake_delegations
-        let node_pubkey = vote_account.vote_state().as_ref().ok()?.node_pubkey;
+            .into_iter()
-        Some(node_pubkey)
+            .map(|(pubkey, stake_account)| match stake_account.delegation() {
+                None => Err(Error::InvalidDelegation(pubkey)),
+                Some(delegation) => Ok((pubkey, delegation)),
+            })
+            .collect::<Result<_, _>>()?;
+        Ok(Self {
+            vote_accounts: stakes.vote_accounts,
+            stake_delegations,
+            unused: stakes.unused,
+            epoch: stakes.epoch,
+            stake_history: stakes.stake_history,
+        })
     }
 }
 
@@ -341,7 +401,7 @@ pub mod tests {
     use {
         super::*,
         rayon::ThreadPoolBuilder,
-        solana_sdk::{account::WritableAccount, pubkey::Pubkey, rent::Rent},
+        solana_sdk::{account::WritableAccount, pubkey::Pubkey, rent::Rent, stake},
         solana_stake_program::stake_state,
         solana_vote_program::vote_state::{self, VoteState, VoteStateVersions},
     };
@@ -707,14 +767,14 @@ pub mod tests {
 
     #[test]
     fn test_vote_balance_and_staked_empty() {
-        let stakes = Stakes::default();
+        let stakes = Stakes::<StakeAccount>::default();
         assert_eq!(stakes.vote_balance_and_staked(), 0);
     }
 
     #[test]
     fn test_vote_balance_and_staked_normal() {
         let stakes_cache = StakesCache::default();
-        impl Stakes {
+        impl Stakes<StakeAccount> {
             pub fn vote_balance_and_warmed_staked(&self) -> u64 {
                 self.vote_accounts
                     .iter()

runtime/src/vote_account.rs

@@ -196,6 +196,12 @@ impl From<AccountSharedData> for VoteAccount {
     }
 }
 
+impl From<VoteAccount> for AccountSharedData {
+    fn from(account: VoteAccount) -> Self {
+        Self::from(account.0.account.clone())
+    }
+}
+
 impl From<Account> for VoteAccount {
     fn from(account: Account) -> Self {
         Self(Arc::new(VoteAccountInner::from(account)))