Christian/epoch update (#47)

Starting to add the solution for keeping track of removed validators

---------

Co-authored-by: Christian <[email protected]>

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,
 }

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<AutoAddValidator>) -> 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<AutoAddValidator>) -> 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(),

programs/steward/src/instructions/auto_remove_validator_from_pool.rs

@@ -143,6 +143,12 @@ pub fn handler(ctx: Context<AutoRemoveValidator>, 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<AutoRemoveValidator>, 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(

programs/steward/src/instructions/compute_delegations.rs

@@ -28,6 +28,11 @@ pub fn handler(ctx: Context<ComputeDelegations>) -> 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());
     }

programs/steward/src/instructions/compute_instant_unstake.rs

@@ -49,6 +49,11 @@ pub fn handler(ctx: Context<ComputeInstantUnstake>, 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<ComputeInstantUnstake>, validator_list_index: usize)
         &cluster,
         &config,
     )?;
+
     maybe_transition_and_emit(
         &mut state_account.state,
         &clock,

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<ComputeScore>, 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());

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<EpochMaintenance>,
+    validator_index_to_remove: Option<usize>,
+) -> 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(())
+}

programs/steward/src/instructions/idle.rs

@@ -33,6 +33,11 @@ pub fn handler(ctx: Context<Idle>) -> Result<()> {
         StewardError::InvalidState
     );
 
+    require!(
+        clock.epoch == state_account.state.current_epoch,
+        StewardError::EpochMaintenanceNotComplete
+    );
+
     if config.is_paused() {
         return Err(StewardError::StateMachinePaused.into());
     }

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::*;

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<ReallocState>) -> 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(())

programs/steward/src/instructions/rebalance.rs

@@ -161,6 +161,11 @@ pub fn handler(ctx: Context<Rebalance>, 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());
     }

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<ResetStewardState>) -> 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(())
 }