adds StakesEnum type representing Stakes<StakeAccount|Delegation>

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.

This commit adds StakesEnum which allows EpochStakes to include either a
Stakes<StakeAccount> or Stakes<Delegation> and so bypass the conversion
cost between the two at the epoch boundary.

runtime/src/bank.rs

@@ -61,7 +61,7 @@ use {
             calculate_stake_weighted_timestamp, MaxAllowableDrift, MAX_ALLOWABLE_DRIFT_PERCENTAGE,
             MAX_ALLOWABLE_DRIFT_PERCENTAGE_FAST, MAX_ALLOWABLE_DRIFT_PERCENTAGE_SLOW,
         },
-        stakes::{InvalidCacheEntryReason, Stakes, StakesCache},
+        stakes::{InvalidCacheEntryReason, Stakes, StakesCache, StakesEnum},
         status_cache::{SlotDelta, StatusCache},
         system_instruction_processor::{get_system_account_kind, SystemAccountKind},
         transaction_batch::TransactionBatch,
@@ -1499,8 +1499,7 @@ impl Bank {
         //  slot = 0 and genesis configuration
         {
             let stakes = bank.stakes_cache.stakes().clone();
-            let stakes = Stakes::<Delegation>::try_from(stakes).unwrap();
-            let stakes = Arc::new(stakes);
+            let stakes = Arc::new(StakesEnum::from(stakes));
             for epoch in 0..=bank.get_leader_schedule_epoch(bank.slot) {
                 bank.epoch_stakes
                     .insert(epoch, EpochStakes::new(stakes.clone(), epoch));
@@ -2404,8 +2403,8 @@ impl Bank {
                 epoch >= leader_schedule_epoch.saturating_sub(MAX_LEADER_SCHEDULE_STAKES)
             });
             let stakes = self.stakes_cache.stakes().clone();
-            let stakes = Stakes::<Delegation>::try_from(stakes).unwrap();
-            let new_epoch_stakes = EpochStakes::new(Arc::new(stakes), leader_schedule_epoch);
+            let stakes = Arc::new(StakesEnum::from(stakes));
+            let new_epoch_stakes = EpochStakes::new(stakes, leader_schedule_epoch);
             {
                 let vote_stakes: HashMap<_, _> = self
                     .stakes_cache

runtime/src/epoch_stakes.rs

@@ -1,7 +1,7 @@
 use {
-    crate::{stakes::Stakes, vote_account::VoteAccountsHashMap},
+    crate::{stakes::StakesEnum, vote_account::VoteAccountsHashMap},
     serde::{Deserialize, Serialize},
-    solana_sdk::{clock::Epoch, pubkey::Pubkey, stake::state::Delegation},
+    solana_sdk::{clock::Epoch, pubkey::Pubkey},
     std::{collections::HashMap, sync::Arc},
 };
 
@@ -16,14 +16,15 @@ pub struct NodeVoteAccounts {
 
 #[derive(Clone, Debug, Serialize, Deserialize, AbiExample, PartialEq)]
 pub struct EpochStakes {
-    stakes: Arc<Stakes<Delegation>>,
+    #[serde(with = "crate::stakes::serde_stakes_enum_compat")]
+    stakes: Arc<StakesEnum>,
     total_stake: u64,
     node_id_to_vote_accounts: Arc<NodeIdToVoteAccounts>,
     epoch_authorized_voters: Arc<EpochAuthorizedVoters>,
 }
 
 impl EpochStakes {
-    pub(crate) fn new(stakes: Arc<Stakes<Delegation>>, leader_schedule_epoch: Epoch) -> Self {
+    pub(crate) fn new(stakes: Arc<StakesEnum>, 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);
@@ -35,7 +36,7 @@ impl EpochStakes {
         }
     }
 
-    pub fn stakes(&self) -> &Stakes<Delegation> {
+    pub fn stakes(&self) -> &StakesEnum {
         &self.stakes
     }
 

runtime/src/serde_snapshot/newer.rs

@@ -4,7 +4,10 @@ use {
         utils::{serialize_iter_as_map, serialize_iter_as_seq},
         *,
     },
-    crate::ancestors::AncestorsForSerialization,
+    crate::{
+        ancestors::AncestorsForSerialization,
+        stakes::{serde_stakes_enum_compat, StakesEnum},
+    },
     solana_measure::measure::Measure,
     solana_sdk::stake::state::Delegation,
     std::{cell::RefCell, collections::HashSet, sync::RwLock},
@@ -129,7 +132,8 @@ struct SerializableVersionedBank<'a> {
     rent_collector: RentCollector,
     epoch_schedule: EpochSchedule,
     inflation: Inflation,
-    stakes: Stakes<Delegation>,
+    #[serde(serialize_with = "serde_stakes_enum_compat::serialize")]
+    stakes: StakesEnum,
     unused_accounts: UnusedAccounts,
     epoch_stakes: &'a HashMap<Epoch, EpochStakes>,
     is_delta: bool,
@@ -166,10 +170,7 @@ impl<'a> From<crate::bank::BankFieldsToSerialize<'a>> for SerializableVersionedB
             rent_collector: rhs.rent_collector,
             epoch_schedule: rhs.epoch_schedule,
             inflation: rhs.inflation,
-            stakes: {
-                let stakes = rhs.stakes.stakes().clone();
-                Stakes::<Delegation>::try_from(stakes).unwrap()
-            },
+            stakes: StakesEnum::from(rhs.stakes.stakes().clone()),
             unused_accounts: UnusedAccounts::default(),
             epoch_stakes: rhs.epoch_stakes,
             is_delta: rhs.is_delta,
@@ -339,7 +340,7 @@ impl<'a> TypeContext<'a> for Context {
             rent_collector: rhs.rent_collector,
             epoch_schedule: rhs.epoch_schedule,
             inflation: rhs.inflation,
-            stakes: rhs.stakes,
+            stakes: StakesEnum::from(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 = "2RqbikYfG4rXV8gtqsKCNSn6g4g6nfkegXDRe1G6RGFQ")]
+    #[frozen_abi(digest = "HT9yewU4zJ6ZAgJ7aDSbHPtzZGZqASpq6rkq6ET42Kki")]
     #[derive(Serialize, AbiExample)]
     pub struct BankAbiTestWrapperNewer {
         #[serde(serialize_with = "wrapper_newer")]

runtime/src/stake_account.rs

@@ -13,10 +13,10 @@ use {
 
 /// An account and a stake state deserialized from the account.
 #[derive(Clone, Debug, Default, PartialEq)]
-pub(crate) struct StakeAccount(AccountSharedData, StakeState);
+pub struct StakeAccount(AccountSharedData, StakeState);
 
 #[derive(Debug, Error)]
-pub(crate) enum Error {
+pub enum Error {
     #[error(transparent)]
     InstructionError(#[from] InstructionError),
     #[error("Invalid stake account owner: {owner:?}")]

runtime/src/stakes.rs

@@ -27,7 +27,7 @@ use {
 };
 
 #[derive(Debug, Error)]
-pub(crate) enum Error {
+pub enum Error {
     #[error("Invalid delegation: {0}")]
     InvalidDelegation(Pubkey),
     #[error(transparent)]
@@ -156,6 +156,19 @@ pub struct Stakes<T: Clone> {
     stake_history: StakeHistory,
 }
 
+// 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.
+#[derive(Debug, AbiExample)]
+pub enum StakesEnum {
+    Accounts(Stakes<StakeAccount>),
+    Delegations(Stakes<Delegation>),
+}
+
 impl<T: Clone> Stakes<T> {
     pub fn vote_accounts(&self) -> &VoteAccounts {
         &self.vote_accounts
@@ -375,6 +388,22 @@ impl Stakes<StakeAccount> {
     }
 }
 
+impl StakesEnum {
+    pub fn vote_accounts(&self) -> &VoteAccounts {
+        match self {
+            StakesEnum::Accounts(stakes) => stakes.vote_accounts(),
+            StakesEnum::Delegations(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(),
+        }
+    }
+}
+
 impl TryFrom<Stakes<StakeAccount>> for Stakes<Delegation> {
     type Error = Error;
     fn try_from(stakes: Stakes<StakeAccount>) -> Result<Self, Self::Error> {
@@ -396,10 +425,85 @@ impl TryFrom<Stakes<StakeAccount>> for Stakes<Delegation> {
     }
 }
 
+impl From<Stakes<StakeAccount>> for StakesEnum {
+    fn from(stakes: Stakes<StakeAccount>) -> Self {
+        Self::Accounts(stakes)
+    }
+}
+
+impl From<Stakes<Delegation>> for StakesEnum {
+    fn from(stakes: Stakes<Delegation>) -> Self {
+        Self::Delegations(stakes)
+    }
+}
+
+// Two StakesEnums are equal as long as they represent the same delegations;
+// whether these delegations are stored as StakeAccounts or Delegations.
+// Therefore, if one side is Stakes<StakeAccount> and the other is a
+// Stakes<Delegation> we convert the former one to Stakes<Delegation> before
+// comparing for equality.
+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)) => {
+                match Stakes::<Delegation>::try_from(stakes.clone()) {
+                    Ok(stakes) => stakes == *other,
+                    Err(_) => false,
+                }
+            }
+            (Self::Delegations(stakes), Self::Accounts(other)) => {
+                match Stakes::<Delegation>::try_from(other.clone()) {
+                    Ok(other) => *stakes == other,
+                    Err(_) => false,
+                }
+            }
+            (Self::Delegations(stakes), Self::Delegations(other)) => stakes == other,
+        }
+    }
+}
+
+// In order to maintain backward compatibility, the StakesEnum in EpochStakes
+// and SerializableVersionedBank should be serialized as Stakes<Delegation>.
+pub(crate) mod serde_stakes_enum_compat {
+    use {
+        super::*,
+        serde::{Deserialize, Deserializer, Serialize, Serializer},
+    };
+
+    pub(crate) fn serialize<S>(stakes: &StakesEnum, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        match stakes {
+            StakesEnum::Accounts(stakes) => {
+                let stakes = match Stakes::<Delegation>::try_from(stakes.clone()) {
+                    Ok(stakes) => stakes,
+                    Err(err) => {
+                        let err = format!("{:?}", err);
+                        return Err(serde::ser::Error::custom(err));
+                    }
+                };
+                stakes.serialize(serializer)
+            }
+            StakesEnum::Delegations(stakes) => stakes.serialize(serializer),
+        }
+    }
+
+    pub(crate) fn deserialize<'de, D>(deserializer: D) -> Result<Arc<StakesEnum>, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let stakes = Stakes::<Delegation>::deserialize(deserializer)?;
+        Ok(Arc::new(StakesEnum::Delegations(stakes)))
+    }
+}
+
 #[cfg(test)]
 pub mod tests {
     use {
         super::*,
+        rand::Rng,
         rayon::ThreadPoolBuilder,
         solana_sdk::{account::WritableAccount, pubkey::Pubkey, rent::Rent, stake},
         solana_stake_program::stake_state,
@@ -808,4 +912,63 @@ pub mod tests {
             );
         }
     }
+
+    #[test]
+    fn test_serde_stakes_enum_compat() {
+        #[derive(Debug, PartialEq, Deserialize, Serialize)]
+        struct Dummy {
+            head: String,
+            #[serde(with = "serde_stakes_enum_compat")]
+            stakes: Arc<StakesEnum>,
+            tail: String,
+        }
+        let mut rng = rand::thread_rng();
+        let stakes_cache = StakesCache::new(Stakes {
+            unused: rng.gen(),
+            epoch: rng.gen(),
+            ..Stakes::default()
+        });
+        for _ in 0..rng.gen_range(5usize, 10) {
+            let vote_pubkey = solana_sdk::pubkey::new_rand();
+            let vote_account = vote_state::create_account(
+                &vote_pubkey,
+                &solana_sdk::pubkey::new_rand(), // node_pubkey
+                rng.gen_range(0, 101),           // commission
+                rng.gen_range(0, 1_000_000),     // lamports
+            );
+            stakes_cache.check_and_store(&vote_pubkey, &vote_account);
+            for _ in 0..rng.gen_range(10usize, 20) {
+                let stake_pubkey = solana_sdk::pubkey::new_rand();
+                let rent = Rent::with_slots_per_epoch(rng.gen());
+                let stake_account = stake_state::create_account(
+                    &stake_pubkey, // authorized
+                    &vote_pubkey,
+                    &vote_account,
+                    &rent,
+                    rng.gen_range(0, 1_000_000), // lamports
+                );
+                stakes_cache.check_and_store(&stake_pubkey, &stake_account);
+            }
+        }
+        let stakes: Stakes<StakeAccount> = stakes_cache.stakes().clone();
+        assert!(stakes.vote_accounts.as_ref().len() >= 5);
+        assert!(stakes.stake_delegations.len() >= 50);
+        let dummy = Dummy {
+            head: String::from("dummy-head"),
+            stakes: Arc::new(StakesEnum::from(stakes.clone())),
+            tail: String::from("dummy-tail"),
+        };
+        assert!(dummy.stakes.vote_accounts().as_ref().len() >= 5);
+        let data = bincode::serialize(&dummy).unwrap();
+        let other: Dummy = bincode::deserialize(&data).unwrap();
+        assert_eq!(other, dummy);
+        let stakes = Stakes::<Delegation>::try_from(stakes).unwrap();
+        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::Delegations(delegations) => delegations,
+        };
+        assert_eq!(other, &stakes)
+    }
 }