Skip to content

Commit

Permalink
feat: serialize new versioned epoch stakes bank snapshot field
Browse files Browse the repository at this point in the history
  • Loading branch information
@jstarry
jstarry committed Jul 26, 2024
1 parent 1806617 commit bad8e8a
Show file tree
Hide file tree
Showing 8 changed files with 373 additions and 70 deletions.
8 changes: 6 additions & 2 deletions runtime/src/bank.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -504,6 +504,7 @@ pub struct BankFieldsToSerialize {
pub epoch_stakes: HashMap<Epoch, EpochStakes>,
pub is_delta: bool,
pub accounts_data_len: u64,
pub versioned_epoch_stakes: HashMap<u64, VersionedEpochStakes>,
}

// Can't derive PartialEq because RwLock doesn't implement PartialEq
Expand Down Expand Up @@ -647,6 +648,7 @@ impl BankFieldsToSerialize {
epoch_stakes: HashMap::default(),
is_delta: bool::default(),
accounts_data_len: u64::default(),
versioned_epoch_stakes: HashMap::default(),
}
}
}
Expand Down Expand Up @@ -1674,6 +1676,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),
Expand Down Expand Up @@ -1702,9 +1705,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,
}
}

Expand Down
116 changes: 76 additions & 40 deletions runtime/src/bank/serde_snapshot.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,14 +12,15 @@ mod tests {
runtime_config::RuntimeConfig,
serde_snapshot::{
self, BankIncrementalSnapshotPersistence, ExtraFieldsToSerialize,
SerdeAccountsHash, SerdeIncrementalAccountsHash, SnapshotStreams,
SerdeAccountsHash, SerdeIncrementalAccountsHash, SerializableVersionedBank,
SnapshotStreams,
},
snapshot_bank_utils,
snapshot_utils::{
create_tmp_accounts_dir_for_tests, get_storages_to_serialize, ArchiveFormat,
StorageAndNextAccountsFileId,
},
stakes::Stakes,
stakes::{Stakes, StakesEnum},
},
solana_accounts_db::{
account_storage::{AccountStorageMap, AccountStorageReference},
Expand All @@ -37,8 +38,8 @@ mod tests {
pubkey::Pubkey, stake::state::Stake,
},
std::{
collections::HashMap,
io::{BufReader, BufWriter, Cursor},
mem,
ops::RangeFull,
path::Path,
sync::{atomic::Ordering, Arc},
Expand Down Expand Up @@ -125,7 +126,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();
Expand Down Expand Up @@ -159,25 +160,60 @@ 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);
let serializable_bank = SerializableVersionedBank::from(bank_fields);
serde_snapshot::serialize_bank_snapshot_into(
&mut writer,
serializable_bank,
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
Expand Down Expand Up @@ -236,6 +272,7 @@ mod tests {
dbank.get_epoch_accounts_hash_to_serialize(),
expected_epoch_accounts_hash,
);

assert_eq!(dbank, bank2);
}

Expand Down Expand Up @@ -266,6 +303,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<StakeAccount> so add a new entry for Stakes<Stake>.
bank.epoch_stakes.insert(
42,
EpochStakes::from(VersionedEpochStakes::Current {
stakes: Stakes::<Stake>::default(),
total_stake: 42,
node_id_to_vote_accounts: Arc::<NodeIdToVoteAccounts>::default(),
epoch_authorized_voters: Arc::<EpochAuthorizedVoters>::default(),
}),
);
assert_eq!(bank.epoch_stakes.len(), 3);

// Serialize
let snapshot_storages = bank.get_snapshot_storages(None);
Expand All @@ -279,18 +329,6 @@ 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..]);
Expand Down Expand Up @@ -324,13 +362,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
Expand Down Expand Up @@ -506,7 +538,7 @@ mod tests {
#[cfg_attr(
feature = "frozen-abi",
derive(AbiExample),
frozen_abi(digest = "AMm4uzGQ6E7fj8MkDjUtFR7kYAjtUyWddXAPLjwaqKqV")
frozen_abi(digest = "CeNFPePrUfgJT2GNr7zYfMQVuJwGyU46bz1Skq7hAPht")
)]
#[derive(Serialize)]
pub struct BankAbiTestWrapper {
Expand All @@ -531,6 +563,9 @@ mod tests {
incremental_capitalization: u64::default(),
};

let bank_fields = bank.get_fields_to_serialize();
let versioned_epoch_stakes = bank_fields.versioned_epoch_stakes;
let serializable_bank = SerializableVersionedBank::from(bank_fields);
serde_snapshot::serialize_bank_snapshot_with(
serializer,
bank.get_fields_to_serialize(),
Expand All @@ -542,6 +577,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(),
)
Expand Down
Loading

0 comments on commit bad8e8a

Please sign in to comment.