use std::{
collections::{HashMap, HashSet, VecDeque},
sync::{
atomic::{AtomicBool, Ordering},
RwLock,
},
};
use rayon::prelude::*;
use rialo_s_account::ReadableAccount;
use rialo_s_clock::Epoch;
use rialo_s_pubkey::Pubkey;
use rialo_s_type_overrides::sync::Arc;
use rialo_stake_manager_interface::instruction::StakeInfo;
pub use rialo_validator_registry_interface::instruction::ValidatorInfo;
pub mod pda;
pub use pda::{derive_self_bond_address, derive_self_bond_address_with_bump, SELF_BOND_SEED};
#[derive(Debug, Clone)]
pub struct StakeCache(Arc<RwLock<StakeCacheData>>);
impl Default for StakeCache {
fn default() -> Self {
Self(Arc::new(RwLock::new(StakeCacheData::default())))
}
}
impl StakeCache {
pub fn new() -> Self {
Self::default()
}
pub fn with_data(data: StakeCacheData) -> Self {
Self(Arc::new(RwLock::new(data)))
}
pub fn from_arc(arc: Arc<RwLock<StakeCacheData>>) -> Self {
Self(arc)
}
pub fn arc_clone(&self) -> Arc<RwLock<StakeCacheData>> {
Arc::clone(&self.0)
}
pub fn read(&self) -> std::sync::RwLockReadGuard<'_, StakeCacheData> {
self.0.read().expect("Failed to acquire read lock")
}
pub fn write(&self) -> std::sync::RwLockWriteGuard<'_, StakeCacheData> {
self.0.write().expect("Failed to acquire write lock")
}
pub fn get_stake_account(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
let data = self.read();
data.stake_accounts.get(pubkey).and_then(|opt| opt.clone())
}
pub fn get_validator_account(&self, pubkey: &Pubkey) -> Option<ValidatorAccount> {
let data = self.read();
data.validator_accounts
.get(pubkey)
.and_then(|opt| opt.clone())
}
pub fn get_all_validator_accounts(&self) -> Vec<(Pubkey, ValidatorAccount)> {
let data = self.read();
data.validator_accounts
.iter()
.filter_map(|(k, v)| v.as_ref().map(|account| (*k, account.clone())))
.collect()
}
pub fn contains_stake_account(&self, pubkey: &Pubkey) -> bool {
let data = self.read();
matches!(data.stake_accounts.get(pubkey), Some(Some(_)))
}
pub fn contains_validator_account(&self, pubkey: &Pubkey) -> bool {
let data = self.read();
matches!(data.validator_accounts.get(pubkey), Some(Some(_)))
}
pub fn insert_stake_account(&self, pubkey: Pubkey, account: StakeAccount) {
let mut data = self.write();
data.stake_accounts.insert(pubkey, Some(account));
data.modified_stake_pubkeys.insert(pubkey);
}
pub fn insert_validator_account(&self, pubkey: Pubkey, account: ValidatorAccount) {
let mut data = self.write();
data.validator_accounts.insert(pubkey, Some(account));
data.modified_validator_pubkeys.insert(pubkey);
}
pub fn tombstone_stake_account(&self, pubkey: Pubkey) {
let mut data = self.write();
data.stake_accounts.insert(pubkey, None);
data.modified_stake_pubkeys.insert(pubkey);
}
pub fn tombstone_validator_account(&self, pubkey: Pubkey) {
let mut data = self.write();
data.validator_accounts.insert(pubkey, None);
data.modified_validator_pubkeys.insert(pubkey);
}
pub fn epoch(&self) -> Epoch {
self.read().epoch
}
pub fn timestamp(&self) -> u64 {
self.read().timestamp
}
pub fn set_epoch(&self, epoch: Epoch) {
self.write().epoch = epoch;
}
pub fn set_timestamp(&self, timestamp: u64) {
self.write().timestamp = timestamp;
}
pub fn check_and_update(&self, pubkey: &Pubkey, account: &impl ReadableAccount) {
let owner = account.owner();
if account.kelvins() == 0 {
if rialo_stake_manager_interface::check_id(owner) {
self.tombstone_stake_account(*pubkey);
} else if rialo_validator_registry_interface::check_id(owner) {
self.tombstone_validator_account(*pubkey);
}
} else if rialo_stake_manager_interface::check_id(owner) {
if let Ok(stake_info) = bincode::deserialize::<StakeInfo>(account.data()) {
self.insert_stake_account(
*pubkey,
StakeAccount {
kelvins: account.kelvins(),
data: stake_info,
},
);
}
} else if rialo_validator_registry_interface::check_id(owner) {
if let Ok(validator_info) = bincode::deserialize::<ValidatorInfo>(account.data()) {
self.insert_validator_account(
*pubkey,
ValidatorAccount {
kelvins: account.kelvins(),
data: validator_info,
},
);
}
}
}
}
#[derive(Debug, Default, Clone)]
pub struct StakeCacheData {
pub stake_accounts: HashMap<Pubkey, Option<StakeAccount>>,
pub validator_accounts: HashMap<Pubkey, Option<ValidatorAccount>>,
pub epoch: Epoch,
pub timestamp: u64,
pub modified_stake_pubkeys: HashSet<Pubkey>,
pub modified_validator_pubkeys: HashSet<Pubkey>,
}
impl StakeCacheData {
pub fn drain_modified(&mut self) -> (HashSet<Pubkey>, HashSet<Pubkey>) {
let stake_pubkeys = std::mem::take(&mut self.modified_stake_pubkeys);
let validator_pubkeys = std::mem::take(&mut self.modified_validator_pubkeys);
(stake_pubkeys, validator_pubkeys)
}
pub fn has_modified(&self) -> bool {
!self.modified_stake_pubkeys.is_empty() || !self.modified_validator_pubkeys.is_empty()
}
}
#[derive(Debug, Clone)]
pub struct StakeHistory(Arc<RwLock<VecDeque<StakeCacheData>>>);
impl Default for StakeHistory {
fn default() -> Self {
Self(Arc::new(RwLock::new(VecDeque::new())))
}
}
impl StakeHistory {
pub fn new() -> Self {
Self::default()
}
pub fn with_entry(data: StakeCacheData) -> Self {
let mut deque = VecDeque::new();
deque.push_back(data);
Self(Arc::new(RwLock::new(deque)))
}
pub fn from_arc(arc: Arc<RwLock<VecDeque<StakeCacheData>>>) -> Self {
Self(arc)
}
pub fn arc_clone(&self) -> Arc<RwLock<VecDeque<StakeCacheData>>> {
Arc::clone(&self.0)
}
pub fn read(&self) -> std::sync::RwLockReadGuard<'_, VecDeque<StakeCacheData>> {
self.0.read().expect("Failed to acquire read lock")
}
pub fn write_lock(&self) -> std::sync::RwLockWriteGuard<'_, VecDeque<StakeCacheData>> {
self.0.write().expect("Failed to acquire write lock")
}
pub fn push_back(&self, data: StakeCacheData) {
self.0
.write()
.expect("Failed to acquire lock")
.push_back(data);
}
pub fn pop_front(&self) -> Option<StakeCacheData> {
self.0.write().expect("Failed to acquire lock").pop_front()
}
pub fn len(&self) -> usize {
self.0.read().expect("Failed to acquire lock").len()
}
pub fn is_empty(&self) -> bool {
self.0.read().expect("Failed to acquire lock").is_empty()
}
pub fn front(&self) -> Option<StakeCacheData> {
self.0
.read()
.expect("Failed to acquire lock")
.front()
.cloned()
}
pub fn back(&self) -> Option<StakeCacheData> {
self.0
.read()
.expect("Failed to acquire lock")
.back()
.cloned()
}
pub fn iter_cloned(&self) -> Vec<StakeCacheData> {
self.0
.read()
.expect("Failed to acquire lock")
.iter()
.cloned()
.collect()
}
}
#[derive(Debug, Clone)]
pub struct StakeAccount {
pub kelvins: u64,
pub data: StakeInfo,
}
#[derive(Debug, Clone)]
pub struct ValidatorAccount {
pub kelvins: u64,
pub data: ValidatorInfo,
}
pub struct StakesHandle {
baseline: StakeCache,
pending: StakeCache,
frozen: StakeHistory,
epoch_rewards_init: Arc<RwLock<Option<EpochRewardsInitRequest>>>,
epoch_stakes_frozen: Arc<AtomicBool>,
epoch_rewards_exists_fn: Arc<dyn Fn(u64) -> bool + Send + Sync>,
}
#[derive(Debug, Clone)]
pub struct EpochRewardsInitRequest {
pub epoch: Epoch,
pub total_rewards: u64,
}
impl std::fmt::Debug for StakesHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("StakesHandle")
.field("baseline", &self.baseline)
.field("pending", &self.pending)
.field("frozen", &self.frozen)
.field("epoch_rewards_init", &self.epoch_rewards_init)
.field("epoch_stakes_frozen", &self.epoch_stakes_frozen)
.field("epoch_rewards_exists_fn", &"<callback>")
.finish()
}
}
impl Clone for StakesHandle {
fn clone(&self) -> Self {
Self {
baseline: self.baseline.clone(),
pending: self.pending.clone(),
frozen: self.frozen.clone(),
epoch_rewards_init: self.epoch_rewards_init.clone(),
epoch_stakes_frozen: Arc::clone(&self.epoch_stakes_frozen),
epoch_rewards_exists_fn: Arc::clone(&self.epoch_rewards_exists_fn),
}
}
}
impl Default for StakesHandle {
fn default() -> Self {
Self {
baseline: StakeCache::default(),
pending: StakeCache::default(),
frozen: StakeHistory::default(),
epoch_rewards_init: Arc::new(RwLock::new(None)),
epoch_stakes_frozen: Arc::new(AtomicBool::new(false)),
epoch_rewards_exists_fn: Arc::new(|_| false),
}
}
}
impl StakesHandle {
pub fn new_shared(
baseline: StakeCache,
pending: StakeCache,
frozen: StakeHistory,
epoch_rewards_exists_fn: Arc<dyn Fn(u64) -> bool + Send + Sync>,
) -> Self {
Self {
baseline,
pending,
frozen,
epoch_rewards_init: Arc::new(RwLock::new(None)),
epoch_stakes_frozen: Arc::new(AtomicBool::new(false)),
epoch_rewards_exists_fn,
}
}
pub fn epoch_rewards_exists(&self, epoch: u64) -> bool {
(self.epoch_rewards_exists_fn)(epoch)
}
pub fn set_epoch_stakes_frozen(&self) {
self.epoch_stakes_frozen.store(true, Ordering::Release);
}
pub fn take_epoch_stakes_frozen(&self) -> bool {
self.epoch_stakes_frozen.swap(false, Ordering::AcqRel)
}
pub fn is_epoch_stakes_frozen(&self) -> bool {
self.epoch_stakes_frozen.load(Ordering::Acquire)
}
pub fn get_stake_account_from_pending(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
{
let pending_data = self.pending.read();
if let Some(value) = pending_data.stake_accounts.get(pubkey) {
return value.clone(); }
}
{
let frozen_data = self.frozen.read();
for frozen_entry in frozen_data.iter().rev() {
if let Some(value) = frozen_entry.stake_accounts.get(pubkey) {
return value.clone();
}
}
}
{
let baseline_data = self.baseline.read();
baseline_data
.stake_accounts
.get(pubkey)
.and_then(|v| v.clone())
}
}
pub fn get_validator_account_from_pending(&self, pubkey: &Pubkey) -> Option<ValidatorAccount> {
{
let pending_data = self.pending.read();
if let Some(value) = pending_data.validator_accounts.get(pubkey) {
return value.clone(); }
}
{
let frozen_data = self.frozen.read();
for frozen_entry in frozen_data.iter().rev() {
if let Some(value) = frozen_entry.validator_accounts.get(pubkey) {
return value.clone();
}
}
}
{
let baseline_data = self.baseline.read();
baseline_data
.validator_accounts
.get(pubkey)
.and_then(|v| v.clone())
}
}
pub fn get_all_validator_accounts_from_pending(&self) -> Vec<(Pubkey, ValidatorAccount)> {
let mut result: HashMap<Pubkey, Option<ValidatorAccount>> = HashMap::new();
{
let baseline_data = self.baseline.read();
for (pubkey, value) in baseline_data.validator_accounts.iter() {
result.insert(*pubkey, value.clone());
}
}
{
let frozen_data = self.frozen.read();
for frozen_entry in frozen_data.iter() {
for (pubkey, value) in frozen_entry.validator_accounts.iter() {
result.insert(*pubkey, value.clone());
}
}
}
{
let pending_data = self.pending.read();
for (pubkey, value) in pending_data.validator_accounts.iter() {
result.insert(*pubkey, value.clone());
}
}
let mut sorted: Vec<_> = result
.into_iter()
.filter_map(|(k, v)| v.map(|account| (k, account)))
.collect();
sorted.sort_by_key(|(pubkey, _)| *pubkey);
sorted
}
pub fn get_stake_account_from_last_frozen(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
{
let frozen_data = self.frozen.read();
for frozen_entry in frozen_data.iter().rev() {
if let Some(value) = frozen_entry.stake_accounts.get(pubkey) {
return value.clone();
}
}
}
{
let baseline_data = self.baseline.read();
baseline_data
.stake_accounts
.get(pubkey)
.and_then(|v| v.clone())
}
}
pub fn get_validator_account_from_last_frozen(
&self,
pubkey: &Pubkey,
) -> Option<ValidatorAccount> {
{
let frozen_data = self.frozen.read();
for frozen_entry in frozen_data.iter().rev() {
if let Some(value) = frozen_entry.validator_accounts.get(pubkey) {
return value.clone();
}
}
}
{
let baseline_data = self.baseline.read();
baseline_data
.validator_accounts
.get(pubkey)
.and_then(|v| v.clone())
}
}
pub fn get_all_validator_accounts_from_last_frozen(&self) -> Vec<(Pubkey, ValidatorAccount)> {
let mut result: HashMap<Pubkey, Option<ValidatorAccount>> = HashMap::new();
{
let baseline_data = self.baseline.read();
for (pubkey, value) in baseline_data.validator_accounts.iter() {
result.insert(*pubkey, value.clone());
}
}
{
let frozen_data = self.frozen.read();
for frozen_entry in frozen_data.iter() {
for (pubkey, value) in frozen_entry.validator_accounts.iter() {
result.insert(*pubkey, value.clone());
}
}
}
let mut sorted: Vec<_> = result
.into_iter()
.filter_map(|(k, v)| v.map(|account| (k, account)))
.collect();
sorted.sort_by_key(|(pubkey, _)| *pubkey);
sorted
}
pub fn get_stake_account_from_first_frozen(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
{
let frozen_data = self.frozen.read();
if let Some(first_frozen) = frozen_data.front() {
if let Some(value) = first_frozen.stake_accounts.get(pubkey) {
return value.clone();
}
}
}
{
let baseline_data = self.baseline.read();
baseline_data
.stake_accounts
.get(pubkey)
.and_then(|v| v.clone())
}
}
pub fn get_all_stake_accounts_from_pending(&self) -> Vec<(Pubkey, StakeAccount)> {
let mut result: HashMap<Pubkey, Option<StakeAccount>> = HashMap::new();
{
let baseline_data = self.baseline.read();
for (pubkey, value) in baseline_data.stake_accounts.iter() {
result.insert(*pubkey, value.clone());
}
}
{
let frozen_data = self.frozen.read();
for frozen_entry in frozen_data.iter() {
for (pubkey, value) in frozen_entry.stake_accounts.iter() {
result.insert(*pubkey, value.clone());
}
}
}
{
let pending_data = self.pending.read();
for (pubkey, value) in pending_data.stake_accounts.iter() {
result.insert(*pubkey, value.clone());
}
}
let mut sorted: Vec<_> = result
.into_iter()
.filter_map(|(k, v)| v.map(|account| (k, account)))
.collect();
sorted.sort_by_key(|(pubkey, _)| *pubkey);
sorted
}
pub fn get_all_stake_accounts_from_first_frozen(&self) -> Vec<(Pubkey, StakeAccount)> {
let mut result: HashMap<Pubkey, Option<StakeAccount>> = HashMap::new();
{
let baseline_data = self.baseline.read();
for (pubkey, value) in baseline_data.stake_accounts.iter() {
result.insert(*pubkey, value.clone());
}
}
{
let frozen_data = self.frozen.read();
if let Some(first_frozen) = frozen_data.front() {
for (pubkey, value) in first_frozen.stake_accounts.iter() {
result.insert(*pubkey, value.clone());
}
}
}
let mut sorted: Vec<_> = result
.into_iter()
.filter_map(|(k, v)| v.map(|account| (k, account)))
.collect();
sorted.sort_by_key(|(pubkey, _)| *pubkey);
sorted
}
pub fn get_all_stake_accounts_from_frozen_epoch(
&self,
target_epoch: Epoch,
) -> Vec<(Pubkey, StakeAccount)> {
let mut result: HashMap<Pubkey, Option<StakeAccount>> = HashMap::new();
{
let baseline_data = self.baseline.read();
for (pubkey, value) in baseline_data.stake_accounts.iter() {
result.insert(*pubkey, value.clone());
}
}
{
let frozen_data = self.frozen.read();
for frozen_entry in frozen_data.iter() {
if frozen_entry.epoch > target_epoch {
break; }
for (pubkey, value) in frozen_entry.stake_accounts.iter() {
result.insert(*pubkey, value.clone());
}
}
}
let mut sorted: Vec<_> = result
.into_iter()
.filter_map(|(k, v)| v.map(|account| (k, account)))
.collect();
sorted.sort_by_key(|(pubkey, _)| *pubkey);
sorted
}
pub fn get_all_stake_accounts_from_baseline(&self) -> Vec<(Pubkey, StakeAccount)> {
let baseline_data = self.baseline.read();
let mut sorted: Vec<_> = baseline_data
.stake_accounts
.iter()
.filter_map(|(k, v)| v.as_ref().map(|account| (*k, account.clone())))
.collect();
sorted.sort_by_key(|(pubkey, _)| *pubkey);
sorted
}
pub fn get_all_validator_accounts_from_baseline(&self) -> Vec<(Pubkey, ValidatorAccount)> {
let baseline_data = self.baseline.read();
let mut sorted: Vec<_> = baseline_data
.validator_accounts
.iter()
.filter_map(|(k, v)| v.as_ref().map(|account| (*k, account.clone())))
.collect();
sorted.sort_by_key(|(pubkey, _)| *pubkey);
sorted
}
pub fn freeze_stakes(&self) {
let frozen_data = {
let mut pending_guard = self.pending.write();
let frozen_data = std::mem::take(&mut *pending_guard);
pending_guard.epoch = frozen_data.epoch + 1;
pending_guard.timestamp = frozen_data.timestamp;
frozen_data
};
self.frozen.push_back(frozen_data);
self.epoch_stakes_frozen.store(true, Ordering::Release);
}
pub fn pending_epoch(&self) -> Epoch {
self.pending.epoch()
}
pub fn set_pending_epoch(&self, epoch: Epoch) {
self.pending.set_epoch(epoch);
}
pub fn set_pending_timestamp(&self, timestamp: u64) {
self.pending.set_timestamp(timestamp);
}
pub fn last_frozen_timestamp(&self) -> Option<u64> {
self.frozen.read().back().map(|data| data.timestamp)
}
pub fn last_frozen_epoch(&self) -> Option<Epoch> {
self.frozen.read().back().map(|data| data.epoch)
}
pub fn pending_timestamp(&self) -> u64 {
self.pending.read().timestamp
}
pub fn push_frozen(&self, data: StakeCacheData) {
self.frozen.push_back(data);
}
pub fn frozen_len(&self) -> usize {
self.frozen.len()
}
pub fn front_frozen_epoch(&self) -> Option<Epoch> {
self.frozen.front().map(|data| data.epoch)
}
pub fn request_epoch_rewards_init(&self, epoch: Epoch, total_rewards: u64) {
*self
.epoch_rewards_init
.write()
.expect("Failed to acquire lock") = Some(EpochRewardsInitRequest {
epoch,
total_rewards,
});
}
pub fn take_epoch_rewards_init_request(&self) -> Option<EpochRewardsInitRequest> {
self.epoch_rewards_init
.write()
.expect("Failed to acquire lock")
.take()
}
pub fn is_epoch_rewards_init_pending(&self) -> bool {
self.epoch_rewards_init
.read()
.expect("Failed to acquire lock")
.is_some()
}
pub fn completed_frozen_epochs(&self) -> Vec<Epoch> {
let frozen_data = self.frozen.read();
let len = frozen_data.len();
if len < 2 {
return vec![];
}
frozen_data
.iter()
.take(len - 1) .map(|data| data.epoch)
.collect()
}
pub fn is_validator_referenced(
&self,
validator: &Pubkey,
validator_info: &ValidatorInfo,
last_freeze_timestamp: u64,
current_timestamp: u64,
) -> bool {
let all_stake_accounts = self.get_all_stake_accounts_from_pending();
all_stake_accounts.par_iter().any(|(_, stake_account)| {
if stake_account.data.validator.as_ref() != Some(validator) {
return false;
}
let Some(deactivation_timestamp) = stake_account.data.deactivation_requested else {
return true;
};
if deactivation_timestamp >= last_freeze_timestamp {
return true;
}
let unbonding_end = validator_info.end_of_unbonding(deactivation_timestamp);
unbonding_end >= current_timestamp
})
}
pub fn has_locked_stakers(
&self,
validator: &Pubkey,
lockup_period: u64,
current_timestamp: u64,
) -> bool {
let self_bond_pubkey = derive_self_bond_address(validator);
let all_stake_accounts = self.get_all_stake_accounts_from_pending();
all_stake_accounts
.par_iter()
.any(|(pubkey, stake_account)| {
if *pubkey == self_bond_pubkey {
return false;
}
if stake_account.data.validator.as_ref() != Some(validator) {
return false;
}
let Some(activation_requested) = stake_account.data.activation_requested else {
return false;
};
let lockup_end = activation_requested.saturating_add(lockup_period);
lockup_end > current_timestamp
})
}
pub fn is_validator_referenced_excluding_self_bond(
&self,
validator_pubkey: &Pubkey,
validator_info: &ValidatorInfo,
last_freeze_timestamp: u64,
current_timestamp: u64,
) -> bool {
let self_bond_pubkey = derive_self_bond_address(validator_pubkey);
let all_stake_accounts = self.get_all_stake_accounts_from_pending();
all_stake_accounts
.par_iter()
.any(|(pubkey, stake_account)| {
if *pubkey == self_bond_pubkey {
return false;
}
if stake_account.data.validator.as_ref() != Some(validator_pubkey) {
return false;
}
let Some(deactivation_timestamp) = stake_account.data.deactivation_requested else {
return true;
};
if deactivation_timestamp >= last_freeze_timestamp {
return true;
}
let unbonding_end = validator_info.end_of_unbonding(deactivation_timestamp);
unbonding_end >= current_timestamp
})
}
pub fn insert_stake_account(&self, pubkey: Pubkey, account: StakeAccount) {
self.pending.insert_stake_account(pubkey, account);
}
pub fn insert_validator_account(&self, pubkey: Pubkey, account: ValidatorAccount) {
self.pending.insert_validator_account(pubkey, account);
}
}
pub struct StakesView(StakesHandle);
impl StakesView {
pub fn new(handle: StakesHandle) -> Self {
Self(handle)
}
pub fn get_stake_account_from_pending(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
self.0.get_stake_account_from_pending(pubkey)
}
pub fn get_validator_account_from_pending(&self, pubkey: &Pubkey) -> Option<ValidatorAccount> {
self.0.get_validator_account_from_pending(pubkey)
}
pub fn get_all_validator_accounts_from_pending(&self) -> Vec<(Pubkey, ValidatorAccount)> {
self.0.get_all_validator_accounts_from_pending()
}
pub fn get_stake_account_from_last_frozen(&self, pubkey: &Pubkey) -> Option<StakeAccount> {
self.0.get_stake_account_from_last_frozen(pubkey)
}
pub fn get_validator_account_from_last_frozen(
&self,
pubkey: &Pubkey,
) -> Option<ValidatorAccount> {
self.0.get_validator_account_from_last_frozen(pubkey)
}
pub fn get_all_validator_accounts_from_last_frozen(&self) -> Vec<(Pubkey, ValidatorAccount)> {
self.0.get_all_validator_accounts_from_last_frozen()
}
pub fn last_frozen_timestamp(&self) -> Option<u64> {
self.0.last_frozen_timestamp()
}
pub fn last_frozen_epoch(&self) -> Option<Epoch> {
self.0.last_frozen_epoch()
}
pub fn pending_epoch(&self) -> Epoch {
self.0.pending_epoch()
}
pub fn pending_timestamp(&self) -> u64 {
self.0.pending_timestamp()
}
}
#[cfg(test)]
impl StakesHandle {
pub fn raw_baseline(&self) -> &StakeCache {
&self.baseline
}
pub fn raw_pending(&self) -> &StakeCache {
&self.pending
}
pub fn raw_frozen(&self) -> &StakeHistory {
&self.frozen
}
}
#[cfg(test)]
mod tests {
use rialo_stake_manager_interface::instruction::StakeInfo;
use rialo_validator_registry_interface::instruction::ValidatorInfo;
use super::*;
fn create_test_stake_account(kelvins: u64, validator: Pubkey) -> StakeAccount {
StakeAccount {
kelvins,
data: StakeInfo {
activation_requested: Some(0),
deactivation_requested: None,
delegated_balance: kelvins,
validator: Some(validator),
admin_authority: Pubkey::new_unique(),
withdraw_authority: Pubkey::new_unique(),
reward_receiver: None,
},
}
}
fn create_test_validator_account(kelvins: u64, stake: u64) -> ValidatorAccount {
ValidatorAccount {
kelvins,
data: ValidatorInfo {
signing_key: Pubkey::new_unique(),
withdrawal_key: Pubkey::new_unique(),
registration_time: 0,
stake,
address: vec![],
state_sync_address: vec![],
hostname: String::new(),
authority_key: vec![0u8; 96],
protocol_key: Pubkey::new_unique(),
network_key: Pubkey::new_unique(),
last_update: 0,
unbonding_periods: std::collections::BTreeMap::from([(0, 0)]),
lockup_period: 0,
commission_rate: 500,
new_commission_rate: None,
earliest_shutdown: None,
},
}
}
#[test]
fn test_layered_lookup_stake_account_from_pending() {
let pubkey = Pubkey::new_unique();
let validator = Pubkey::new_unique();
let handle = StakesHandle::default();
let pending_account = create_test_stake_account(1000, validator);
handle.insert_stake_account(pubkey, pending_account.clone());
let found = handle.get_stake_account_from_pending(&pubkey);
assert!(found.is_some());
assert_eq!(found.unwrap().kelvins, 1000);
}
#[test]
fn test_layered_lookup_stake_account_from_frozen() {
let pubkey = Pubkey::new_unique();
let validator = Pubkey::new_unique();
let handle = StakesHandle::default();
let account = create_test_stake_account(2000, validator);
handle.insert_stake_account(pubkey, account);
handle.freeze_stakes();
let found = handle.get_stake_account_from_pending(&pubkey);
assert!(found.is_some());
assert_eq!(found.unwrap().kelvins, 2000);
assert!(handle.raw_pending().get_stake_account(&pubkey).is_none());
}
#[test]
fn test_layered_lookup_stake_account_from_baseline() {
let pubkey = Pubkey::new_unique();
let validator = Pubkey::new_unique();
let mut baseline_data = StakeCacheData::default();
baseline_data
.stake_accounts
.insert(pubkey, Some(create_test_stake_account(3000, validator)));
let baseline = StakeCache::with_data(baseline_data);
let handle = StakesHandle::new_shared(
baseline,
StakeCache::default(),
StakeHistory::default(),
Arc::new(|_| false),
);
let found = handle.get_stake_account_from_pending(&pubkey);
assert!(found.is_some());
assert_eq!(found.unwrap().kelvins, 3000);
}
#[test]
fn test_layered_lookup_priority_pending_over_frozen() {
let pubkey = Pubkey::new_unique();
let validator = Pubkey::new_unique();
let handle = StakesHandle::default();
handle.insert_stake_account(pubkey, create_test_stake_account(1000, validator));
handle.freeze_stakes();
handle.insert_stake_account(pubkey, create_test_stake_account(2000, validator));
let found = handle.get_stake_account_from_pending(&pubkey);
assert!(found.is_some());
assert_eq!(found.unwrap().kelvins, 2000);
let found_frozen = handle.get_stake_account_from_last_frozen(&pubkey);
assert!(found_frozen.is_some());
assert_eq!(found_frozen.unwrap().kelvins, 1000);
}
#[test]
fn test_layered_lookup_priority_frozen_over_baseline() {
let pubkey = Pubkey::new_unique();
let validator = Pubkey::new_unique();
let mut baseline_data = StakeCacheData::default();
baseline_data
.stake_accounts
.insert(pubkey, Some(create_test_stake_account(1000, validator)));
let baseline = StakeCache::with_data(baseline_data);
let handle = StakesHandle::new_shared(
baseline,
StakeCache::default(),
StakeHistory::default(),
Arc::new(|_| false),
);
handle.insert_stake_account(pubkey, create_test_stake_account(2000, validator));
handle.freeze_stakes();
let found = handle.get_stake_account_from_pending(&pubkey);
assert!(found.is_some());
assert_eq!(found.unwrap().kelvins, 2000);
}
#[test]
fn test_layered_lookup_multiple_frozen_epochs() {
let pubkey = Pubkey::new_unique();
let validator = Pubkey::new_unique();
let handle = StakesHandle::default();
handle.insert_stake_account(pubkey, create_test_stake_account(1000, validator));
handle.freeze_stakes();
handle.insert_stake_account(pubkey, create_test_stake_account(2000, validator));
handle.freeze_stakes();
handle.insert_stake_account(pubkey, create_test_stake_account(3000, validator));
handle.freeze_stakes();
let found = handle.get_stake_account_from_last_frozen(&pubkey);
assert!(found.is_some());
assert_eq!(found.unwrap().kelvins, 3000);
assert_eq!(handle.frozen_len(), 3);
}
#[test]
fn test_layered_lookup_validator_account() {
let pubkey = Pubkey::new_unique();
let mut baseline_data = StakeCacheData::default();
baseline_data
.validator_accounts
.insert(pubkey, Some(create_test_validator_account(1000, 500)));
let baseline = StakeCache::with_data(baseline_data);
let handle = StakesHandle::new_shared(
baseline,
StakeCache::default(),
StakeHistory::default(),
Arc::new(|_| false),
);
let found = handle.get_validator_account_from_pending(&pubkey);
assert!(found.is_some());
assert_eq!(found.unwrap().kelvins, 1000);
handle.insert_validator_account(pubkey, create_test_validator_account(2000, 600));
let found = handle.get_validator_account_from_pending(&pubkey);
assert!(found.is_some());
assert_eq!(found.unwrap().kelvins, 2000);
}
#[test]
fn test_tombstone_in_pending_hides_frozen() {
let pubkey = Pubkey::new_unique();
let validator = Pubkey::new_unique();
let handle = StakesHandle::default();
handle.insert_stake_account(pubkey, create_test_stake_account(1000, validator));
handle.freeze_stakes();
handle.raw_pending().tombstone_stake_account(pubkey);
let found = handle.get_stake_account_from_pending(&pubkey);
assert!(
found.is_none(),
"Tombstone in pending should hide frozen value"
);
let found_frozen = handle.get_stake_account_from_last_frozen(&pubkey);
assert!(found_frozen.is_some());
assert_eq!(found_frozen.unwrap().kelvins, 1000);
}
#[test]
fn test_tombstone_in_frozen_hides_baseline() {
let pubkey = Pubkey::new_unique();
let validator = Pubkey::new_unique();
let mut baseline_data = StakeCacheData::default();
baseline_data
.stake_accounts
.insert(pubkey, Some(create_test_stake_account(1000, validator)));
let baseline = StakeCache::with_data(baseline_data);
let handle = StakesHandle::new_shared(
baseline,
StakeCache::default(),
StakeHistory::default(),
Arc::new(|_| false),
);
handle.raw_pending().tombstone_stake_account(pubkey);
handle.freeze_stakes();
let found = handle.get_stake_account_from_last_frozen(&pubkey);
assert!(
found.is_none(),
"Tombstone in frozen should hide baseline value"
);
let found_first = handle.get_stake_account_from_first_frozen(&pubkey);
assert!(found_first.is_none());
}
#[test]
fn test_tombstone_validator_account() {
let pubkey = Pubkey::new_unique();
let mut baseline_data = StakeCacheData::default();
baseline_data
.validator_accounts
.insert(pubkey, Some(create_test_validator_account(1000, 500)));
let baseline = StakeCache::with_data(baseline_data);
let handle = StakesHandle::new_shared(
baseline,
StakeCache::default(),
StakeHistory::default(),
Arc::new(|_| false),
);
assert!(handle.get_validator_account_from_pending(&pubkey).is_some());
handle.raw_pending().tombstone_validator_account(pubkey);
let found = handle.get_validator_account_from_pending(&pubkey);
assert!(found.is_none(), "Tombstone should hide baseline validator");
}
#[test]
fn test_get_all_validators_excludes_tombstones() {
let pubkey1 = Pubkey::new_unique();
let pubkey2 = Pubkey::new_unique();
let mut baseline_data = StakeCacheData::default();
baseline_data
.validator_accounts
.insert(pubkey1, Some(create_test_validator_account(1000, 100)));
baseline_data
.validator_accounts
.insert(pubkey2, Some(create_test_validator_account(2000, 200)));
let baseline = StakeCache::with_data(baseline_data);
let handle = StakesHandle::new_shared(
baseline,
StakeCache::default(),
StakeHistory::default(),
Arc::new(|_| false),
);
let all = handle.get_all_validator_accounts_from_pending();
assert_eq!(all.len(), 2);
handle.raw_pending().tombstone_validator_account(pubkey1);
let all = handle.get_all_validator_accounts_from_pending();
assert_eq!(all.len(), 1);
assert_eq!(all[0].0, pubkey2);
}
#[test]
fn test_tombstone_then_readd() {
let pubkey = Pubkey::new_unique();
let validator = Pubkey::new_unique();
let mut baseline_data = StakeCacheData::default();
baseline_data
.stake_accounts
.insert(pubkey, Some(create_test_stake_account(1000, validator)));
let baseline = StakeCache::with_data(baseline_data);
let handle = StakesHandle::new_shared(
baseline,
StakeCache::default(),
StakeHistory::default(),
Arc::new(|_| false),
);
handle.raw_pending().tombstone_stake_account(pubkey);
handle.freeze_stakes();
let found = handle.get_stake_account_from_last_frozen(&pubkey);
assert!(found.is_none());
handle.insert_stake_account(pubkey, create_test_stake_account(5000, validator));
handle.freeze_stakes();
let found = handle.get_stake_account_from_last_frozen(&pubkey);
assert!(found.is_some());
assert_eq!(found.unwrap().kelvins, 5000);
}
#[test]
fn test_empty_pending_freeze() {
let handle = StakesHandle::default();
handle.freeze_stakes();
assert_eq!(handle.frozen_len(), 1);
let pubkey = Pubkey::new_unique();
assert!(handle.get_stake_account_from_pending(&pubkey).is_none());
}
#[test]
fn test_empty_frozen_epochs() {
let pubkey = Pubkey::new_unique();
let validator = Pubkey::new_unique();
let mut baseline_data = StakeCacheData::default();
baseline_data
.stake_accounts
.insert(pubkey, Some(create_test_stake_account(1000, validator)));
let baseline = StakeCache::with_data(baseline_data);
let handle = StakesHandle::new_shared(
baseline,
StakeCache::default(),
StakeHistory::default(),
Arc::new(|_| false),
);
handle.freeze_stakes();
handle.freeze_stakes();
handle.freeze_stakes();
let found = handle.get_stake_account_from_pending(&pubkey);
assert!(found.is_some());
assert_eq!(found.unwrap().kelvins, 1000);
}
#[test]
fn test_no_frozen_epochs_falls_through_to_baseline() {
let pubkey = Pubkey::new_unique();
let validator = Pubkey::new_unique();
let mut baseline_data = StakeCacheData::default();
baseline_data
.stake_accounts
.insert(pubkey, Some(create_test_stake_account(1000, validator)));
let baseline = StakeCache::with_data(baseline_data);
let handle = StakesHandle::new_shared(
baseline,
StakeCache::default(),
StakeHistory::default(),
Arc::new(|_| false),
);
let found = handle.get_stake_account_from_last_frozen(&pubkey);
assert!(found.is_some());
assert_eq!(found.unwrap().kelvins, 1000);
}
#[test]
fn test_get_all_stake_accounts_from_frozen_epoch() {
let validator = Pubkey::new_unique();
let baseline_stake = Pubkey::new_unique();
let mut baseline_data = StakeCacheData::default();
baseline_data.stake_accounts.insert(
baseline_stake,
Some(create_test_stake_account(1000, validator)),
);
let baseline = StakeCache::with_data(baseline_data);
let handle = StakesHandle::new_shared(
baseline,
StakeCache::default(),
StakeHistory::default(),
Arc::new(|_| false),
);
let stake_epoch5 = Pubkey::new_unique();
handle.set_pending_epoch(5);
handle.insert_stake_account(stake_epoch5, create_test_stake_account(2000, validator));
handle.freeze_stakes();
let stake_epoch6 = Pubkey::new_unique();
handle.insert_stake_account(stake_epoch6, create_test_stake_account(3000, validator));
handle.freeze_stakes();
let stake_epoch7 = Pubkey::new_unique();
handle.insert_stake_account(stake_epoch7, create_test_stake_account(4000, validator));
handle.freeze_stakes();
let accounts_epoch5 = handle.get_all_stake_accounts_from_frozen_epoch(5);
assert_eq!(accounts_epoch5.len(), 2); assert!(accounts_epoch5.iter().any(|(k, _)| *k == baseline_stake));
assert!(accounts_epoch5.iter().any(|(k, _)| *k == stake_epoch5));
assert!(!accounts_epoch5.iter().any(|(k, _)| *k == stake_epoch6));
let accounts_epoch6 = handle.get_all_stake_accounts_from_frozen_epoch(6);
assert_eq!(accounts_epoch6.len(), 3);
assert!(accounts_epoch6.iter().any(|(k, _)| *k == stake_epoch6));
assert!(!accounts_epoch6.iter().any(|(k, _)| *k == stake_epoch7));
let accounts_epoch7 = handle.get_all_stake_accounts_from_frozen_epoch(7);
assert_eq!(accounts_epoch7.len(), 4);
assert!(accounts_epoch7.iter().any(|(k, _)| *k == stake_epoch7));
}
#[test]
fn test_get_all_validators_with_no_validators() {
let handle = StakesHandle::default();
let all = handle.get_all_validator_accounts_from_pending();
assert!(all.is_empty());
handle.freeze_stakes();
let all = handle.get_all_validator_accounts_from_last_frozen();
assert!(all.is_empty());
}
}