fix: remove expensive versioned epoch stakes clone (#2453)

* fix: remove expensive versioned epoch stakes clone

* Add custom partialeq impl for dcou

Author: jstarry

runtime/src/bank/serde_snapshot.rs

@@ -6,7 +6,8 @@ mod tests {
                 epoch_accounts_hash_utils, test_utils as bank_test_utils, Bank, EpochRewardStatus,
             },
             epoch_stakes::{
-                EpochAuthorizedVoters, EpochStakes, NodeIdToVoteAccounts, VersionedEpochStakes,
+                EpochAuthorizedVoters, EpochStakes, NodeIdToVoteAccounts, StakesSerdeWrapper,
+                VersionedEpochStakes,
             },
             genesis_utils::activate_all_features,
             runtime_config::RuntimeConfig,
@@ -306,7 +307,7 @@ mod tests {
         bank.epoch_stakes.insert(
             42,
             EpochStakes::from(VersionedEpochStakes::Current {
-                stakes: Stakes::<Stake>::default(),
+                stakes: StakesSerdeWrapper::Stake(Stakes::<Stake>::default()),
                 total_stake: 42,
                 node_id_to_vote_accounts: Arc::<NodeIdToVoteAccounts>::default(),
                 epoch_authorized_voters: Arc::<EpochAuthorizedVoters>::default(),
@@ -535,7 +536,7 @@ mod tests {
         #[cfg_attr(
             feature = "frozen-abi",
             derive(AbiExample),
-            frozen_abi(digest = "CeNFPePrUfgJT2GNr7zYfMQVuJwGyU46bz1Skq7hAPht")
+            frozen_abi(digest = "7a6C1oFtgZiMtZig7FbX9289xn55QadQ962rX61Gheef")
         )]
         #[derive(Serialize)]
         pub struct BankAbiTestWrapper {

runtime/src/epoch_stakes.rs

@@ -1,7 +1,11 @@
 use {
-    crate::stakes::{Stakes, StakesEnum},
-    serde::{Deserialize, Serialize},
+    crate::{
+        stake_account::StakeAccount,
+        stakes::{Stakes, StakesEnum},
+    },
+    serde::{Deserialize, Deserializer, Serialize, Serializer},
     solana_sdk::{clock::Epoch, pubkey::Pubkey, stake::state::Stake},
+    solana_stake_program::stake_state::Delegation,
     solana_vote::vote_account::VoteAccountsHashMap,
     std::{collections::HashMap, sync::Arc},
 };
@@ -131,16 +135,78 @@ impl EpochStakes {
 }
 
 #[cfg_attr(feature = "frozen-abi", derive(AbiExample, AbiEnumVisitor))]
-#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum VersionedEpochStakes {
     Current {
-        stakes: Stakes<Stake>,
+        stakes: StakesSerdeWrapper,
         total_stake: u64,
         node_id_to_vote_accounts: Arc<NodeIdToVoteAccounts>,
         epoch_authorized_voters: Arc<EpochAuthorizedVoters>,
     },
 }
 
+/// Wrapper struct with custom serialization to support serializing
+/// `Stakes<StakeAccount>` as `Stakes<Stake>` without doing a full deep clone of
+/// the stake data. Serialization works by building a `Stakes<&Stake>` map which
+/// borrows `&Stake` from `StakeAccount` entries in `Stakes<StakeAccount>`. Note
+/// that `Stakes<&Stake>` still copies `Pubkey` keys so the `Stakes<&Stake>`
+/// data structure still allocates a fair amount of memory but the memory only
+/// remains allocated during serialization.
+#[cfg_attr(feature = "frozen-abi", derive(AbiExample, AbiEnumVisitor))]
+#[derive(Debug, Clone)]
+pub enum StakesSerdeWrapper {
+    Stake(Stakes<Stake>),
+    Account(Stakes<StakeAccount<Delegation>>),
+}
+
+#[cfg(feature = "dev-context-only-utils")]
+impl PartialEq<Self> for StakesSerdeWrapper {
+    fn eq(&self, other: &Self) -> bool {
+        match (self, other) {
+            (Self::Stake(stakes), Self::Stake(other)) => stakes == other,
+            (Self::Account(stakes), Self::Account(other)) => stakes == other,
+            (Self::Stake(stakes), Self::Account(other)) => {
+                stakes == &Stakes::<Stake>::from(other.clone())
+            }
+            (Self::Account(stakes), Self::Stake(other)) => {
+                other == &Stakes::<Stake>::from(stakes.clone())
+            }
+        }
+    }
+}
+
+impl From<StakesSerdeWrapper> for StakesEnum {
+    fn from(stakes: StakesSerdeWrapper) -> Self {
+        match stakes {
+            StakesSerdeWrapper::Stake(stakes) => Self::Stakes(stakes),
+            StakesSerdeWrapper::Account(stakes) => Self::Accounts(stakes),
+        }
+    }
+}
+
+impl Serialize for StakesSerdeWrapper {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        match self {
+            Self::Stake(stakes) => stakes.serialize(serializer),
+            Self::Account(stakes) => Stakes::<&Stake>::from(stakes).serialize(serializer),
+        }
+    }
+}
+
+impl<'de> Deserialize<'de> for StakesSerdeWrapper {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        let stakes = Stakes::<Stake>::deserialize(deserializer)?;
+        Ok(Self::Stake(stakes))
+    }
+}
+
 impl From<VersionedEpochStakes> for EpochStakes {
     fn from(versioned: VersionedEpochStakes) -> Self {
         let VersionedEpochStakes::Current {
@@ -151,7 +217,7 @@ impl From<VersionedEpochStakes> for EpochStakes {
         } = versioned;
 
         Self {
-            stakes: Arc::new(StakesEnum::Stakes(stakes)),
+            stakes: Arc::new(stakes.into()),
             total_stake,
             node_id_to_vote_accounts,
             epoch_authorized_voters,
@@ -196,7 +262,7 @@ pub(crate) fn split_epoch_stakes(
                 versioned_epoch_stakes.insert(
                     epoch,
                     VersionedEpochStakes::Current {
-                        stakes: Stakes::<Stake>::from(stakes.clone()),
+                        stakes: StakesSerdeWrapper::Account(stakes.clone()),
                         total_stake,
                         node_id_to_vote_accounts,
                         epoch_authorized_voters,
@@ -207,7 +273,7 @@ pub(crate) fn split_epoch_stakes(
                 versioned_epoch_stakes.insert(
                     epoch,
                     VersionedEpochStakes::Current {
-                        stakes: stakes.clone(),
+                        stakes: StakesSerdeWrapper::Stake(stakes.clone()),
                         total_stake,
                         node_id_to_vote_accounts,
                         epoch_authorized_voters,
@@ -426,7 +492,7 @@ pub(crate) mod tests {
         assert_eq!(
             versioned.get(&epoch),
             Some(&VersionedEpochStakes::Current {
-                stakes: Stakes::<Stake>::from(test_stakes),
+                stakes: StakesSerdeWrapper::Account(test_stakes),
                 total_stake: epoch_stakes.total_stake,
                 node_id_to_vote_accounts: epoch_stakes.node_id_to_vote_accounts,
                 epoch_authorized_voters: epoch_stakes.epoch_authorized_voters,
@@ -455,7 +521,7 @@ pub(crate) mod tests {
         assert_eq!(
             versioned.get(&epoch),
             Some(&VersionedEpochStakes::Current {
-                stakes: test_stakes,
+                stakes: StakesSerdeWrapper::Stake(test_stakes),
                 total_stake: epoch_stakes.total_stake,
                 node_id_to_vote_accounts: epoch_stakes.node_id_to_vote_accounts,
                 epoch_authorized_voters: epoch_stakes.epoch_authorized_voters,
@@ -506,8 +572,24 @@ pub(crate) mod tests {
         assert!(old.contains_key(&epoch1));
 
         assert_eq!(versioned.len(), 2);
-        assert!(versioned.contains_key(&epoch2));
-        assert!(versioned.contains_key(&epoch3));
+        assert_eq!(
+            versioned.get(&epoch2),
+            Some(&VersionedEpochStakes::Current {
+                stakes: StakesSerdeWrapper::Account(Stakes::default()),
+                total_stake: 200,
+                node_id_to_vote_accounts: Arc::default(),
+                epoch_authorized_voters: Arc::default(),
+            })
+        );
+        assert_eq!(
+            versioned.get(&epoch3),
+            Some(&VersionedEpochStakes::Current {
+                stakes: StakesSerdeWrapper::Stake(Stakes::default()),
+                total_stake: 300,
+                node_id_to_vote_accounts: Arc::default(),
+                epoch_authorized_voters: Arc::default(),
+            })
+        );
     }
 
     #[test]

runtime/src/serde_snapshot.rs

@@ -389,7 +389,8 @@ where
 /// added to this struct a minor release before they are added to the serialize
 /// struct.
 #[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
-#[derive(Clone, Debug, Deserialize, PartialEq)]
+#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
+#[derive(Clone, Debug, Deserialize)]
 struct ExtraFieldsToDeserialize {
     #[serde(deserialize_with = "default_on_eof")]
     lamports_per_signature: u64,
@@ -408,8 +409,8 @@ struct ExtraFieldsToDeserialize {
 /// be added to the deserialize struct a minor release before they are added to
 /// this one.
 #[cfg_attr(feature = "frozen-abi", derive(AbiExample))]
-#[cfg_attr(feature = "dev-context-only-utils", derive(Default))]
-#[derive(Debug, Serialize, PartialEq)]
+#[cfg_attr(feature = "dev-context-only-utils", derive(Default, PartialEq))]
+#[derive(Debug, Serialize)]
 pub struct ExtraFieldsToSerialize<'a> {
     pub lamports_per_signature: u64,
     pub incremental_snapshot_persistence: Option<&'a BankIncrementalSnapshotPersistence>,

runtime/src/stake_account.rs

@@ -55,10 +55,10 @@ impl StakeAccount<Delegation> {
     }
 
     #[inline]
-    pub(crate) fn stake(&self) -> Stake {
+    pub(crate) fn stake(&self) -> &Stake {
         // Safe to unwrap here because StakeAccount<Delegation> will always
         // only wrap a stake-state.
-        self.stake_state.stake().unwrap()
+        self.stake_state.stake_ref().unwrap()
     }
 }
 

runtime/src/stakes.rs

@@ -544,12 +544,15 @@ impl From<Stakes<StakeAccount>> for Stakes<Delegation> {
     }
 }
 
+/// This conversion is very memory intensive so should only be used in
+/// development contexts.
+#[cfg(feature = "dev-context-only-utils")]
 impl From<Stakes<StakeAccount>> for Stakes<Stake> {
     fn from(stakes: Stakes<StakeAccount>) -> Self {
         let stake_delegations = stakes
             .stake_delegations
             .into_iter()
-            .map(|(pubkey, stake_account)| (pubkey, stake_account.stake()))
+            .map(|(pubkey, stake_account)| (pubkey, *stake_account.stake()))
             .collect();
         Self {
             vote_accounts: stakes.vote_accounts,
@@ -561,6 +564,23 @@ impl From<Stakes<StakeAccount>> for Stakes<Stake> {
     }
 }
 
+impl<'a> From<&'a Stakes<StakeAccount>> for Stakes<&'a Stake> {
+    fn from(stakes: &'a Stakes<StakeAccount>) -> Self {
+        let stake_delegations = stakes
+            .stake_delegations
+            .iter()
+            .map(|(pubkey, stake_account)| (*pubkey, stake_account.stake()))
+            .collect();
+        Self {
+            vote_accounts: stakes.vote_accounts.clone(),
+            stake_delegations,
+            unused: stakes.unused,
+            epoch: stakes.epoch,
+            stake_history: stakes.stake_history.clone(),
+        }
+    }
+}
+
 impl From<Stakes<Stake>> for Stakes<Delegation> {
     fn from(stakes: Stakes<Stake>) -> Self {
         let stake_delegations = stakes

sdk/program/src/stake/state.rs

@@ -213,6 +213,13 @@ impl StakeStateV2 {
         }
     }
 
+    pub fn stake_ref(&self) -> Option<&Stake> {
+        match self {
+            StakeStateV2::Stake(_meta, stake, _stake_flags) => Some(stake),
+            _ => None,
+        }
+    }
+
     pub fn delegation(&self) -> Option<Delegation> {
         match self {
             StakeStateV2::Stake(_meta, stake, _stake_flags) => Some(stake.delegation),