blockworks-foundation
/
solana
mirror of https://github.com/blockworks-foundation/solana.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

enforces type-safety in Stakes<StakeAccount> using phantom data 

StakeAccount<Delegation> can only wrap a stake-state which is a
Delegation; whereas StakeAccount<()> wraps any account with stake state.

As a result, StakeAccount::<Delegation>::delegation() will return
Delegation instead of Option<Delegation>.
This commit is contained in:
behzad nouri 2022-04-15 10:17:13 -04:00
parent f937fcbd95
commit 108aa23d90
3 changed files with 103 additions and 60 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
runtime/src/bank.rs
View File

@ -1265,7 +1265,8 @@ pub struct Bank {
struct VoteWithStakeDelegations {
vote_state: Arc<VoteState>,
vote_account: AccountSharedData,
delegations: Vec<(Pubkey, StakeAccount)>,
// TODO: use StakeAccount<Delegation> once the old code is deleted.
delegations: Vec<(Pubkey, StakeAccount<()>)>,
}
struct LoadVoteAndStakeAccountsResult {
@ -2000,7 +2001,7 @@ impl Bank {
// from Stakes<Delegation> by reading the full account state from
// accounts-db. Note that it is crucial that these accounts are loaded
// at the right slot and match precisely with serialized Delegations.
let stakes = Stakes::<StakeAccount>::new(&fields.stakes, |pubkey| {
let stakes = Stakes::new(&fields.stakes, |pubkey| {
let (account, _slot) = bank_rc.accounts.load_with_fixed_root(&ancestors, pubkey)?;
Some(account)
})
@ -2648,7 +2649,7 @@ impl Bank {
stake_delegations
.into_par_iter()
.for_each(|(stake_pubkey, stake_account)| {
let delegation = stake_account.delegation().unwrap();
let delegation = stake_account.delegation();
let vote_pubkey = &delegation.voter_pubkey;
if invalid_vote_keys.contains_key(vote_pubkey) {
return;
@ -2661,7 +2662,7 @@ impl Bank {
return;
}
};
let stake_account = match StakeAccount::try_from(stake_account) {
let stake_account = match StakeAccount::<()>::try_from(stake_account) {
Ok(stake_account) => stake_account,
Err(stake_account::Error::InvalidOwner { .. }) => {
invalid_stake_keys
@ -2673,6 +2674,15 @@ impl Bank {
.insert(*stake_pubkey, InvalidCacheEntryReason::BadState);
return;
}
Err(stake_account::Error::InvalidDelegation(_)) => {
// This should not happen.
error!(
"Unexpected code path! StakeAccount<()> \
should not check if stake-state is a \
Delegation."
);
return;
}
};
let stake_delegation = (*stake_pubkey, stake_account);
let mut vote_delegations = if let Some(vote_delegations) =
@ -2761,7 +2771,7 @@ impl Bank {
.fold(
HashSet::default,
|mut voter_pubkeys, (_stake_pubkey, stake_account)| {
let delegation = stake_account.delegation().unwrap();
let delegation = stake_account.delegation();
voter_pubkeys.insert(delegation.voter_pubkey);
voter_pubkeys
},
@ -2823,8 +2833,8 @@ impl Bank {
.collect()
});
// Join stake accounts with vote-accounts.
let push_stake_delegation = |(stake_pubkey, stake_account): (&Pubkey, &StakeAccount)| {
let delegation = stake_account.delegation().unwrap();
let push_stake_delegation = |(stake_pubkey, stake_account): (&Pubkey, &StakeAccount<_>)| {
let delegation = stake_account.delegation();
let mut vote_delegations =
match vote_with_stake_delegations_map.get_mut(&delegation.voter_pubkey) {
Some(vote_delegations) => vote_delegations,
@ -2836,7 +2846,8 @@ impl Bank {
let event = RewardCalculationEvent::Staking(stake_pubkey, &delegation);
reward_calc_tracer(&event);
}
let stake_delegation = (*stake_pubkey, stake_account.clone());
let stake_account = StakeAccount::from(stake_account.clone());
let stake_delegation = (*stake_pubkey, stake_account);
vote_delegations.delegations.push(stake_delegation);
};
thread_pool.install(|| {

76
runtime/src/stake_account.rs
View File

@ -8,39 +8,54 @@ use {
pubkey::Pubkey,
stake::state::{Delegation, StakeState},
},
std::marker::PhantomData,
thiserror::Error,
};
/// An account and a stake state deserialized from the account.
/// Generic type T enforces type-saftey so that StakeAccount<Delegation> can
/// only wrap a stake-state which is a Delegation; whereas StakeAccount<()>
/// wraps any account with stake state.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct StakeAccount(AccountSharedData, StakeState);
pub struct StakeAccount<T> {
account: AccountSharedData,
stake_state: StakeState,
_phantom: PhantomData<T>,
}
#[derive(Debug, Error)]
#[allow(clippy::enum_variant_names)]
pub enum Error {
#[error(transparent)]
InstructionError(#[from] InstructionError),
#[error("Invalid delegation: {0:?}")]
InvalidDelegation(StakeState),
#[error("Invalid stake account owner: {owner:?}")]
InvalidOwner { owner: Pubkey },
}
impl StakeAccount {
impl<T> StakeAccount<T> {
#[inline]
pub(crate) fn lamports(&self) -> u64 {
self.0.lamports()
self.account.lamports()
}
#[inline]
pub(crate) fn stake_state(&self) -> &StakeState {
&self.1
}
#[inline]
pub(crate) fn delegation(&self) -> Option<Delegation> {
self.1.delegation()
&self.stake_state
}
}
impl TryFrom<AccountSharedData> for StakeAccount {
impl StakeAccount<Delegation> {
#[inline]
pub(crate) fn delegation(&self) -> Delegation {
// Safe to unwrap here becasue StakeAccount<Delegation> will always
// only wrap a stake-state which is a delegation.
self.stake_state.delegation().unwrap()
}
}
impl TryFrom<AccountSharedData> for StakeAccount<()> {
type Error = Error;
fn try_from(account: AccountSharedData) -> Result<Self, Self::Error> {
if account.owner() != &solana_stake_program::id() {
@ -49,18 +64,49 @@ impl TryFrom<AccountSharedData> for StakeAccount {
});
}
let stake_state = account.state()?;
Ok(Self(account, stake_state))
Ok(Self {
account,
stake_state,
_phantom: PhantomData::default(),
})
}
}
impl From<StakeAccount> for (AccountSharedData, StakeState) {
fn from(stake_account: StakeAccount) -> Self {
(stake_account.0, stake_account.1)
impl TryFrom<AccountSharedData> for StakeAccount<Delegation> {
type Error = Error;
fn try_from(account: AccountSharedData) -> Result<Self, Self::Error> {
let stake_account = StakeAccount::<()>::try_from(account)?;
if stake_account.stake_state.delegation().is_none() {
return Err(Error::InvalidDelegation(stake_account.stake_state));
}
Ok(Self {
account: stake_account.account,
stake_state: stake_account.stake_state,
_phantom: PhantomData::default(),
})
}
}
impl From<StakeAccount<Delegation>> for StakeAccount<()> {
#[inline]
fn from(stake_account: StakeAccount<Delegation>) -> Self {
Self {
account: stake_account.account,
stake_state: stake_account.stake_state,
_phantom: PhantomData::default(),
}
}
}
impl<T> From<StakeAccount<T>> for (AccountSharedData, StakeState) {
#[inline]
fn from(stake_account: StakeAccount<T>) -> Self {
(stake_account.account, stake_account.stake_state)
}
}
#[cfg(RUSTC_WITH_SPECIALIZATION)]
impl AbiExample for StakeAccount {
impl AbiExample for StakeAccount<Delegation> {
fn example() -> Self {
use solana_sdk::{
account::Account,

60
runtime/src/stakes.rs
View File

@ -2,7 +2,7 @@
//! node stakes
use {
crate::{
stake_account::{self, StakeAccount},
stake_account,
stake_history::StakeHistory,
vote_account::{VoteAccount, VoteAccounts},
},
@ -43,6 +43,8 @@ pub enum InvalidCacheEntryReason {
WrongOwner,
}
type StakeAccount = stake_account::StakeAccount<Delegation>;
#[derive(Default, Debug, AbiExample)]
pub(crate) struct StakesCache(RwLock<Stakes<StakeAccount>>);
@ -200,7 +202,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() == Some(*delegation) {
if stake_account.delegation() == *delegation {
Ok((*pubkey, stake_account))
} else {
Err(Error::InvalidDelegation(*pubkey))
@ -237,7 +239,7 @@ impl Stakes<StakeAccount> {
stake_delegations
.par_iter()
.fold(StakeActivationStatus::default, |acc, stake_account| {
let delegation = stake_account.delegation().unwrap();
let delegation = stake_account.delegation();
acc + delegation
.stake_activating_and_deactivating(self.epoch, Some(&self.stake_history))
})
@ -251,7 +253,7 @@ impl Stakes<StakeAccount> {
stake_delegations
.par_iter()
.fold(HashMap::default, |mut delegated_stakes, stake_account| {
let delegation = stake_account.delegation().unwrap();
let delegation = stake_account.delegation();
let entry = delegated_stakes.entry(delegation.voter_pubkey).or_default();
*entry += delegation.stake(self.epoch, Some(&self.stake_history));
delegated_stakes
@ -280,7 +282,7 @@ impl Stakes<StakeAccount> {
) -> u64 {
self.stake_delegations
.values()
.map(|stake_account| stake_account.delegation().unwrap())
.map(StakeAccount::delegation)
.filter(|delegation| &delegation.voter_pubkey == voter_pubkey)
.map(|delegation| delegation.stake(epoch, Some(stake_history)))
.sum()
@ -288,7 +290,7 @@ impl Stakes<StakeAccount> {
/// Sum the lamports of the vote accounts and the delegated stake
pub fn vote_balance_and_staked(&self) -> u64 {
let get_stake = |stake_account: &StakeAccount| stake_account.delegation().unwrap().stake;
let get_stake = |stake_account: &StakeAccount| stake_account.delegation().stake;
let get_lamports = |(_, (_, vote_account)): (_, &(_, VoteAccount))| vote_account.lamports();
self.stake_delegations.values().map(get_stake).sum::<u64>()
@ -301,7 +303,7 @@ impl Stakes<StakeAccount> {
pub fn remove_stake_delegation(&mut self, stake_pubkey: &Pubkey) {
if let Some(stake_account) = self.stake_delegations.remove(stake_pubkey) {
let removed_delegation = stake_account.delegation().unwrap();
let removed_delegation = stake_account.delegation();
let removed_stake = removed_delegation.stake(self.epoch, Some(&self.stake_history));
self.vote_accounts
.sub_stake(&removed_delegation.voter_pubkey, removed_stake);
@ -339,15 +341,13 @@ impl Stakes<StakeAccount> {
.stake_delegations
.get(stake_pubkey)
.map(|stake_account| {
let delegation = stake_account.delegation().unwrap();
let delegation = stake_account.delegation();
(
delegation.voter_pubkey,
delegation.stake(self.epoch, Some(&self.stake_history)),
)
});
let new_delegation = new_stake_account
.as_ref()
.and_then(StakeAccount::delegation);
let new_delegation = new_stake_account.as_ref().map(StakeAccount::delegation);
let new_stake = new_stake_account.as_ref().and_then(|new_stake_account| {
let new_delegation = new_delegation?;
// When account is removed (lamports == 0), this check ensures
@ -408,24 +408,20 @@ impl StakesEnum {
}
}
impl TryFrom<Stakes<StakeAccount>> for Stakes<Delegation> {
type Error = Error;
fn try_from(stakes: Stakes<StakeAccount>) -> Result<Self, Self::Error> {
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)| match stake_account.delegation() {
None => Err(Error::InvalidDelegation(pubkey)),
Some(delegation) => Ok((pubkey, delegation)),
})
.collect::<Result<_, _>>()?;
Ok(Self {
.map(|(pubkey, stake_account)| (pubkey, stake_account.delegation()))
.collect();
Self {
vote_accounts: stakes.vote_accounts,
stake_delegations,
unused: stakes.unused,
epoch: stakes.epoch,
stake_history: stakes.stake_history,
})
}
}
}
@ -451,16 +447,12 @@ impl PartialEq<StakesEnum> for StakesEnum {
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,
}
let stakes = Stakes::<Delegation>::from(stakes.clone());
&stakes == other
}
(Self::Delegations(stakes), Self::Accounts(other)) => {
match Stakes::<Delegation>::try_from(other.clone()) {
Ok(other) => *stakes == other,
Err(_) => false,
}
let other = Stakes::<Delegation>::from(other.clone());
stakes == &other
}
(Self::Delegations(stakes), Self::Delegations(other)) => stakes == other,
}
@ -481,13 +473,7 @@ pub(crate) mod serde_stakes_enum_compat {
{
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));
}
};
let stakes = Stakes::<Delegation>::from(stakes.clone());
stakes.serialize(serializer)
}
StakesEnum::Delegations(stakes) => stakes.serialize(serializer),
@ -966,7 +952,7 @@ pub mod tests {
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();
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 {