From d3d3d873547eeb7e445f536c8ff37fd617d175f6 Mon Sep 17 00:00:00 2001 From: Coach Chuck <169060940+coachchucksol@users.noreply.github.com> Date: Thu, 13 Jun 2024 09:40:11 -0600 Subject: [PATCH] Christian/epoch update (#47) Starting to add the solution for keeping track of removed validators --------- Co-authored-by: Christian --- programs/steward/src/errors.rs | 10 ++ .../auto_add_validator_to_pool.rs | 22 ++- .../auto_remove_validator_from_pool.rs | 10 +- .../src/instructions/compute_delegations.rs | 5 + .../instructions/compute_instant_unstake.rs | 6 + .../steward/src/instructions/compute_score.rs | 16 ++- .../src/instructions/epoch_maintenance.rs | 95 +++++++++++++ programs/steward/src/instructions/idle.rs | 5 + programs/steward/src/instructions/mod.rs | 2 + .../steward/src/instructions/realloc_state.rs | 7 +- .../steward/src/instructions/rebalance.rs | 5 + .../src/instructions/reset_steward_state.rs | 7 +- .../src/instructions/spl_passthrough.rs | 36 ++++- programs/steward/src/lib.rs | 8 ++ programs/steward/src/state/steward_state.rs | 134 +++++++++++++++--- programs/steward/src/utils.rs | 39 ++++- tests/src/steward_fixtures.rs | 32 ++++- tests/tests/steward/test_integration.rs | 51 ++++++- tests/tests/steward/test_spl_passthrough.rs | 33 ++++- tests/tests/steward/test_state_methods.rs | 69 +++++++-- tests/tests/steward/test_steward.rs | 1 + 21 files changed, 532 insertions(+), 61 deletions(-) create mode 100644 programs/steward/src/instructions/epoch_maintenance.rs diff --git a/programs/steward/src/errors.rs b/programs/steward/src/errors.rs index 7380a4ef..6554baf9 100644 --- a/programs/steward/src/errors.rs +++ b/programs/steward/src/errors.rs @@ -50,4 +50,14 @@ pub enum StewardError { MaxValidatorsReached, #[msg("Validator history account does not match vote account")] ValidatorHistoryMismatch, + #[msg("Epoch Maintenance must be called before continuing")] + EpochMaintenanceNotComplete, + #[msg("The stake pool must be updated before continuing")] + StakePoolNotUpdated, + #[msg("Validator not marked for removal")] + ValidatorNotMarkedForRemoval, + #[msg("Validators have not been removed")] + ValidatorsHaveNotBeenRemoved, + #[msg("Validator List count does not match state machine")] + ListStateMismatch, } diff --git a/programs/steward/src/instructions/auto_add_validator_to_pool.rs b/programs/steward/src/instructions/auto_add_validator_to_pool.rs index 8ecb5b47..123e620f 100644 --- a/programs/steward/src/instructions/auto_add_validator_to_pool.rs +++ b/programs/steward/src/instructions/auto_add_validator_to_pool.rs @@ -1,6 +1,6 @@ use crate::constants::{MAX_VALIDATORS, STAKE_POOL_WITHDRAW_SEED}; use crate::errors::StewardError; -use crate::state::{Config, Staker}; +use crate::state::{Config, Staker, StewardStateAccount}; use crate::utils::{deserialize_stake_pool, get_stake_pool_address}; use anchor_lang::prelude::*; use anchor_lang::solana_program::{program::invoke_signed, stake, sysvar, vote}; @@ -10,6 +10,15 @@ use validator_history::state::ValidatorHistory; #[derive(Accounts)] pub struct AutoAddValidator<'info> { + pub config: AccountLoader<'info, Config>, + + #[account( + mut, + seeds = [StewardStateAccount::SEED, config.key().as_ref()], + bump + )] + pub steward_state: AccountLoader<'info, StewardStateAccount>, + // Only adding validators where this exists #[account( seeds = [ValidatorHistory::SEED, vote_account.key().as_ref()], @@ -18,8 +27,6 @@ pub struct AutoAddValidator<'info> { )] pub validator_history_account: AccountLoader<'info, ValidatorHistory>, - pub config: AccountLoader<'info, Config>, - /// CHECK: CPI address #[account( address = spl_stake_pool::ID @@ -102,10 +109,17 @@ all the validators we want to be eligible for delegation, as well as to accept s Performs some eligibility checks in order to not fill up the validator list with offline or malicious validators. */ pub fn handler(ctx: Context) -> Result<()> { + let mut state_account = ctx.accounts.steward_state.load_mut()?; let config = ctx.accounts.config.load()?; let validator_history = ctx.accounts.validator_history_account.load()?; let epoch = Clock::get()?.epoch; + // Should not be able to add a validator if update is not complete + require!( + epoch == state_account.state.current_epoch, + StewardError::EpochMaintenanceNotComplete + ); + { let validator_list_data = &mut ctx.accounts.validator_list.try_borrow_mut_data()?; let (_, validator_list) = ValidatorListHeader::deserialize_vec(validator_list_data)?; @@ -140,6 +154,8 @@ pub fn handler(ctx: Context) -> Result<()> { return Err(StewardError::ValidatorBelowLivenessMinimum.into()); } + state_account.state.increment_validator_to_add()?; + invoke_signed( &spl_stake_pool::instruction::add_validator_to_pool( &ctx.accounts.stake_pool_program.key(), diff --git a/programs/steward/src/instructions/auto_remove_validator_from_pool.rs b/programs/steward/src/instructions/auto_remove_validator_from_pool.rs index 15b8b67d..fa4bd358 100644 --- a/programs/steward/src/instructions/auto_remove_validator_from_pool.rs +++ b/programs/steward/src/instructions/auto_remove_validator_from_pool.rs @@ -143,6 +143,12 @@ pub fn handler(ctx: Context, validator_list_index: usize) - StewardError::ValidatorNotInList ); + // Should not be able to remove a validator if update is not complete + require!( + epoch == state_account.state.current_epoch, + StewardError::EpochMaintenanceNotComplete + ); + // Checks state for deactivate delinquent status, preventing pool from merging stake with activating let stake_account_deactivated = { let stake_account_data = &mut ctx.accounts.stake_account.data.borrow_mut(); @@ -164,7 +170,9 @@ pub fn handler(ctx: Context, validator_list_index: usize) - StewardError::ValidatorNotRemovable ); - state_account.state.remove_validator(validator_list_index)?; + state_account + .state + .mark_validator_for_removal(validator_list_index)?; invoke_signed( &spl_stake_pool::instruction::remove_validator_from_pool( diff --git a/programs/steward/src/instructions/compute_delegations.rs b/programs/steward/src/instructions/compute_delegations.rs index 679144df..33f29b78 100644 --- a/programs/steward/src/instructions/compute_delegations.rs +++ b/programs/steward/src/instructions/compute_delegations.rs @@ -28,6 +28,11 @@ pub fn handler(ctx: Context) -> Result<()> { let clock = Clock::get()?; let epoch_schedule = EpochSchedule::get()?; + require!( + clock.epoch == state_account.state.current_epoch, + StewardError::EpochMaintenanceNotComplete + ); + if config.is_paused() { return Err(StewardError::StateMachinePaused.into()); } diff --git a/programs/steward/src/instructions/compute_instant_unstake.rs b/programs/steward/src/instructions/compute_instant_unstake.rs index 9bcd64f4..75eb282f 100644 --- a/programs/steward/src/instructions/compute_instant_unstake.rs +++ b/programs/steward/src/instructions/compute_instant_unstake.rs @@ -49,6 +49,11 @@ pub fn handler(ctx: Context, validator_list_index: usize) StewardError::ValidatorNotInList ); + require!( + clock.epoch == state_account.state.current_epoch, + StewardError::EpochMaintenanceNotComplete + ); + if config.is_paused() { return Err(StewardError::StateMachinePaused.into()); } @@ -61,6 +66,7 @@ pub fn handler(ctx: Context, validator_list_index: usize) &cluster, &config, )?; + maybe_transition_and_emit( &mut state_account.state, &clock, diff --git a/programs/steward/src/instructions/compute_score.rs b/programs/steward/src/instructions/compute_score.rs index 81125003..c93f618a 100644 --- a/programs/steward/src/instructions/compute_score.rs +++ b/programs/steward/src/instructions/compute_score.rs @@ -1,8 +1,9 @@ use anchor_lang::prelude::*; -use spl_stake_pool::state::ValidatorListHeader; use crate::{ - errors::StewardError, maybe_transition_and_emit, utils::get_validator_stake_info_at_index, + errors::StewardError, + maybe_transition_and_emit, + utils::{get_validator_list_length, get_validator_stake_info_at_index}, Config, StewardStateAccount, StewardStateEnum, }; use validator_history::{ClusterHistory, ValidatorHistory}; @@ -51,11 +52,12 @@ pub fn handler(ctx: Context, validator_list_index: usize) -> Resul StewardError::ValidatorNotInList ); - let num_pool_validators = { - let mut validator_list_data = validator_list.try_borrow_mut_data()?; - let (_, validator_list) = ValidatorListHeader::deserialize_vec(&mut validator_list_data)?; - validator_list.len() as u64 - }; + let num_pool_validators = get_validator_list_length(validator_list)?; + + require!( + clock.epoch == state_account.state.current_epoch, + StewardError::EpochMaintenanceNotComplete + ); if config.is_paused() { return Err(StewardError::StateMachinePaused.into()); diff --git a/programs/steward/src/instructions/epoch_maintenance.rs b/programs/steward/src/instructions/epoch_maintenance.rs new file mode 100644 index 00000000..480a3157 --- /dev/null +++ b/programs/steward/src/instructions/epoch_maintenance.rs @@ -0,0 +1,95 @@ +use crate::{ + errors::StewardError, + utils::{ + check_validator_list_has_stake_status, get_stake_pool, get_validator_list_length, StakePool, + }, + Config, StewardStateAccount, +}; +use anchor_lang::prelude::*; +use spl_stake_pool::state::StakeStatus; + +#[derive(Accounts)] +pub struct EpochMaintenance<'info> { + pub config: AccountLoader<'info, Config>, + + #[account( + mut, + seeds = [StewardStateAccount::SEED, config.key().as_ref()], + bump + )] + pub state_account: AccountLoader<'info, StewardStateAccount>, + + #[account(mut, address = stake_pool.validator_list)] + pub validator_list: AccountInfo<'info>, + + #[account( + address = get_stake_pool(&config)? + )] + pub stake_pool: Account<'info, StakePool>, +} + +/// Runs maintenance tasks at the start of each epoch, needs to be run multiple times +/// Routines: +/// - Remove delinquent validators +pub fn handler( + ctx: Context, + validator_index_to_remove: Option, +) -> Result<()> { + let stake_pool = &ctx.accounts.stake_pool; + let mut state_account = ctx.accounts.state_account.load_mut()?; + + let clock = Clock::get()?; + + require!( + clock.epoch == stake_pool.last_update_epoch, + StewardError::StakePoolNotUpdated + ); + + if (!state_account.state.checked_validators_removed_from_list).into() { + // Ensure there are no validators in the list that have not been removed, that should be + require!( + !check_validator_list_has_stake_status( + &ctx.accounts.validator_list, + StakeStatus::ReadyForRemoval + )?, + StewardError::ValidatorsHaveNotBeenRemoved + ); + state_account.state.checked_validators_removed_from_list = true.into(); + } + + { + // Routine - Remove marked validators + // We still want these checks to run even if we don't specify a validator to remove + let validators_in_list = get_validator_list_length(&ctx.accounts.validator_list)?; + let validators_to_remove = state_account.state.validators_to_remove.count(); + + // Ensure we have a 1-1 mapping between the number of validators in the list and the number of validators in the state + require!( + state_account.state.num_pool_validators + state_account.state.validators_added as usize + - validators_to_remove + == validators_in_list, + StewardError::ListStateMismatch + ); + + if let Some(validator_index_to_remove) = validator_index_to_remove { + state_account + .state + .remove_validator(validator_index_to_remove)?; + } + } + + { + // Routine - Update state + let okay_to_update = state_account.state.validators_to_remove.is_empty() + && state_account + .state + .checked_validators_removed_from_list + .into(); + if okay_to_update { + state_account.state.current_epoch = clock.epoch; + state_account.state.checked_validators_removed_from_list = false.into(); + } + } + + Ok(()) +} diff --git a/programs/steward/src/instructions/idle.rs b/programs/steward/src/instructions/idle.rs index 4d647612..47d594aa 100644 --- a/programs/steward/src/instructions/idle.rs +++ b/programs/steward/src/instructions/idle.rs @@ -33,6 +33,11 @@ pub fn handler(ctx: Context) -> Result<()> { StewardError::InvalidState ); + require!( + clock.epoch == state_account.state.current_epoch, + StewardError::EpochMaintenanceNotComplete + ); + if config.is_paused() { return Err(StewardError::StateMachinePaused.into()); } diff --git a/programs/steward/src/instructions/mod.rs b/programs/steward/src/instructions/mod.rs index ec849977..ad9c7c2c 100644 --- a/programs/steward/src/instructions/mod.rs +++ b/programs/steward/src/instructions/mod.rs @@ -6,6 +6,7 @@ pub mod close_steward_accounts; pub mod compute_delegations; pub mod compute_instant_unstake; pub mod compute_score; +pub mod epoch_maintenance; pub mod idle; pub mod initialize_config; pub mod initialize_state; @@ -26,6 +27,7 @@ pub use close_steward_accounts::*; pub use compute_delegations::*; pub use compute_instant_unstake::*; pub use compute_score::*; +pub use epoch_maintenance::*; pub use idle::*; pub use initialize_config::*; pub use initialize_state::*; diff --git a/programs/steward/src/instructions/realloc_state.rs b/programs/steward/src/instructions/realloc_state.rs index fadaf40c..1f543815 100644 --- a/programs/steward/src/instructions/realloc_state.rs +++ b/programs/steward/src/instructions/realloc_state.rs @@ -3,7 +3,7 @@ use crate::{ constants::{MAX_ALLOC_BYTES, MAX_VALIDATORS, SORTED_INDEX_DEFAULT}, errors::StewardError, state::{Config, StewardStateAccount}, - Delegation, StewardStateEnum, + Delegation, StewardStateEnum, STATE_PADDING_0_SIZE, }; use anchor_lang::prelude::*; use spl_stake_pool::state::ValidatorListHeader; @@ -89,7 +89,10 @@ pub fn handler(ctx: Context) -> Result<()> { state_account.state.rebalance_completed = false.into(); state_account.state.instant_unstake = BitMask::default(); state_account.state.start_computing_scores_slot = clock.slot; - state_account.state._padding0 = [0; 6 + MAX_VALIDATORS * 8]; + state_account.state.validators_to_remove = BitMask::default(); + state_account.state.validators_added = 0; + state_account.state.checked_validators_removed_from_list = false.into(); + state_account.state._padding0 = [0; STATE_PADDING_0_SIZE]; } Ok(()) diff --git a/programs/steward/src/instructions/rebalance.rs b/programs/steward/src/instructions/rebalance.rs index 92ef5a46..4f6e0ff6 100644 --- a/programs/steward/src/instructions/rebalance.rs +++ b/programs/steward/src/instructions/rebalance.rs @@ -161,6 +161,11 @@ pub fn handler(ctx: Context, validator_list_index: usize) -> Result<( ); let transient_seed = u64::from(validator_stake_info.transient_seed_suffix); + require!( + clock.epoch == state_account.state.current_epoch, + StewardError::EpochMaintenanceNotComplete + ); + if config.is_paused() { return Err(StewardError::StateMachinePaused.into()); } diff --git a/programs/steward/src/instructions/reset_steward_state.rs b/programs/steward/src/instructions/reset_steward_state.rs index 6b9bdfa2..0140f7d6 100644 --- a/programs/steward/src/instructions/reset_steward_state.rs +++ b/programs/steward/src/instructions/reset_steward_state.rs @@ -3,7 +3,7 @@ use crate::{ errors::StewardError, state::{Config, StewardStateAccount}, utils::{get_config_authority, get_stake_pool, StakePool}, - BitMask, Delegation, StewardStateEnum, + BitMask, Delegation, StewardStateEnum, STATE_PADDING_0_SIZE, }; use anchor_lang::prelude::*; use spl_stake_pool::state::ValidatorListHeader; @@ -59,6 +59,9 @@ pub fn handler(ctx: Context) -> Result<()> { state_account.state.rebalance_completed = false.into(); state_account.state.instant_unstake = BitMask::default(); state_account.state.start_computing_scores_slot = clock.slot; - state_account.state._padding0 = [0; 6 + MAX_VALIDATORS * 8]; + state_account.state.validators_to_remove = BitMask::default(); + state_account.state.validators_added = 0; + state_account.state.checked_validators_removed_from_list = false.into(); + state_account.state._padding0 = [0; STATE_PADDING_0_SIZE]; Ok(()) } diff --git a/programs/steward/src/instructions/spl_passthrough.rs b/programs/steward/src/instructions/spl_passthrough.rs index ed4224a8..fb6771d0 100644 --- a/programs/steward/src/instructions/spl_passthrough.rs +++ b/programs/steward/src/instructions/spl_passthrough.rs @@ -23,6 +23,13 @@ use validator_history::ValidatorHistory; #[derive(Accounts)] pub struct AddValidatorToPool<'info> { pub config: AccountLoader<'info, Config>, + + #[account( + mut, + seeds = [StewardStateAccount::SEED, config.key().as_ref()], + bump + )] + pub steward_state: AccountLoader<'info, StewardStateAccount>, /// CHECK: CPI program #[account( address = spl_stake_pool::ID @@ -73,6 +80,15 @@ pub fn add_validator_to_pool_handler( ctx: Context, validator_seed: Option, ) -> Result<()> { + let mut state_account = ctx.accounts.steward_state.load_mut()?; + let epoch = Clock::get()?.epoch; + + // Should not be able to add a validator if update is not complete + require!( + epoch == state_account.state.current_epoch, + StewardError::EpochMaintenanceNotComplete + ); + { let validator_list_data = &mut ctx.accounts.validator_list.try_borrow_mut_data()?; let (_, validator_list) = ValidatorListHeader::deserialize_vec(validator_list_data)?; @@ -81,6 +97,9 @@ pub fn add_validator_to_pool_handler( return Err(StewardError::MaxValidatorsReached.into()); } } + + state_account.state.increment_validator_to_add()?; + invoke_signed( &spl_stake_pool::instruction::add_validator_to_pool( &ctx.accounts.stake_pool_program.key(), @@ -168,6 +187,13 @@ pub fn remove_validator_from_pool_handler( validator_list_index: usize, ) -> Result<()> { let mut state_account = ctx.accounts.steward_state.load_mut()?; + let epoch = Clock::get()?.epoch; + + // Should not be able to remove a validator if update is not complete + require!( + epoch == state_account.state.current_epoch, + StewardError::EpochMaintenanceNotComplete + ); if validator_list_index < state_account.state.num_pool_validators as usize { let validator_list_stake_info = get_validator_stake_info_at_index( @@ -185,10 +211,12 @@ pub fn remove_validator_from_pool_handler( if validator_list_stake_account != ctx.accounts.stake_account.key() { return Err(StewardError::ValidatorNotInList.into()); } - - state_account.state.remove_validator(validator_list_index)?; } + state_account + .state + .mark_validator_for_removal(validator_list_index)?; + invoke_signed( &spl_stake_pool::instruction::remove_validator_from_pool( &ctx.accounts.stake_pool_program.key(), @@ -307,7 +335,7 @@ pub struct IncreaseValidatorStake<'info> { /// CHECK: passing through, checks are done by spl-stake-pool pub withdraw_authority: AccountInfo<'info>, /// CHECK: passing through, checks are done by spl-stake-pool - #[account(mut)] + #[account(mut, address = stake_pool.validator_list)] pub validator_list: AccountInfo<'info>, /// CHECK: passing through, checks are done by spl-stake-pool #[account( @@ -436,7 +464,7 @@ pub struct DecreaseValidatorStake<'info> { /// CHECK: passing through, checks are done by spl-stake-pool pub withdraw_authority: AccountInfo<'info>, /// CHECK: passing through, checks are done by spl-stake-pool - #[account(mut)] + #[account(mut, address = stake_pool.validator_list)] pub validator_list: AccountInfo<'info>, /// CHECK: passing through, checks are done by spl-stake-pool #[account( diff --git a/programs/steward/src/lib.rs b/programs/steward/src/lib.rs index bed7bea1..4f9578b7 100644 --- a/programs/steward/src/lib.rs +++ b/programs/steward/src/lib.rs @@ -91,6 +91,14 @@ pub mod steward { instructions::auto_remove_validator_from_pool::handler(ctx, validator_list_index as usize) } + /// Housekeeping, run at the start of any new epoch before any other instructions + pub fn epoch_maintenance( + ctx: Context, + validator_index_to_remove: Option, + ) -> Result<()> { + instructions::epoch_maintenance::handler(ctx, validator_index_to_remove) + } + /// Computes score for a the validator at `validator_list_index` for the current cycle. pub fn compute_score(ctx: Context, validator_list_index: u64) -> Result<()> { instructions::compute_score::handler(ctx, validator_list_index as usize) diff --git a/programs/steward/src/state/steward_state.rs b/programs/steward/src/state/steward_state.rs index eab3f46a..1a575641 100644 --- a/programs/steward/src/state/steward_state.rs +++ b/programs/steward/src/state/steward_state.rs @@ -35,26 +35,28 @@ pub struct StateTransition { } pub fn maybe_transition_and_emit( - state_account: &mut StewardState, + steward_state: &mut StewardState, clock: &Clock, params: &Parameters, epoch_schedule: &EpochSchedule, ) -> Result<()> { - let initial_state = state_account.state_tag.to_string(); - state_account.transition(clock, params, epoch_schedule)?; - if initial_state != state_account.state_tag.to_string() { + let initial_state = steward_state.state_tag.to_string(); + steward_state.transition(clock, params, epoch_schedule)?; + + if initial_state != steward_state.state_tag.to_string() { emit!(StateTransition { epoch: clock.epoch, slot: clock.slot, previous_state: initial_state, - new_state: state_account.state_tag.to_string(), + new_state: steward_state.state_tag.to_string(), }); } Ok(()) } /// Tracks state of the stake pool. -/// Follow state transitions here: [TODO add link to github diagram] +/// Follow state transitions here: +/// https://github.com/jito-foundation/stakenet/blob/master/programs/steward/state-machine-diagram.png #[derive(BorshSerialize)] #[zero_copy] pub struct StewardState { @@ -87,6 +89,10 @@ pub struct StewardState { /// Tracks progress of states that require one instruction per validator pub progress: BitMask, + /// Marks a validator for removal after `remove_validator_from_pool` has been called on the stake pool + /// This is cleaned up in the next epoch + pub validators_to_remove: BitMask, + ////// Cycle metadata fields ////// /// Slot of the first ComputeScores instruction in the current cycle pub start_computing_scores_slot: u64, @@ -110,17 +116,25 @@ pub struct StewardState { /// Total lamports that have been due to stake deposits this cycle pub stake_deposit_unstake_total: u64, + /// Number of validators added to the pool in the current cycle + pub validators_added: u16, + /// Tracks whether delegation computation has been completed pub compute_delegations_completed: U8Bool, /// Tracks whether unstake and delegate steps have completed pub rebalance_completed: U8Bool, + /// So we only have to check the validator list once for `ReadyToRemove` + pub checked_validators_removed_from_list: U8Bool, + /// Future state and #[repr(C)] alignment - pub _padding0: [u8; 6 + MAX_VALIDATORS * 8], + pub _padding0: [u8; STATE_PADDING_0_SIZE], // TODO ADD MORE PADDING } +pub const STATE_PADDING_0_SIZE: usize = MAX_VALIDATORS * 8 + 3; + #[derive(Clone, Copy)] #[repr(u64)] pub enum StewardStateEnum { @@ -239,6 +253,7 @@ impl StewardState { let current_epoch = clock.epoch; let current_slot = clock.slot; let epoch_progress = epoch_progress(clock, epoch_schedule)?; + match self.state_tag { StewardStateEnum::ComputeScores => self.transition_compute_scores( current_epoch, @@ -308,7 +323,6 @@ impl StewardState { )?; } else if self.compute_delegations_completed.into() { self.state_tag = StewardStateEnum::Idle; - self.current_epoch = current_epoch; self.rebalance_completed = false.into(); } Ok(()) @@ -356,7 +370,6 @@ impl StewardState { )?; } else if current_epoch > self.current_epoch { self.state_tag = StewardStateEnum::Idle; - self.current_epoch = current_epoch; self.instant_unstake = BitMask::default(); self.progress = BitMask::default(); } else if self.progress.is_complete(self.num_pool_validators)? { @@ -382,12 +395,10 @@ impl StewardState { )?; } else if current_epoch > self.current_epoch { self.state_tag = StewardStateEnum::Idle; - self.current_epoch = current_epoch; self.progress = BitMask::default(); self.rebalance_completed = false.into(); } else if self.progress.is_complete(self.num_pool_validators)? { self.state_tag = StewardStateEnum::Idle; - self.current_epoch = current_epoch; self.rebalance_completed = true.into(); } Ok(()) @@ -404,7 +415,6 @@ impl StewardState { self.scores = [0; MAX_VALIDATORS]; self.yield_scores = [0; MAX_VALIDATORS]; self.progress = BitMask::default(); - self.current_epoch = current_epoch; self.next_cycle_epoch = current_epoch .checked_add(num_epochs_between_scoring) .ok_or(StewardError::ArithmeticError)?; @@ -421,6 +431,11 @@ impl StewardState { /// Update internal state when a validator is removed from the pool pub fn remove_validator(&mut self, index: usize) -> Result<()> { + require!( + self.validators_to_remove.get(index)?, + StewardError::ValidatorNotMarkedForRemoval + ); + self.num_pool_validators = self .num_pool_validators .checked_sub(1) @@ -437,6 +452,8 @@ impl StewardState { self.instant_unstake .set(i, self.instant_unstake.get(next_i)?)?; self.progress.set(i, self.progress.get(next_i)?)?; + self.validators_to_remove + .set(i, self.validators_to_remove.get(next_i)?)?; } // Update score indices @@ -477,18 +494,37 @@ impl StewardState { } // Clear values on empty last index - self.validator_lamport_balances[num_pool_validators] = 0; - self.scores[num_pool_validators] = 0; - self.yield_scores[num_pool_validators] = 0; - self.sorted_score_indices[num_pool_validators] = SORTED_INDEX_DEFAULT; - self.sorted_yield_score_indices[num_pool_validators] = SORTED_INDEX_DEFAULT; - self.delegations[num_pool_validators] = Delegation::default(); - self.instant_unstake.set(num_pool_validators, false)?; - self.progress.set(num_pool_validators, false)?; + self.validator_lamport_balances[self.num_pool_validators] = 0; + self.scores[self.num_pool_validators] = 0; + self.yield_scores[self.num_pool_validators] = 0; + self.sorted_score_indices[self.num_pool_validators] = SORTED_INDEX_DEFAULT; + self.sorted_yield_score_indices[self.num_pool_validators] = SORTED_INDEX_DEFAULT; + self.delegations[self.num_pool_validators] = Delegation::default(); + self.instant_unstake.set(self.num_pool_validators, false)?; + self.validators_to_remove + .set(self.num_pool_validators, false)?; + self.progress.set(self.num_pool_validators, false)?; Ok(()) } + /// Mark a validator for removal from the pool - this happens right after + /// `remove_validator_from_pool` has been called on the stake pool + /// This is cleaned up in the next epoch + pub fn mark_validator_for_removal(&mut self, index: usize) -> Result<()> { + self.validators_to_remove.set(index, true) + } + + /// Called when adding a validator to the pool so that we can ensure a 1-1 mapping between + /// the validator list and the steward state + pub fn increment_validator_to_add(&mut self) -> Result<()> { + self.validators_added = self + .validators_added + .checked_add(1) + .ok_or(StewardError::ArithmeticError)?; + Ok(()) + } + /// One instruction per validator. Can be done in any order. /// Computes score for a validator for the current epoch, stores score, and yield score component. /// Inserts this validator's index into sorted_score_indices and sorted_yield_score_indices, sorted by @@ -544,7 +580,42 @@ impl StewardState { config.parameters.num_epochs_between_scoring, )?; // Updates num_pool_validators at the start of the cycle so validator additions later won't be considered + + require!( + num_pool_validators + == self.num_pool_validators + self.validators_added as usize, + StewardError::ListStateMismatch + ); self.num_pool_validators = num_pool_validators; + self.validators_added = 0; + } + + // Skip scoring if already processed + if self.progress.get(index)? { + return Ok(()); + } + + // Skip scoring if marked for deletion + if self.validators_to_remove.get(index)? { + self.scores[index] = 0_u32; + self.yield_scores[index] = 0_u32; + + let num_scores_calculated = self.progress.count(); + insert_sorted_index( + &mut self.sorted_score_indices, + &self.scores, + index as u16, + self.scores[index], + num_scores_calculated, + )?; + insert_sorted_index( + &mut self.sorted_yield_score_indices, + &self.yield_scores, + index as u16, + self.yield_scores[index], + num_scores_calculated, + )?; + return Ok(()); } let score = validator_score(validator, index, cluster, config, current_epoch as u16)?; @@ -646,6 +717,17 @@ impl StewardState { return Err(StewardError::InstantUnstakeNotReady.into()); } + // Skip if already processed + if self.progress.get(index)? { + return Ok(()); + } + + // Skip if marked for deletion + if self.validators_to_remove.get(index)? { + self.progress.set(index, true)?; + return Ok(()); + } + let first_slot = epoch_schedule.get_first_slot_in_epoch(clock.epoch); // Epoch credits and cluster history must be updated in the current epoch and after the midpoint of the epoch @@ -715,6 +797,18 @@ impl StewardState { self.state_tag.to_string(), )); } + + // Skip if already processed + if self.progress.get(index)? { + return Ok(RebalanceType::None); + } + + // Skip if marked for deletion + if self.validators_to_remove.get(index)? { + self.progress.set(index, true)?; + return Ok(RebalanceType::None); + } + let base_lamport_balance = minimum_delegation .checked_add(stake_rent) .ok_or(StewardError::ArithmeticError)?; diff --git a/programs/steward/src/utils.rs b/programs/steward/src/utils.rs index 96f281fd..78d52735 100644 --- a/programs/steward/src/utils.rs +++ b/programs/steward/src/utils.rs @@ -5,7 +5,7 @@ use borsh::{BorshDeserialize, BorshSerialize}; use spl_pod::{bytemuck::pod_from_bytes, primitives::PodU64, solana_program::program_pack::Pack}; use spl_stake_pool::{ big_vec::BigVec, - state::{ValidatorListHeader, ValidatorStakeInfo}, + state::{StakeStatus, ValidatorListHeader, ValidatorStakeInfo}, }; use crate::{errors::StewardError, Config, Delegation}; @@ -86,6 +86,43 @@ pub fn get_validator_stake_info_at_index( Ok(validator_stake_info) } +pub fn check_validator_list_has_stake_status( + validator_list_account_info: &AccountInfo, + flag: StakeStatus, +) -> Result { + let mut validator_list_data = validator_list_account_info.try_borrow_mut_data()?; + let (header, validator_list) = ValidatorListHeader::deserialize_vec(&mut validator_list_data)?; + require!( + header.account_type == spl_stake_pool::state::AccountType::ValidatorList, + StewardError::ValidatorListTypeMismatch + ); + + for index in 0..validator_list.len() as usize { + let stake_status_index = VEC_SIZE_BYTES + .saturating_add(index.saturating_mul(ValidatorStakeInfo::LEN)) + .checked_add(40) + .ok_or(StewardError::ArithmeticError)?; + + let stake_status = validator_list.data[stake_status_index]; + + if stake_status == flag as u8 { + return Ok(true); + } + } + + Ok(false) +} + +pub fn get_validator_list_length(validator_list_account_info: &AccountInfo) -> Result { + let mut validator_list_data = validator_list_account_info.try_borrow_mut_data()?; + let (header, validator_list) = ValidatorListHeader::deserialize_vec(&mut validator_list_data)?; + require!( + header.account_type == spl_stake_pool::state::AccountType::ValidatorList, + StewardError::ValidatorListTypeMismatch + ); + Ok(validator_list.len() as usize) +} + /// A boolean type stored as a u8. #[derive(BorshSerialize, BorshDeserialize, Debug, PartialEq, Eq)] #[zero_copy] diff --git a/tests/src/steward_fixtures.rs b/tests/src/steward_fixtures.rs index f7a97489..1ca223b4 100644 --- a/tests/src/steward_fixtures.rs +++ b/tests/src/steward_fixtures.rs @@ -1,3 +1,4 @@ +#![allow(clippy::await_holding_refcell_ref)] use std::{cell::RefCell, rc::Rc, str::FromStr, vec}; use crate::spl_stake_pool_cli; @@ -15,7 +16,7 @@ use jito_steward::{ constants::{MAX_VALIDATORS, SORTED_INDEX_DEFAULT, STAKE_POOL_WITHDRAW_SEED}, utils::StakePool, Config, Delegation, Parameters, Staker, StewardState, StewardStateAccount, StewardStateEnum, - UpdateParametersArgs, + UpdateParametersArgs, STATE_PADDING_0_SIZE, }; use solana_program_test::*; use solana_sdk::{ @@ -178,6 +179,30 @@ impl TestFixture { account } + pub async fn simulate_stake_pool_update(&self) { + let stake_pool: StakePool = self + .load_and_deserialize(&self.stake_pool_meta.stake_pool) + .await; + + let mut stake_pool_spl = stake_pool.as_ref().clone(); + + let current_epoch = self + .ctx + .borrow_mut() + .banks_client + .get_sysvar::() + .await + .unwrap() + .epoch; + + stake_pool_spl.last_update_epoch = current_epoch; + + self.ctx.borrow_mut().set_account( + &self.stake_pool_meta.stake_pool, + &serialized_stake_pool_account(stake_pool_spl, std::mem::size_of::()).into(), + ); + } + pub async fn initialize_stake_pool(&self) { // Call command_create_pool and execute transactions responded let mint = Keypair::new(); @@ -929,7 +954,10 @@ impl Default for StateMachineFixtures { instant_unstake: BitMask::default(), compute_delegations_completed: false.into(), rebalance_completed: false.into(), - _padding0: [0; 6 + 8 * MAX_VALIDATORS], + validators_added: 0, + checked_validators_removed_from_list: false.into(), + validators_to_remove: BitMask::default(), + _padding0: [0; STATE_PADDING_0_SIZE], }; StateMachineFixtures { diff --git a/tests/tests/steward/test_integration.rs b/tests/tests/steward/test_integration.rs index 8e158ef6..e8dcbc00 100644 --- a/tests/tests/steward/test_integration.rs +++ b/tests/tests/steward/test_integration.rs @@ -245,6 +245,7 @@ async fn test_compute_scores() { steward_state_account.state.current_epoch = clock.epoch; steward_state_account.state.next_cycle_epoch = clock.epoch + steward_config.parameters.num_epochs_between_scoring; + // steward_state_account.state.validators_added = MAX_VALIDATORS as u16; // Setup validator list let mut validator_list_validators = (0..MAX_VALIDATORS) @@ -284,6 +285,23 @@ async fn test_compute_scores() { &serialized_config(steward_config).into(), ); + fixture.simulate_stake_pool_update().await; + + let epoch_maintenance_ix = Instruction { + program_id: jito_steward::id(), + accounts: jito_steward::accounts::EpochMaintenance { + config: fixture.steward_config.pubkey(), + state_account: fixture.steward_state, + validator_list: fixture.stake_pool_meta.validator_list, + stake_pool: fixture.stake_pool_meta.stake_pool, + } + .to_account_metas(None), + data: jito_steward::instruction::EpochMaintenance { + validator_index_to_remove: None, + } + .data(), + }; + // Basic test - test score computation that requires most compute let compute_scores_ix = Instruction { program_id: jito_steward::id(), @@ -305,8 +323,9 @@ async fn test_compute_scores() { let tx = Transaction::new_signed_with_payer( &[ // Only high because we are averaging 512 epochs - ComputeBudgetInstruction::set_compute_unit_limit(600_000), + ComputeBudgetInstruction::set_compute_unit_limit(800_000), ComputeBudgetInstruction::request_heap_frame(128 * 1024), + epoch_maintenance_ix.clone(), compute_scores_ix.clone(), ], Some(&fixture.keypair.pubkey()), @@ -316,6 +335,8 @@ async fn test_compute_scores() { fixture.submit_transaction_assert_success(tx).await; + println!("Okay!"); + let mut steward_state_account: StewardStateAccount = fixture.load_and_deserialize(&fixture.steward_state).await; @@ -332,6 +353,30 @@ async fn test_compute_scores() { // Transition out of this state // Reset current state, set progress[1] to true, progress[0] to false + + { + // Reset Validator List, such that there are only 2 validators + let mut validator_list_validators = (0..2) + .map(|_| ValidatorStakeInfo { + vote_account_address: Pubkey::new_unique(), + ..ValidatorStakeInfo::default() + }) + .collect::>(); + validator_list_validators[0].vote_account_address = vote_account; + let validator_list = spl_stake_pool::state::ValidatorList { + header: ValidatorListHeader { + account_type: AccountType::ValidatorList, + max_validators: MAX_VALIDATORS as u32, + }, + validators: validator_list_validators, + }; + + fixture.ctx.borrow_mut().set_account( + &fixture.stake_pool_meta.validator_list, + &serialized_validator_list_account(validator_list, None).into(), + ); + } + steward_state_account.state.num_pool_validators = 2; steward_state_account.state.scores[..2].copy_from_slice(&[0, 0]); steward_state_account.state.yield_scores[..2].copy_from_slice(&[0, 0]); @@ -919,6 +964,7 @@ async fn test_rebalance_increase() { let add_validator_to_pool_ix = Instruction { program_id: jito_steward::id(), accounts: jito_steward::accounts::AutoAddValidator { + steward_state: fixture.steward_state, validator_history_account: validator_history_address, config: fixture.steward_config.pubkey(), stake_pool_program: spl_stake_pool::id(), @@ -1151,6 +1197,8 @@ async fn test_rebalance_decrease() { let add_validator_to_pool_ix = Instruction { program_id: jito_steward::id(), accounts: jito_steward::accounts::AutoAddValidator { + steward_state: fixture.steward_state, + validator_history_account: validator_history_address, config: fixture.steward_config.pubkey(), stake_pool_program: spl_stake_pool::id(), @@ -1392,6 +1440,7 @@ async fn test_rebalance_other_cases() { let add_validator_to_pool_ix = Instruction { program_id: jito_steward::id(), accounts: jito_steward::accounts::AutoAddValidator { + steward_state: fixture.steward_state, validator_history_account: validator_history_address, config: fixture.steward_config.pubkey(), stake_pool_program: spl_stake_pool::id(), diff --git a/tests/tests/steward/test_spl_passthrough.rs b/tests/tests/steward/test_spl_passthrough.rs index a3448572..b0a7b080 100644 --- a/tests/tests/steward/test_spl_passthrough.rs +++ b/tests/tests/steward/test_spl_passthrough.rs @@ -158,10 +158,28 @@ async fn _add_test_validator(fixture: &TestFixture, vote_account: Pubkey) { let (stake_account_address, _, withdraw_authority) = fixture.stake_accounts_for_validator(vote_account).await; + fixture.simulate_stake_pool_update().await; + + let epoch_maintenance_ix = Instruction { + program_id: jito_steward::id(), + accounts: jito_steward::accounts::EpochMaintenance { + config: fixture.steward_config.pubkey(), + state_account: fixture.steward_state, + validator_list: fixture.stake_pool_meta.validator_list, + stake_pool: fixture.stake_pool_meta.stake_pool, + } + .to_account_metas(None), + data: jito_steward::instruction::EpochMaintenance { + validator_index_to_remove: None, + } + .data(), + }; + // Add Validator let instruction = Instruction { program_id: jito_steward::id(), accounts: jito_steward::accounts::AddValidatorToPool { + steward_state: fixture.steward_state, config: fixture.steward_config.pubkey(), stake_pool_program: spl_stake_pool::id(), stake_pool: fixture.stake_pool_meta.stake_pool, @@ -189,7 +207,7 @@ async fn _add_test_validator(fixture: &TestFixture, vote_account: Pubkey) { let latest_blockhash = _get_latest_blockhash(fixture).await; let transaction = Transaction::new_signed_with_payer( - &[instruction], + &[epoch_maintenance_ix, instruction], Some(&fixture.keypair.pubkey()), &[&fixture.keypair], latest_blockhash, @@ -500,7 +518,7 @@ async fn test_add_validator_to_pool() { } { - // Add 5 validators + // Add 10 validators for _ in 0..10 { _add_test_validator(&fixture, Pubkey::new_unique()).await; } @@ -518,7 +536,7 @@ async fn test_remove_validator_from_pool() { fixture.initialize_steward_state().await; // Setup the steward state - _setup_test_steward_state(&fixture, MAX_VALIDATORS, 1_000_000_000).await; + // _setup_test_steward_state(&fixture, MAX_VALIDATORS, 1_000_000_000).await; // Assert the validator was added to the validator list _add_test_validator(&fixture, Pubkey::new_unique()).await; @@ -805,6 +823,7 @@ async fn test_decrease_additional_validator_stake() { let validator_list_account_raw = fixture .get_account(&fixture.stake_pool_meta.validator_list) .await; + let validator_list_account: ValidatorList = ValidatorList::try_deserialize_unchecked(&mut validator_list_account_raw.data.as_slice()) .expect("Failed to deserialize validator list account"); @@ -814,14 +833,20 @@ async fn test_decrease_additional_validator_stake() { .get(validator_list_index) .expect("Validator is not in list"); + println!("4"); + let vote_account = validator_to_increase_stake.vote_account_address; let (stake_account_address, transient_stake_account_address, withdraw_authority) = fixture.stake_accounts_for_validator(vote_account).await; + println!("5"); + _simulate_stake_deposit(&fixture, stake_account_address, 2_000_000_000).await; + println!("6"); let validator_history = fixture.initialize_validator_history_with_credits(vote_account, validator_list_index); + println!("7"); let (ephemeral_stake_account, _) = find_ephemeral_stake_program_address( &spl_stake_pool::id(), @@ -861,6 +886,7 @@ async fn test_decrease_additional_validator_stake() { }; let latest_blockhash = _get_latest_blockhash(&fixture).await; + println!("8"); let transaction = Transaction::new_signed_with_payer( &[instruction], @@ -869,6 +895,7 @@ async fn test_decrease_additional_validator_stake() { latest_blockhash, ); fixture.submit_transaction_assert_success(transaction).await; + println!("9"); drop(fixture); } diff --git a/tests/tests/steward/test_state_methods.rs b/tests/tests/steward/test_state_methods.rs index 6ec24a1e..da660681 100644 --- a/tests/tests/steward/test_state_methods.rs +++ b/tests/tests/steward/test_state_methods.rs @@ -71,6 +71,7 @@ fn test_compute_scores() { assert!(state.current_epoch == current_epoch); // Test invalid state + state.progress.reset(); state.state_tag = StewardStateEnum::Idle; let res = state.compute_score( clock, @@ -187,25 +188,26 @@ fn test_compute_scores() { assert!(res.is_ok()); assert!(state.start_computing_scores_slot == clock.slot); assert!(state.next_cycle_epoch == current_epoch + parameters.num_epochs_between_scoring); - assert!(state.current_epoch == current_epoch); assert!(state.num_pool_validators == 4); // 2) Progress stalled and time moved into next epoch // Conditions: clock.epoch > state.current_epoch and !state.progress.is_empty() - state.current_epoch = current_epoch - 1; - assert!(!state.progress.is_empty()); - assert!(state.current_epoch < clock.epoch); - let res = state.compute_score( - clock, - epoch_schedule, - &validators[0], - validators[0].index as usize, - cluster_history, - config, - state.num_pool_validators, - ); - assert!(res.is_ok()); - assert!(state.current_epoch == current_epoch); + // REDACTED: The epoch is now updated in the epoch_maintenance method + + // state.current_epoch = current_epoch - 1; + // assert!(!state.progress.is_empty()); + // assert!(state.current_epoch < clock.epoch); + // let res = state.compute_score( + // clock, + // epoch_schedule, + // &validators[0], + // validators[0].index as usize, + // cluster_history, + // config, + // state.num_pool_validators, + // ); + // assert!(res.is_ok()); + // assert!(state.current_epoch == current_epoch); // 3) Progress started, but took >1000 slots to complete // Conditions: start_computing_scores_slot > 1000 slots ago, !progress.is_empty(), and clock.epoch == state.current_epoch @@ -225,6 +227,8 @@ fn test_compute_scores() { state.num_pool_validators, ); assert!(res.is_ok()); + println!("{:?}", state.start_computing_scores_slot); + println!("{:?}", clock.slot); assert!(state.start_computing_scores_slot == clock.slot); } @@ -427,7 +431,19 @@ fn test_compute_instant_unstake_success() { .get(validators[0].index as usize) .unwrap()); + // Should skip validator since it's already been computed + let res = state.compute_instant_unstake( + clock, + epoch_schedule, + &validators[0], + validators[0].index as usize, + cluster_history, + config, + ); + assert!(res.is_ok()); + // Instant unstakeable validator + state.progress.reset(); state.instant_unstake.reset(); config .blacklist @@ -450,6 +466,7 @@ fn test_compute_instant_unstake_success() { // Instant unstakeable validator with no delegation amount state.delegations[validators[0].index as usize] = Delegation::new(0, 1); + state.progress.reset(); state.instant_unstake.reset(); let res = state.compute_instant_unstake( clock, @@ -571,6 +588,24 @@ fn test_rebalance() { _ => panic!("Expected RebalanceType::Decrease"), } + // Test that rebalance will be skipped if validator has already been run + let res = state.rebalance( + fixtures.current_epoch, + 1, + &validator_list_bigvec, + 4000 * LAMPORTS_PER_SOL, + 1000 * LAMPORTS_PER_SOL, + 0, + 0, + &fixtures.config.parameters, + ); + + assert!(res.is_ok()); + match res.unwrap() { + RebalanceType::None => {} + _ => panic!("Expected RebalanceType::None"), + } + // Instant unstake validator, but no delegation, so other delegations are not affected // Same scenario as above but out-of-band validator state.delegations[0..3].copy_from_slice(&[ @@ -589,6 +624,7 @@ fn test_rebalance() { // Validator index 1: 1000 SOL, 0.5 score, 0 delegation, -> Decrease stake, from "instant unstake" category, and set delegation to 0 // Validator index 2: 1000 SOL, 0 score, 0 delegation -> Decrease stake, from "regular unstake" category + state.progress.reset(); let res = state.rebalance( fixtures.current_epoch, 1, @@ -639,6 +675,8 @@ fn test_rebalance() { state.sorted_yield_score_indices[0..3].copy_from_slice(&[0, 1, 2]); state.instant_unstake.reset(); state.instant_unstake.set(0, true).unwrap(); + + state.progress.reset(); let res = state.rebalance( fixtures.current_epoch, 0, @@ -670,6 +708,7 @@ fn test_rebalance() { state.scores[0..3].copy_from_slice(&[1_000_000_000, 1_000_000_000, 1_000_000_000]); state.sorted_score_indices[0..3].copy_from_slice(&[0, 1, 2]); state.sorted_yield_score_indices[0..3].copy_from_slice(&[0, 1, 2]); + state.progress.reset(); let res = state.rebalance( fixtures.current_epoch, 0, diff --git a/tests/tests/steward/test_steward.rs b/tests/tests/steward/test_steward.rs index 5e7673d1..6f759cb0 100644 --- a/tests/tests/steward/test_steward.rs +++ b/tests/tests/steward/test_steward.rs @@ -41,6 +41,7 @@ async fn _auto_add_validator_to_pool(fixture: &TestFixture, vote_account: &Pubke let add_validator_to_pool_ix = Instruction { program_id: jito_steward::id(), accounts: jito_steward::accounts::AutoAddValidator { + steward_state: fixture.steward_state, validator_history_account, config: fixture.steward_config.pubkey(), stake_pool_program: spl_stake_pool::id(),