feat: deserialize new epoch stakes bank snapshot field (solana-labs#1397

)

runtime/src/bank/serde_snapshot.rs

@@ -5,6 +5,9 @@ mod tests {
             bank::{
                 epoch_accounts_hash_utils, test_utils as bank_test_utils, Bank, EpochRewardStatus,
             },
+            epoch_stakes::{
+                EpochAuthorizedVoters, EpochStakes, NodeIdToVoteAccounts, VersionedEpochStakes,
+            },
             genesis_utils::activate_all_features,
             runtime_config::RuntimeConfig,
             serde_snapshot::{
@@ -16,6 +19,7 @@ mod tests {
                 self, create_tmp_accounts_dir_for_tests, get_storages_to_serialize, ArchiveFormat,
                 StorageAndNextAccountsFileId, BANK_SNAPSHOT_PRE_FILENAME_EXTENSION,
             },
+            stakes::Stakes,
             status_cache::StatusCache,
         },
         solana_accounts_db::{
@@ -35,8 +39,10 @@ mod tests {
             hash::Hash,
             pubkey::Pubkey,
             signature::{Keypair, Signer},
+            stake::state::Stake,
         },
         std::{
+            collections::HashMap,
             io::{Cursor, Read, Write},
             ops::RangeFull,
             path::Path,
@@ -369,6 +375,18 @@ mod tests {
         )
         .unwrap();
 
+        let mut new_epoch_stakes: HashMap<u64, VersionedEpochStakes> = HashMap::new();
+        new_epoch_stakes.insert(
+            42,
+            VersionedEpochStakes::Current {
+                stakes: Stakes::<Stake>::default(),
+                total_stake: 42,
+                node_id_to_vote_accounts: Arc::<NodeIdToVoteAccounts>::default(),
+                epoch_authorized_voters: Arc::<EpochAuthorizedVoters>::default(),
+            },
+        );
+        bincode::serialize_into(&mut writer, &new_epoch_stakes).unwrap();
+
         // Deserialize
         let rdr = Cursor::new(&buf[..]);
         let mut reader = std::io::BufReader::new(&buf[rdr.position() as usize..]);
@@ -402,6 +420,13 @@ mod tests {
         )
         .unwrap();
 
+        assert_eq!(
+            dbank.epoch_stakes(42),
+            Some(&EpochStakes::from(
+                new_epoch_stakes.get(&42).unwrap().clone()
+            ))
+        );
+
         assert_eq!(
             bank.fee_rate_governor.lamports_per_signature,
             dbank.fee_rate_governor.lamports_per_signature

runtime/src/epoch_stakes.rs

@@ -1,7 +1,7 @@
 use {
-    crate::stakes::StakesEnum,
+    crate::stakes::{Stakes, StakesEnum},
     serde::{Deserialize, Serialize},
-    solana_sdk::{clock::Epoch, pubkey::Pubkey},
+    solana_sdk::{clock::Epoch, pubkey::Pubkey, stake::state::Stake},
     solana_vote::vote_account::VoteAccountsHashMap,
     std::{collections::HashMap, sync::Arc},
 };
@@ -124,6 +124,34 @@ impl EpochStakes {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub(crate) enum VersionedEpochStakes {
+    Current {
+        stakes: Stakes<Stake>,
+        total_stake: u64,
+        node_id_to_vote_accounts: Arc<NodeIdToVoteAccounts>,
+        epoch_authorized_voters: Arc<EpochAuthorizedVoters>,
+    },
+}
+
+impl From<VersionedEpochStakes> for EpochStakes {
+    fn from(versioned: VersionedEpochStakes) -> Self {
+        let VersionedEpochStakes::Current {
+            stakes,
+            total_stake,
+            node_id_to_vote_accounts,
+            epoch_authorized_voters,
+        } = versioned;
+
+        Self {
+            stakes: Arc::new(StakesEnum::Stakes(stakes)),
+            total_stake,
+            node_id_to_vote_accounts,
+            epoch_authorized_voters,
+        }
+    }
+}
+
 #[cfg(test)]
 pub(crate) mod tests {
     use {

runtime/src/serde_snapshot.rs

@@ -6,7 +6,7 @@ use {
             builtins::BuiltinPrototype, Bank, BankFieldsToDeserialize, BankFieldsToSerialize,
             BankRc,
         },
-        epoch_stakes::EpochStakes,
+        epoch_stakes::{EpochStakes, VersionedEpochStakes},
         runtime_config::RuntimeConfig,
         serde_snapshot::storage::SerializableAccountStorageEntry,
         snapshot_utils::{
@@ -412,6 +412,16 @@ where
     let epoch_accounts_hash = ignore_eof_error(deserialize_from(&mut stream))?;
     bank_fields.epoch_accounts_hash = epoch_accounts_hash;
 
+    // If we deserialize the new epoch stakes, add all of the entries into the
+    // other deserialized map which could still have old epoch stakes entries
+    let new_epoch_stakes: HashMap<u64, VersionedEpochStakes> =
+        ignore_eof_error(deserialize_from(&mut stream))?;
+    bank_fields.epoch_stakes.extend(
+        new_epoch_stakes
+            .into_iter()
+            .map(|(epoch, versioned_epoch_stakes)| (epoch, versioned_epoch_stakes.into())),
+    );
+
     Ok((bank_fields, accounts_db_fields))
 }
 

runtime/src/stakes.rs

@@ -16,6 +16,7 @@ use {
         stake::state::{Delegation, StakeActivationStatus},
         vote::state::VoteStateVersions,
     },
+    solana_stake_program::stake_state::Stake,
     solana_vote::vote_account::{VoteAccount, VoteAccounts},
     std::{
         collections::HashMap,
@@ -200,17 +201,19 @@ pub struct Stakes<T: Clone> {
 }
 
 // For backward compatibility, we can only serialize and deserialize
-// Stakes<Delegation>. However Bank caches Stakes<StakeAccount>. This type
-// mismatch incurs a conversion cost at epoch boundary when updating
-// EpochStakes.
-// Below type allows EpochStakes to include either a Stakes<StakeAccount> or
-// Stakes<Delegation> and so bypass the conversion cost between the two at the
-// epoch boundary.
+// Stakes<Delegation> in the old `epoch_stakes` bank snapshot field. However,
+// Stakes<StakeAccount> entries are added to the bank's epoch stakes hashmap
+// when crossing epoch boundaries and Stakes<Stake> entries are added when
+// starting up from bank snapshots that have the new epoch stakes field. By
+// using this enum, the cost of converting all entries to Stakes<Delegation> is
+// put off until serializing new snapshots. This helps avoid bogging down epoch
+// boundaries and startup with the conversion overhead.
 #[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum StakesEnum {
     Accounts(Stakes<StakeAccount>),
     Delegations(Stakes<Delegation>),
+    Stakes(Stakes<Stake>),
 }
 
 impl<T: Clone> Stakes<T> {
@@ -483,13 +486,15 @@ impl StakesEnum {
         match self {
             StakesEnum::Accounts(stakes) => stakes.vote_accounts(),
             StakesEnum::Delegations(stakes) => stakes.vote_accounts(),
+            StakesEnum::Stakes(stakes) => stakes.vote_accounts(),
         }
     }
 
     pub(crate) fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
         match self {
             StakesEnum::Accounts(stakes) => stakes.staked_nodes(),
             StakesEnum::Delegations(stakes) => stakes.staked_nodes(),
+            StakesEnum::Stakes(stakes) => stakes.staked_nodes(),
         }
     }
 }
@@ -511,6 +516,33 @@ impl From<Stakes<StakeAccount>> for Stakes<Delegation> {
     }
 }
 
+impl From<Stakes<Stake>> for Stakes<Delegation> {
+    fn from(stakes: Stakes<Stake>) -> Self {
+        let stake_delegations = stakes
+            .stake_delegations
+            .into_iter()
+            .map(|(pubkey, stake)| (pubkey, stake.delegation))
+            .collect();
+        Self {
+            vote_accounts: stakes.vote_accounts,
+            stake_delegations,
+            unused: stakes.unused,
+            epoch: stakes.epoch,
+            stake_history: stakes.stake_history,
+        }
+    }
+}
+
+impl From<StakesEnum> for Stakes<Delegation> {
+    fn from(stakes: StakesEnum) -> Self {
+        match stakes {
+            StakesEnum::Accounts(stakes) => stakes.into(),
+            StakesEnum::Delegations(stakes) => stakes,
+            StakesEnum::Stakes(stakes) => stakes.into(),
+        }
+    }
+}
+
 impl From<Stakes<StakeAccount>> for StakesEnum {
     fn from(stakes: Stakes<StakeAccount>) -> Self {
         Self::Accounts(stakes)
@@ -533,15 +565,13 @@ impl PartialEq<StakesEnum> for StakesEnum {
     fn eq(&self, other: &StakesEnum) -> bool {
         match (self, other) {
             (Self::Accounts(stakes), Self::Accounts(other)) => stakes == other,
-            (Self::Accounts(stakes), Self::Delegations(other)) => {
+            (Self::Delegations(stakes), Self::Delegations(other)) => stakes == other,
+            (Self::Stakes(stakes), Self::Stakes(other)) => stakes == other,
+            (stakes, other) => {
                 let stakes = Stakes::<Delegation>::from(stakes.clone());
-                &stakes == other
-            }
-            (Self::Delegations(stakes), Self::Accounts(other)) => {
                 let other = Stakes::<Delegation>::from(other.clone());
-                stakes == &other
+                stakes == other
             }
-            (Self::Delegations(stakes), Self::Delegations(other)) => stakes == other,
         }
     }
 }
@@ -559,11 +589,11 @@ pub(crate) mod serde_stakes_enum_compat {
         S: Serializer,
     {
         match stakes {
-            StakesEnum::Accounts(stakes) => {
+            StakesEnum::Delegations(stakes) => stakes.serialize(serializer),
+            stakes => {
                 let stakes = Stakes::<Delegation>::from(stakes.clone());
                 stakes.serialize(serializer)
             }
-            StakesEnum::Delegations(stakes) => stakes.serialize(serializer),
         }
     }
 
@@ -1101,7 +1131,7 @@ pub(crate) mod tests {
         assert!(stakes.vote_accounts.as_ref().len() >= 5);
         assert!(stakes.stake_delegations.len() >= 50);
         let other = match &*other.stakes {
-            StakesEnum::Accounts(_) => panic!("wrong type!"),
+            StakesEnum::Accounts(_) | StakesEnum::Stakes(_) => panic!("wrong type!"),
             StakesEnum::Delegations(delegations) => delegations,
         };
         assert_eq!(other, &stakes)