Skip to content

Commit

Permalink
v2.0: fix: borrow stakes delegation during snapshot serialization (ba…
Browse files Browse the repository at this point in the history
…ckport of #2455) (#2500)
  • Loading branch information
mergify[bot] authored Aug 13, 2024
1 parent 858279a commit e86b3cd
Show file tree
Hide file tree
Showing 7 changed files with 239 additions and 102 deletions.
5 changes: 2 additions & 3 deletions runtime/src/bank.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2296,8 +2296,7 @@ impl Bank {
.fold(
HashSet::default,
|mut voter_pubkeys, (_stake_pubkey, stake_account)| {
let delegation = stake_account.delegation();
voter_pubkeys.insert(delegation.voter_pubkey);
voter_pubkeys.insert(stake_account.delegation().voter_pubkey);
voter_pubkeys
},
)
Expand Down Expand Up @@ -2361,7 +2360,7 @@ impl Bank {
};
if let Some(reward_calc_tracer) = reward_calc_tracer.as_ref() {
let delegation =
InflationPointCalculationEvent::Delegation(delegation, solana_vote_program);
InflationPointCalculationEvent::Delegation(*delegation, solana_vote_program);
let event = RewardCalculationEvent::Staking(stake_pubkey, &delegation);
reward_calc_tracer(&event);
}
Expand Down
6 changes: 2 additions & 4 deletions runtime/src/bank/partitioned_epoch_rewards/calculation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -378,10 +378,9 @@ impl Bank {
let stake_pubkey = **stake_pubkey;
let stake_account = (*stake_account).to_owned();

let delegation = stake_account.delegation();
let vote_pubkey = stake_account.delegation().voter_pubkey;
let (mut stake_account, stake_state) =
<(AccountSharedData, StakeStateV2)>::from(stake_account);
let vote_pubkey = delegation.voter_pubkey;
let vote_account = get_vote_account(&vote_pubkey)?;
if vote_account.owner() != &solana_vote_program {
return None;
Expand Down Expand Up @@ -501,8 +500,7 @@ impl Bank {
stake_delegations
.par_iter()
.map(|(_stake_pubkey, stake_account)| {
let delegation = stake_account.delegation();
let vote_pubkey = delegation.voter_pubkey;
let vote_pubkey = stake_account.delegation().voter_pubkey;

let Some(vote_account) = get_vote_account(&vote_pubkey) else {
return 0;
Expand Down
2 changes: 1 addition & 1 deletion runtime/src/bank/serde_snapshot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -488,7 +488,7 @@ mod tests {
#[cfg_attr(
feature = "frozen-abi",
derive(AbiExample),
frozen_abi(digest = "6riNuebfnAUpS2e3GYb5G8udH5PoEtep48ULchLjRDCB")
frozen_abi(digest = "AzsxaiGEcf3Bxwna5DPS3w5DP4JjyxPDfHnYTvUfyZ5f")
)]
#[derive(Serialize)]
pub struct BankAbiTestWrapper {
Expand Down
4 changes: 2 additions & 2 deletions runtime/src/stake_account.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,10 +48,10 @@ impl<T> StakeAccount<T> {

impl StakeAccount<Delegation> {
#[inline]
pub(crate) fn delegation(&self) -> Delegation {
pub(crate) fn delegation(&self) -> &Delegation {
// Safe to unwrap here because StakeAccount<Delegation> will always
// only wrap a stake-state which is a delegation.
self.stake_state.delegation().unwrap()
self.stake_state.delegation_ref().unwrap()
}
}

Expand Down
106 changes: 14 additions & 92 deletions runtime/src/stakes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ use {
thiserror::Error,
};

mod serde_stakes;
pub(crate) use serde_stakes::serde_stakes_enum_compat;

#[derive(Debug, Error)]
pub enum Error {
#[error("Invalid delegation: {0}")]
Expand Down Expand Up @@ -267,7 +270,7 @@ impl Stakes<StakeAccount> {
let stake_account = StakeAccount::try_from(stake_account)?;
// Sanity check that the delegation is consistent with what is
// stored in the account.
if stake_account.delegation() == *delegation {
if stake_account.delegation() == delegation {
map.insert(*pubkey, stake_account);
Ok(map)
} else {
Expand Down Expand Up @@ -499,12 +502,15 @@ impl StakesEnum {
}
}

/// 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<Delegation> {
fn from(stakes: Stakes<StakeAccount>) -> Self {
let stake_delegations = stakes
.stake_delegations
.into_iter()
.map(|(pubkey, stake_account)| (pubkey, stake_account.delegation()))
.map(|(pubkey, stake_account)| (pubkey, *stake_account.delegation()))
.collect();
Self {
vote_accounts: stakes.vote_accounts,
Expand All @@ -516,6 +522,9 @@ impl From<Stakes<StakeAccount>> for Stakes<Delegation> {
}
}

/// This conversion is memory intensive so should only be used in development
/// contexts.
#[cfg(feature = "dev-context-only-utils")]
impl From<Stakes<Stake>> for Stakes<Delegation> {
fn from(stakes: Stakes<Stake>) -> Self {
let stake_delegations = stakes
Expand All @@ -533,6 +542,9 @@ impl From<Stakes<Stake>> for Stakes<Delegation> {
}
}

/// This conversion is memory intensive so should only be used in development
/// contexts.
#[cfg(feature = "dev-context-only-utils")]
impl From<StakesEnum> for Stakes<Delegation> {
fn from(stakes: StakesEnum) -> Self {
match stakes {
Expand Down Expand Up @@ -576,36 +588,6 @@ impl PartialEq<StakesEnum> for StakesEnum {
}
}

// 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::Delegations(stakes) => stakes.serialize(serializer),
stakes => {
let stakes = Stakes::<Delegation>::from(stakes.clone());
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)))
}
}

fn refresh_vote_accounts(
thread_pool: &ThreadPool,
epoch: Epoch,
Expand Down Expand Up @@ -651,7 +633,6 @@ fn refresh_vote_accounts(
pub(crate) mod tests {
use {
super::*,
rand::Rng,
rayon::ThreadPoolBuilder,
solana_sdk::{account::WritableAccount, pubkey::Pubkey, rent::Rent, stake},
solana_stake_program::stake_state,
Expand Down Expand Up @@ -1078,63 +1059,4 @@ pub(crate) 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, None);
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, None);
}
}
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>::from(stakes);
assert!(stakes.vote_accounts.as_ref().len() >= 5);
assert!(stakes.stake_delegations.len() >= 50);
let other = match &*other.stakes {
StakesEnum::Accounts(_) | StakesEnum::Stakes(_) => panic!("wrong type!"),
StakesEnum::Delegations(delegations) => delegations,
};
assert_eq!(other, &stakes)
}
}
Loading

0 comments on commit e86b3cd

Please sign in to comment.