instant removal

programs/steward/src/errors.rs

@@ -48,6 +48,8 @@ pub enum StewardError {
     ArithmeticError,
     #[msg("Validator not eligible for removal. Must be delinquent or have closed vote account")]
     ValidatorNotRemovable,
+    #[msg("Validator was marked active when it should be deactivating")]
+    ValidatorMarkedActive,
     #[msg("Max validators reached")]
     MaxValidatorsReached,
     #[msg("Validator history account does not match vote account")]
@@ -56,6 +58,12 @@ pub enum StewardError {
     EpochMaintenanceNotComplete,
     #[msg("The stake pool must be updated before continuing")]
     StakePoolNotUpdated,
+    #[msg("Epoch Maintenance has already been completed")]
+    EpochMaintenanceAlreadyComplete,
+    #[msg("Validators are marked for immediate removal")]
+    ValidatorsNeedToBeRemoved,
+    #[msg("No validators are marked for immediate removal")]
+    NoValidatorsNeedToBeRemoved,
     #[msg("Validator not marked for removal")]
     ValidatorNotMarkedForRemoval,
     #[msg("Validators have not been removed")]

programs/steward/src/events.rs

@@ -12,6 +12,7 @@ pub struct AutoRemoveValidatorEvent {
     pub vote_account: Pubkey,
     pub vote_account_closed: bool,
     pub stake_account_deactivated: bool,
+    pub marked_for_immediate_removal: bool,
 }
 
 #[event]

programs/steward/src/instructions/auto_add_validator_to_pool.rs

@@ -107,10 +107,17 @@ pub fn handler(ctx: Context<AutoAddValidator>) -> Result<()> {
     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
-    );
+    {
+        require!(
+            epoch == state_account.state.current_epoch,
+            StewardError::EpochMaintenanceNotComplete
+        );
+
+        require!(
+            state_account.state.validators_for_immediate_removal.count() == 0,
+            StewardError::ValidatorsNeedToBeRemoved
+        );
+    }
 
     let validator_list_len = {
         let validator_list_data = &mut ctx.accounts.validator_list.try_borrow_mut_data()?;

programs/steward/src/instructions/auto_remove_validator_from_pool.rs

@@ -12,6 +12,7 @@ use anchor_lang::solana_program::{program::invoke_signed, stake, sysvar, vote};
 use anchor_lang::{prelude::*, system_program};
 use spl_pod::solana_program::borsh1::try_from_slice_unchecked;
 use spl_pod::solana_program::stake::state::StakeStateV2;
+use spl_stake_pool::state::StakeStatus;
 use spl_stake_pool::{find_stake_program_address, find_transient_stake_program_address};
 use validator_history::state::ValidatorHistory;
 
@@ -124,91 +125,115 @@ pub struct AutoRemoveValidator<'info> {
 
 */
 pub fn handler(ctx: Context<AutoRemoveValidator>, validator_list_index: usize) -> Result<()> {
+    let state_account = ctx.accounts.state_account.load().unwrap();
+    let validator_list = &ctx.accounts.validator_list;
+    let epoch = Clock::get()?.epoch;
+
+    let validator_stake_info =
+        get_validator_stake_info_at_index(validator_list, validator_list_index)?;
+    require!(
+        validator_stake_info.vote_account_address == ctx.accounts.vote_account.key(),
+        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();
+        let stake_state: StakeStateV2 =
+            try_from_slice_unchecked::<StakeStateV2>(stake_account_data)?;
+
+        let deactivation_epoch = match stake_state {
+            StakeStateV2::Stake(_meta, stake, _stake_flags) => stake.delegation.deactivation_epoch,
+            _ => return Err(StewardError::InvalidState.into()), // TODO fix
+        };
+        deactivation_epoch < epoch
+    };
+
+    // Check if vote account closed
+    let vote_account_closed = ctx.accounts.vote_account.owner == &system_program::ID;
+
+    require!(
+        stake_account_deactivated || vote_account_closed,
+        StewardError::ValidatorNotRemovable
+    );
+
+    {
+        invoke_signed(
+            &spl_stake_pool::instruction::remove_validator_from_pool(
+                &ctx.accounts.stake_pool_program.key(),
+                &ctx.accounts.stake_pool.key(),
+                &ctx.accounts.state_account.key(),
+                &ctx.accounts.withdraw_authority.key(),
+                &ctx.accounts.validator_list.key(),
+                &ctx.accounts.stake_account.key(),
+                &ctx.accounts.transient_stake_account.key(),
+            ),
+            &[
+                ctx.accounts.stake_pool.to_account_info(),
+                ctx.accounts.state_account.to_account_info(),
+                ctx.accounts.reserve_stake.to_owned(),
+                ctx.accounts.withdraw_authority.to_owned(),
+                ctx.accounts.validator_list.to_account_info(),
+                ctx.accounts.stake_account.to_account_info(),
+                ctx.accounts.transient_stake_account.to_account_info(),
+                ctx.accounts.vote_account.to_account_info(),
+                ctx.accounts.rent.to_account_info(),
+                ctx.accounts.clock.to_account_info(),
+                ctx.accounts.stake_history.to_account_info(),
+                ctx.accounts.stake_config.to_account_info(),
+                ctx.accounts.system_program.to_account_info(),
+                ctx.accounts.stake_program.to_account_info(),
+            ],
+            &[&[
+                StewardStateAccount::SEED,
+                &ctx.accounts.config.key().to_bytes(),
+                &[ctx.bumps.state_account],
+            ]],
+        )?;
+    }
+
     {
+        // Read the state account again
         let mut state_account = ctx.accounts.state_account.load_mut()?;
         let validator_list = &ctx.accounts.validator_list;
-        let epoch = Clock::get()?.epoch;
-
         let validator_stake_info =
             get_validator_stake_info_at_index(validator_list, validator_list_index)?;
-        require!(
-            validator_stake_info.vote_account_address == ctx.accounts.vote_account.key(),
-            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();
-            let stake_state: StakeStateV2 =
-                try_from_slice_unchecked::<StakeStateV2>(stake_account_data)?;
-
-            let deactivation_epoch = match stake_state {
-                StakeStateV2::Stake(_meta, stake, _stake_flags) => {
-                    stake.delegation.deactivation_epoch
-                }
-                _ => return Err(StewardError::InvalidState.into()), // TODO fix
-            };
-            deactivation_epoch < epoch
-        };
-
-        // Check if vote account closed
-        let vote_account_closed = ctx.accounts.vote_account.owner == &system_program::ID;
 
-        require!(
-            stake_account_deactivated || vote_account_closed,
-            StewardError::ValidatorNotRemovable
-        );
-
-        state_account
-            .state
-            .mark_validator_for_removal(validator_list_index)?;
+        let stake_status = StakeStatus::try_from(validator_stake_info.status).unwrap();
+        let marked_for_immediate_removal: bool;
+
+        match stake_status {
+            StakeStatus::Active => {
+                // Should never happen
+                return Err(StewardError::ValidatorMarkedActive.into());
+            }
+            StakeStatus::DeactivatingValidator | StakeStatus::ReadyForRemoval => {
+                marked_for_immediate_removal = true;
+                // Mark for immediate removal
+            }
+            StakeStatus::DeactivatingAll | StakeStatus::DeactivatingTransient => {
+                marked_for_immediate_removal = false;
+                state_account
+                    .state
+                    .mark_validator_for_removal(validator_list_index)?;
+                // Mark for next Epoch removal
+            }
+        }
 
         emit!(AutoRemoveValidatorEvent {
             vote_account: ctx.accounts.vote_account.key(),
             validator_list_index: validator_list_index as u64,
             stake_account_deactivated,
             vote_account_closed,
+            marked_for_immediate_removal,
         });
     }
 
-    invoke_signed(
-        &spl_stake_pool::instruction::remove_validator_from_pool(
-            &ctx.accounts.stake_pool_program.key(),
-            &ctx.accounts.stake_pool.key(),
-            &ctx.accounts.state_account.key(),
-            &ctx.accounts.withdraw_authority.key(),
-            &ctx.accounts.validator_list.key(),
-            &ctx.accounts.stake_account.key(),
-            &ctx.accounts.transient_stake_account.key(),
-        ),
-        &[
-            ctx.accounts.stake_pool.to_account_info(),
-            ctx.accounts.state_account.to_account_info(),
-            ctx.accounts.reserve_stake.to_owned(),
-            ctx.accounts.withdraw_authority.to_owned(),
-            ctx.accounts.validator_list.to_account_info(),
-            ctx.accounts.stake_account.to_account_info(),
-            ctx.accounts.transient_stake_account.to_account_info(),
-            ctx.accounts.vote_account.to_account_info(),
-            ctx.accounts.rent.to_account_info(),
-            ctx.accounts.clock.to_account_info(),
-            ctx.accounts.stake_history.to_account_info(),
-            ctx.accounts.stake_config.to_account_info(),
-            ctx.accounts.system_program.to_account_info(),
-            ctx.accounts.stake_program.to_account_info(),
-        ],
-        &[&[
-            StewardStateAccount::SEED,
-            &ctx.accounts.config.key().to_bytes(),
-            &[ctx.bumps.state_account],
-        ]],
-    )?;
-
     Ok(())
 }

programs/steward/src/instructions/compute_delegations.rs

@@ -1,5 +1,5 @@
 use crate::errors::StewardError;
-use crate::{maybe_transition_and_emit, Config, StewardStateAccount};
+use crate::{maybe_transition_and_emit, Config, StewardStateAccount, StewardStateEnum};
 use anchor_lang::prelude::*;
 
 #[derive(Accounts)]
@@ -25,13 +25,28 @@ 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
-    );
+    {
+        require!(
+            matches!(
+                state_account.state.state_tag,
+                StewardStateEnum::ComputeDelegations
+            ),
+            StewardError::InvalidState
+        );
 
-    if config.is_paused() {
-        return Err(StewardError::StateMachinePaused.into());
+        require!(
+            clock.epoch == state_account.state.current_epoch,
+            StewardError::EpochMaintenanceNotComplete
+        );
+
+        require!(
+            state_account.state.validators_for_immediate_removal.count() == 0,
+            StewardError::ValidatorsNeedToBeRemoved
+        );
+
+        if config.is_paused() {
+            return Err(StewardError::StateMachinePaused.into());
+        }
     }
 
     state_account

programs/steward/src/instructions/compute_instant_unstake.rs

@@ -2,7 +2,7 @@ use crate::{
     errors::StewardError,
     maybe_transition_and_emit,
     utils::{get_validator_list, get_validator_stake_info_at_index},
-    Config, StewardStateAccount,
+    Config, StewardStateAccount, StewardStateEnum,
 };
 use anchor_lang::prelude::*;
 use validator_history::{ClusterHistory, ValidatorHistory};
@@ -41,22 +41,37 @@ pub fn handler(ctx: Context<ComputeInstantUnstake>, validator_list_index: usize)
     let clock = Clock::get()?;
     let epoch_schedule = EpochSchedule::get()?;
 
+    {
+        require!(
+            matches!(
+                state_account.state.state_tag,
+                StewardStateEnum::ComputeInstantUnstake
+            ),
+            StewardError::InvalidState
+        );
+
+        require!(
+            clock.epoch == state_account.state.current_epoch,
+            StewardError::EpochMaintenanceNotComplete
+        );
+
+        require!(
+            state_account.state.validators_for_immediate_removal.count() == 0,
+            StewardError::ValidatorsNeedToBeRemoved
+        );
+
+        if config.is_paused() {
+            return Err(StewardError::StateMachinePaused.into());
+        }
+    }
+
     let validator_stake_info =
         get_validator_stake_info_at_index(validator_list, validator_list_index)?;
     require!(
         validator_stake_info.vote_account_address == validator_history.vote_account,
         StewardError::ValidatorNotInList
     );
 
-    require!(
-        clock.epoch == state_account.state.current_epoch,
-        StewardError::EpochMaintenanceNotComplete
-    );
-
-    if config.is_paused() {
-        return Err(StewardError::StateMachinePaused.into());
-    }
-
     if let Some(instant_unstake) = state_account.state.compute_instant_unstake(
         &clock,
         &epoch_schedule,

programs/steward/src/instructions/compute_score.rs

@@ -42,6 +42,22 @@ pub fn handler(ctx: Context<ComputeScore>, validator_list_index: usize) -> Resul
     let clock: Clock = Clock::get()?;
     let epoch_schedule = EpochSchedule::get()?;
 
+    {
+        require!(
+            clock.epoch == state_account.state.current_epoch,
+            StewardError::EpochMaintenanceNotComplete
+        );
+
+        require!(
+            state_account.state.validators_for_immediate_removal.count() == 0,
+            StewardError::ValidatorsNeedToBeRemoved
+        );
+
+        if config.is_paused() {
+            return Err(StewardError::StateMachinePaused.into());
+        }
+    }
+
     let validator_stake_info =
         get_validator_stake_info_at_index(validator_list, validator_list_index)?;
     require!(
@@ -51,15 +67,6 @@ pub fn handler(ctx: Context<ComputeScore>, validator_list_index: usize) -> Resul
 
     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());
-    }
-
     // May need to force an extra transition here in case cranking got stuck in any previous state
     // and it's now the start of a new scoring cycle
     if !matches!(
@@ -73,6 +80,7 @@ pub fn handler(ctx: Context<ComputeScore>, validator_list_index: usize) -> Resul
             &epoch_schedule,
         )?;
     }
+
     require!(
         matches!(
             state_account.state.state_tag,