From f39eab171aa620a53b5fa6fd89df550aa8c1d121 Mon Sep 17 00:00:00 2001 From: Justin Starry Date: Tue, 30 Jul 2024 10:00:08 +0800 Subject: [PATCH] feat: serialize new versioned epoch stakes bank snapshot field (#2282) * feat: serialize new versioned epoch stakes bank snapshot field * fix abi test * feedback * feedback --- runtime/src/bank.rs | 8 +- runtime/src/bank/serde_snapshot.rs | 113 +++++++++----- runtime/src/epoch_stakes.rs | 235 ++++++++++++++++++++++++++++- runtime/src/serde_snapshot.rs | 25 +-- runtime/src/snapshot_package.rs | 3 +- runtime/src/snapshot_utils.rs | 5 +- runtime/src/stake_account.rs | 9 +- runtime/src/stakes.rs | 17 +++ 8 files changed, 356 insertions(+), 59 deletions(-) diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index 278ba6009118a1..ca62027030f39b 100644 --- a/runtime/src/bank.rs +++ b/runtime/src/bank.rs @@ -41,7 +41,7 @@ use { partitioned_epoch_rewards::{EpochRewardStatus, StakeRewards, VoteRewardsAccounts}, }, bank_forks::BankForks, - epoch_stakes::{EpochStakes, NodeVoteAccounts}, + epoch_stakes::{split_epoch_stakes, EpochStakes, NodeVoteAccounts, VersionedEpochStakes}, installed_scheduler_pool::{BankWithScheduler, InstalledSchedulerRwLock}, runtime_config::RuntimeConfig, serde_snapshot::BankIncrementalSnapshotPersistence, @@ -499,6 +499,7 @@ pub struct BankFieldsToSerialize { pub epoch_stakes: HashMap, pub is_delta: bool, pub accounts_data_len: u64, + pub versioned_epoch_stakes: HashMap, } // Can't derive PartialEq because RwLock doesn't implement PartialEq @@ -642,6 +643,7 @@ impl BankFieldsToSerialize { epoch_stakes: HashMap::default(), is_delta: bool::default(), accounts_data_len: u64::default(), + versioned_epoch_stakes: HashMap::default(), } } } @@ -1655,6 +1657,7 @@ impl Bank { /// Return subset of bank fields representing serializable state pub(crate) fn get_fields_to_serialize(&self) -> BankFieldsToSerialize { + let (epoch_stakes, versioned_epoch_stakes) = split_epoch_stakes(self.epoch_stakes.clone()); BankFieldsToSerialize { blockhash_queue: self.blockhash_queue.read().unwrap().clone(), ancestors: AncestorsForSerialization::from(&self.ancestors), @@ -1683,9 +1686,10 @@ impl Bank { epoch_schedule: self.epoch_schedule.clone(), inflation: *self.inflation.read().unwrap(), stakes: StakesEnum::from(self.stakes_cache.stakes().clone()), - epoch_stakes: self.epoch_stakes.clone(), + epoch_stakes, is_delta: self.is_delta.load(Relaxed), accounts_data_len: self.load_accounts_data_size(), + versioned_epoch_stakes, } } diff --git a/runtime/src/bank/serde_snapshot.rs b/runtime/src/bank/serde_snapshot.rs index 4be5b833566376..2f633c0910b2d8 100644 --- a/runtime/src/bank/serde_snapshot.rs +++ b/runtime/src/bank/serde_snapshot.rs @@ -19,7 +19,7 @@ mod tests { create_tmp_accounts_dir_for_tests, get_storages_to_serialize, ArchiveFormat, StorageAndNextAccountsFileId, }, - stakes::Stakes, + stakes::{Stakes, StakesEnum}, }, solana_accounts_db::{ account_storage::{AccountStorageMap, AccountStorageReference}, @@ -37,8 +37,8 @@ mod tests { pubkey::Pubkey, stake::state::Stake, }, std::{ - collections::HashMap, io::{BufReader, BufWriter, Cursor}, + mem, ops::RangeFull, path::Path, sync::{atomic::Ordering, Arc}, @@ -124,7 +124,7 @@ mod tests { } else { 0 } + 2; - let bank2 = Bank::new_from_parent(bank0, &Pubkey::default(), bank2_slot); + let mut bank2 = Bank::new_from_parent(bank0, &Pubkey::default(), bank2_slot); // Test new account let key2 = Pubkey::new_unique(); @@ -158,25 +158,59 @@ mod tests { epoch_accounts_hash }); + // Only if a bank was recently recreated from a snapshot will it have an epoch stakes entry + // of type "delegations" which cannot be serialized into the versioned epoch stakes map. Simulate + // this condition by replacing the epoch 0 stakes map of stake accounts with an epoch stakes map + // of delegations. + { + assert_eq!(bank2.epoch_stakes.len(), 2); + assert!(bank2 + .epoch_stakes + .values() + .all(|epoch_stakes| matches!(epoch_stakes.stakes(), &StakesEnum::Accounts(_)))); + + let StakesEnum::Accounts(stake_accounts) = + bank2.epoch_stakes.remove(&0).unwrap().stakes().clone() + else { + panic!("expected the epoch 0 stakes entry to have stake accounts"); + }; + + bank2.epoch_stakes.insert( + 0, + EpochStakes::new(Arc::new(StakesEnum::Delegations(stake_accounts.into())), 0), + ); + } + let mut buf = Vec::new(); let cursor = Cursor::new(&mut buf); let mut writer = BufWriter::new(cursor); - serde_snapshot::serialize_bank_snapshot_into( - &mut writer, - bank2.get_fields_to_serialize(), - accounts_db.get_bank_hash_stats(bank2_slot).unwrap(), - accounts_db.get_accounts_delta_hash(bank2_slot).unwrap(), - expected_accounts_hash, - &get_storages_to_serialize(&bank2.get_snapshot_storages(None)), - ExtraFieldsToSerialize { - lamports_per_signature: bank2.fee_rate_governor.lamports_per_signature, - incremental_snapshot_persistence: expected_incremental_snapshot_persistence - .as_ref(), - epoch_accounts_hash: expected_epoch_accounts_hash, - }, - accounts_db.write_version.load(Ordering::Acquire), - ) - .unwrap(); + { + let mut bank_fields = bank2.get_fields_to_serialize(); + // Ensure that epoch_stakes and versioned_epoch_stakes are each + // serialized with at least one entry to verify that epoch stakes + // entries are combined correctly during deserialization + assert!(!bank_fields.epoch_stakes.is_empty()); + assert!(!bank_fields.versioned_epoch_stakes.is_empty()); + + let versioned_epoch_stakes = mem::take(&mut bank_fields.versioned_epoch_stakes); + serde_snapshot::serialize_bank_snapshot_into( + &mut writer, + bank_fields, + accounts_db.get_bank_hash_stats(bank2_slot).unwrap(), + accounts_db.get_accounts_delta_hash(bank2_slot).unwrap(), + expected_accounts_hash, + &get_storages_to_serialize(&bank2.get_snapshot_storages(None)), + ExtraFieldsToSerialize { + lamports_per_signature: bank2.fee_rate_governor.lamports_per_signature, + incremental_snapshot_persistence: expected_incremental_snapshot_persistence + .as_ref(), + epoch_accounts_hash: expected_epoch_accounts_hash, + versioned_epoch_stakes, + }, + accounts_db.write_version.load(Ordering::Acquire), + ) + .unwrap(); + } drop(writer); // Now deserialize the serialized bank and ensure it matches the original bank @@ -235,6 +269,7 @@ mod tests { dbank.get_epoch_accounts_hash_to_serialize(), expected_epoch_accounts_hash, ); + assert_eq!(dbank, bank2); } @@ -265,6 +300,19 @@ mod tests { // Set extra fields bank.fee_rate_governor.lamports_per_signature = 7000; + // Note that epoch_stakes already has two epoch stakes entries for epochs 0 and 1 + // which will also be serialized to the versioned epoch stakes extra field. Those + // entries are of type Stakes so add a new entry for Stakes. + bank.epoch_stakes.insert( + 42, + EpochStakes::from(VersionedEpochStakes::Current { + stakes: Stakes::::default(), + total_stake: 42, + node_id_to_vote_accounts: Arc::::default(), + epoch_authorized_voters: Arc::::default(), + }), + ); + assert_eq!(bank.epoch_stakes.len(), 3); // Serialize let snapshot_storages = bank.get_snapshot_storages(None); @@ -278,18 +326,6 @@ mod tests { ) .unwrap(); - let mut new_epoch_stakes: HashMap = HashMap::new(); - new_epoch_stakes.insert( - 42, - VersionedEpochStakes::Current { - stakes: Stakes::::default(), - total_stake: 42, - node_id_to_vote_accounts: Arc::::default(), - epoch_authorized_voters: Arc::::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..]); @@ -323,13 +359,7 @@ mod tests { ) .unwrap(); - assert_eq!( - dbank.epoch_stakes(42), - Some(&EpochStakes::from( - new_epoch_stakes.get(&42).unwrap().clone() - )) - ); - + assert_eq!(bank.epoch_stakes, dbank.epoch_stakes); assert_eq!( bank.fee_rate_governor.lamports_per_signature, dbank.fee_rate_governor.lamports_per_signature @@ -505,7 +535,7 @@ mod tests { #[cfg_attr( feature = "frozen-abi", derive(AbiExample), - frozen_abi(digest = "AMm4uzGQ6E7fj8MkDjUtFR7kYAjtUyWddXAPLjwaqKqV") + frozen_abi(digest = "CeNFPePrUfgJT2GNr7zYfMQVuJwGyU46bz1Skq7hAPht") )] #[derive(Serialize)] pub struct BankAbiTestWrapper { @@ -530,9 +560,11 @@ mod tests { incremental_capitalization: u64::default(), }; + let mut bank_fields = bank.get_fields_to_serialize(); + let versioned_epoch_stakes = std::mem::take(&mut bank_fields.versioned_epoch_stakes); serde_snapshot::serialize_bank_snapshot_with( serializer, - bank.get_fields_to_serialize(), + bank_fields, BankHashStats::default(), AccountsDeltaHash(Hash::new_unique()), AccountsHash(Hash::new_unique()), @@ -541,6 +573,7 @@ mod tests { lamports_per_signature: bank.fee_rate_governor.lamports_per_signature, incremental_snapshot_persistence: Some(&incremental_snapshot_persistence), epoch_accounts_hash: Some(EpochAccountsHash::new(Hash::new_unique())), + versioned_epoch_stakes, }, StoredMetaWriteVersion::default(), ) diff --git a/runtime/src/epoch_stakes.rs b/runtime/src/epoch_stakes.rs index 63f40218b482c3..015daabe7f86c3 100644 --- a/runtime/src/epoch_stakes.rs +++ b/runtime/src/epoch_stakes.rs @@ -124,8 +124,9 @@ impl EpochStakes { } } +#[cfg_attr(feature = "frozen-abi", derive(AbiExample, AbiEnumVisitor))] #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub(crate) enum VersionedEpochStakes { +pub enum VersionedEpochStakes { Current { stakes: Stakes, total_stake: u64, @@ -152,11 +153,76 @@ impl From for EpochStakes { } } +/// Only the `StakesEnum::Delegations` variant is unable to be serialized as a +/// `StakesEnum::Stakes` variant, so leave those entries and split off the other +/// epoch stakes enum variants into a new map which will be serialized into the +/// new `versioned_epoch_stakes` snapshot field. After a cluster transitions to +/// serializing epoch stakes in the new format, `StakesEnum::Delegations` +/// variants for recent epochs will no longer be created and can be deprecated. +pub(crate) fn split_epoch_stakes( + bank_epoch_stakes: HashMap, +) -> ( + HashMap, + HashMap, +) { + let mut old_epoch_stakes = HashMap::new(); + let mut versioned_epoch_stakes = HashMap::new(); + for (epoch, epoch_stakes) in bank_epoch_stakes.into_iter() { + let EpochStakes { + stakes, + total_stake, + node_id_to_vote_accounts, + epoch_authorized_voters, + } = epoch_stakes; + match stakes.as_ref() { + StakesEnum::Delegations(_) => { + old_epoch_stakes.insert( + epoch, + EpochStakes { + stakes: stakes.clone(), + total_stake, + node_id_to_vote_accounts, + epoch_authorized_voters, + }, + ); + } + StakesEnum::Accounts(stakes) => { + versioned_epoch_stakes.insert( + epoch, + VersionedEpochStakes::Current { + stakes: Stakes::::from(stakes.clone()), + total_stake, + node_id_to_vote_accounts, + epoch_authorized_voters, + }, + ); + } + StakesEnum::Stakes(stakes) => { + versioned_epoch_stakes.insert( + epoch, + VersionedEpochStakes::Current { + stakes: stakes.clone(), + total_stake, + node_id_to_vote_accounts, + epoch_authorized_voters, + }, + ); + } + } + } + (old_epoch_stakes, versioned_epoch_stakes) +} + #[cfg(test)] pub(crate) mod tests { use { - super::*, solana_sdk::account::AccountSharedData, solana_vote::vote_account::VoteAccount, - solana_vote_program::vote_state::create_account_with_authorized, std::iter, + super::*, + crate::{stake_account::StakeAccount, stakes::StakesCache}, + solana_sdk::{account::AccountSharedData, rent::Rent}, + solana_stake_program::stake_state::{self, Delegation}, + solana_vote::vote_account::VoteAccount, + solana_vote_program::vote_state::{self, create_account_with_authorized}, + std::iter, }; struct VoteAccountInfo { @@ -256,4 +322,167 @@ pub(crate) mod tests { vote_accounts_map.len() as u64 * num_vote_accounts_per_node as u64 * 100 ); } + + fn create_test_stakes() -> Stakes> { + let stakes_cache = StakesCache::new(Stakes::default()); + + let vote_pubkey = Pubkey::new_unique(); + let vote_account = vote_state::create_account_with_authorized( + &Pubkey::new_unique(), + &Pubkey::new_unique(), + &Pubkey::new_unique(), + 0, + 1, + ); + + let stake = 1_000_000_000; + let stake_pubkey = Pubkey::new_unique(); + let stake_account = stake_state::create_account( + &Pubkey::new_unique(), + &vote_pubkey, + &vote_account, + &Rent::default(), + stake, + ); + + stakes_cache.check_and_store(&vote_pubkey, &vote_account, None); + stakes_cache.check_and_store(&stake_pubkey, &stake_account, None); + + let stakes = Stakes::clone(&stakes_cache.stakes()); + + stakes + } + + #[test] + fn test_split_epoch_stakes_empty() { + let bank_epoch_stakes = HashMap::new(); + let (old, versioned) = split_epoch_stakes(bank_epoch_stakes); + assert!(old.is_empty()); + assert!(versioned.is_empty()); + } + + #[test] + fn test_split_epoch_stakes_delegations() { + let mut bank_epoch_stakes = HashMap::new(); + let epoch = 0; + let stakes = Arc::new(StakesEnum::Delegations(create_test_stakes().into())); + let epoch_stakes = EpochStakes { + stakes, + total_stake: 100, + node_id_to_vote_accounts: Arc::new(HashMap::new()), + epoch_authorized_voters: Arc::new(HashMap::new()), + }; + bank_epoch_stakes.insert(epoch, epoch_stakes.clone()); + + let (old, versioned) = split_epoch_stakes(bank_epoch_stakes); + + assert_eq!(old.len(), 1); + assert_eq!(old.get(&epoch), Some(&epoch_stakes)); + assert!(versioned.is_empty()); + } + + #[test] + fn test_split_epoch_stakes_accounts() { + let mut bank_epoch_stakes = HashMap::new(); + let epoch = 0; + let test_stakes = create_test_stakes(); + let stakes = Arc::new(StakesEnum::Accounts(test_stakes.clone())); + let epoch_stakes = EpochStakes { + stakes, + total_stake: 100, + node_id_to_vote_accounts: Arc::new(HashMap::new()), + epoch_authorized_voters: Arc::new(HashMap::new()), + }; + bank_epoch_stakes.insert(epoch, epoch_stakes.clone()); + + let (old, versioned) = split_epoch_stakes(bank_epoch_stakes); + + assert!(old.is_empty()); + assert_eq!(versioned.len(), 1); + assert_eq!( + versioned.get(&epoch), + Some(&VersionedEpochStakes::Current { + stakes: Stakes::::from(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, + }) + ); + } + + #[test] + fn test_split_epoch_stakes_stakes() { + let mut bank_epoch_stakes = HashMap::new(); + let epoch = 0; + let test_stakes: Stakes = create_test_stakes().into(); + let stakes = Arc::new(StakesEnum::Stakes(test_stakes.clone())); + let epoch_stakes = EpochStakes { + stakes, + total_stake: 100, + node_id_to_vote_accounts: Arc::new(HashMap::new()), + epoch_authorized_voters: Arc::new(HashMap::new()), + }; + bank_epoch_stakes.insert(epoch, epoch_stakes.clone()); + + let (old, versioned) = split_epoch_stakes(bank_epoch_stakes); + + assert!(old.is_empty()); + assert_eq!(versioned.len(), 1); + assert_eq!( + versioned.get(&epoch), + Some(&VersionedEpochStakes::Current { + stakes: 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, + }) + ); + } + + #[test] + fn test_split_epoch_stakes_mixed() { + let mut bank_epoch_stakes = HashMap::new(); + + // Delegations + let epoch1 = 0; + let stakes1 = Arc::new(StakesEnum::Delegations(Stakes::default())); + let epoch_stakes1 = EpochStakes { + stakes: stakes1, + total_stake: 100, + node_id_to_vote_accounts: Arc::new(HashMap::new()), + epoch_authorized_voters: Arc::new(HashMap::new()), + }; + bank_epoch_stakes.insert(epoch1, epoch_stakes1); + + // Accounts + let epoch2 = 1; + let stakes2 = Arc::new(StakesEnum::Accounts(Stakes::default())); + let epoch_stakes2 = EpochStakes { + stakes: stakes2, + total_stake: 200, + node_id_to_vote_accounts: Arc::new(HashMap::new()), + epoch_authorized_voters: Arc::new(HashMap::new()), + }; + bank_epoch_stakes.insert(epoch2, epoch_stakes2); + + // Stakes + let epoch3 = 2; + let stakes3 = Arc::new(StakesEnum::Stakes(Stakes::default())); + let epoch_stakes3 = EpochStakes { + stakes: stakes3, + total_stake: 300, + node_id_to_vote_accounts: Arc::new(HashMap::new()), + epoch_authorized_voters: Arc::new(HashMap::new()), + }; + bank_epoch_stakes.insert(epoch3, epoch_stakes3); + + let (old, versioned) = split_epoch_stakes(bank_epoch_stakes); + + assert_eq!(old.len(), 1); + assert!(old.contains_key(&epoch1)); + + assert_eq!(versioned.len(), 2); + assert!(versioned.contains_key(&epoch2)); + assert!(versioned.contains_key(&epoch3)); + } } diff --git a/runtime/src/serde_snapshot.rs b/runtime/src/serde_snapshot.rs index 72ee7c97d76f6f..4bef54eaa9670d 100644 --- a/runtime/src/serde_snapshot.rs +++ b/runtime/src/serde_snapshot.rs @@ -408,12 +408,13 @@ 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)] pub struct ExtraFieldsToSerialize<'a> { pub lamports_per_signature: u64, pub incremental_snapshot_persistence: Option<&'a BankIncrementalSnapshotPersistence>, pub epoch_accounts_hash: Option, - // (v2.1) versioned_epoch_stakes + pub versioned_epoch_stakes: HashMap, } fn deserialize_bank_fields( @@ -683,15 +684,16 @@ impl<'a> Serialize for SerializableBankAndStorage<'a> { S: serde::ser::Serializer, { let slot = self.bank.slot(); - let fields = self.bank.get_fields_to_serialize(); + let mut bank_fields = self.bank.get_fields_to_serialize(); let accounts_db = &self.bank.rc.accounts.accounts_db; let bank_hash_stats = accounts_db.get_bank_hash_stats(slot).unwrap(); let accounts_delta_hash = accounts_db.get_accounts_delta_hash(slot).unwrap(); let accounts_hash = accounts_db.get_accounts_hash(slot).unwrap().0; let write_version = accounts_db.write_version.load(Ordering::Acquire); - let lamports_per_signature = fields.fee_rate_governor.lamports_per_signature; + let lamports_per_signature = bank_fields.fee_rate_governor.lamports_per_signature; + let versioned_epoch_stakes = std::mem::take(&mut bank_fields.versioned_epoch_stakes); let bank_fields_to_serialize = ( - SerializableVersionedBank::from(fields), + SerializableVersionedBank::from(bank_fields), SerializableAccountsDb::<'_> { slot, account_storage_entries: self.snapshot_storages, @@ -700,11 +702,12 @@ impl<'a> Serialize for SerializableBankAndStorage<'a> { accounts_hash, write_version, }, - lamports_per_signature, - None::, - self.bank - .get_epoch_accounts_hash_to_serialize() - .map(|epoch_accounts_hash| *epoch_accounts_hash.as_ref()), + ExtraFieldsToSerialize { + lamports_per_signature, + incremental_snapshot_persistence: None, + epoch_accounts_hash: self.bank.get_epoch_accounts_hash_to_serialize(), + versioned_epoch_stakes, + }, ); bank_fields_to_serialize.serialize(serializer) } @@ -723,14 +726,14 @@ impl<'a> Serialize for SerializableBankAndStorageNoExtra<'a> { S: serde::ser::Serializer, { let slot = self.bank.slot(); - let fields = self.bank.get_fields_to_serialize(); + let bank_fields = self.bank.get_fields_to_serialize(); let accounts_db = &self.bank.rc.accounts.accounts_db; let bank_hash_stats = accounts_db.get_bank_hash_stats(slot).unwrap(); let accounts_delta_hash = accounts_db.get_accounts_delta_hash(slot).unwrap(); let accounts_hash = accounts_db.get_accounts_hash(slot).unwrap().0; let write_version = accounts_db.write_version.load(Ordering::Acquire); ( - SerializableVersionedBank::from(fields), + SerializableVersionedBank::from(bank_fields), SerializableAccountsDb::<'_> { slot, account_storage_entries: self.snapshot_storages, diff --git a/runtime/src/snapshot_package.rs b/runtime/src/snapshot_package.rs index a8e7d92967d3c2..1d929227109772 100644 --- a/runtime/src/snapshot_package.rs +++ b/runtime/src/snapshot_package.rs @@ -80,9 +80,10 @@ impl AccountsPackage { let accounts_delta_hash = accounts_db.get_accounts_delta_hash(slot).unwrap(); // SAFETY: Every slot *must* have a BankHashStats entry in AccountsDb. let bank_hash_stats = accounts_db.get_bank_hash_stats(slot).unwrap(); + let bank_fields_to_serialize = bank.get_fields_to_serialize(); SupplementalSnapshotInfo { status_cache_slot_deltas, - bank_fields_to_serialize: bank.get_fields_to_serialize(), + bank_fields_to_serialize, bank_hash_stats, accounts_delta_hash, epoch_accounts_hash: bank.get_epoch_accounts_hash_to_serialize(), diff --git a/runtime/src/snapshot_utils.rs b/runtime/src/snapshot_utils.rs index 039ca1c8f2142a..e897d9bb729b4b 100644 --- a/runtime/src/snapshot_utils.rs +++ b/runtime/src/snapshot_utils.rs @@ -42,6 +42,7 @@ use { collections::{HashMap, HashSet}, fmt, fs, io::{BufReader, BufWriter, Error as IoError, Read, Result as IoResult, Seek, Write}, + mem, num::NonZeroUsize, ops::RangeInclusive, path::{Path, PathBuf}, @@ -826,7 +827,7 @@ fn serialize_snapshot( snapshot_version: SnapshotVersion, snapshot_storages: &[Arc], slot_deltas: &[BankSlotDelta], - bank_fields: BankFieldsToSerialize, + mut bank_fields: BankFieldsToSerialize, bank_hash_stats: BankHashStats, accounts_delta_hash: AccountsDeltaHash, accounts_hash: AccountsHash, @@ -877,10 +878,12 @@ fn serialize_snapshot( .map_err(AddBankSnapshotError::HardLinkStorages)?); let bank_snapshot_serializer = move |stream: &mut BufWriter| -> Result<()> { + let versioned_epoch_stakes = mem::take(&mut bank_fields.versioned_epoch_stakes); let extra_fields = ExtraFieldsToSerialize { lamports_per_signature: bank_fields.fee_rate_governor.lamports_per_signature, incremental_snapshot_persistence: bank_incremental_snapshot_persistence, epoch_accounts_hash, + versioned_epoch_stakes, }; serde_snapshot::serialize_bank_snapshot_into( stream, diff --git a/runtime/src/stake_account.rs b/runtime/src/stake_account.rs index ea4ed6dd0f624c..7bf30c088d9ea3 100644 --- a/runtime/src/stake_account.rs +++ b/runtime/src/stake_account.rs @@ -6,7 +6,7 @@ use { account_utils::StateMut, instruction::InstructionError, pubkey::Pubkey, - stake::state::{Delegation, StakeStateV2}, + stake::state::{Delegation, Stake, StakeStateV2}, }, std::marker::PhantomData, thiserror::Error, @@ -53,6 +53,13 @@ impl StakeAccount { // only wrap a stake-state which is a delegation. self.stake_state.delegation().unwrap() } + + #[inline] + pub(crate) fn stake(&self) -> Stake { + // Safe to unwrap here because StakeAccount will always + // only wrap a stake-state. + self.stake_state.stake().unwrap() + } } impl TryFrom for StakeAccount { diff --git a/runtime/src/stakes.rs b/runtime/src/stakes.rs index 7555238c57b761..4f2dedf1facb07 100644 --- a/runtime/src/stakes.rs +++ b/runtime/src/stakes.rs @@ -529,6 +529,23 @@ impl From> for Stakes { } } +impl From> for Stakes { + fn from(stakes: Stakes) -> Self { + let stake_delegations = stakes + .stake_delegations + .into_iter() + .map(|(pubkey, stake_account)| (pubkey, stake_account.stake())) + .collect(); + Self { + vote_accounts: stakes.vote_accounts, + stake_delegations, + unused: stakes.unused, + epoch: stakes.epoch, + stake_history: stakes.stake_history, + } + } +} + impl From> for Stakes { fn from(stakes: Stakes) -> Self { let stake_delegations = stakes