can now remove and add validator in the same epoch

programs/steward/src/instructions/auto_add_validator_to_pool.rs

@@ -2,7 +2,7 @@ use crate::constants::{MAX_VALIDATORS, STAKE_POOL_WITHDRAW_SEED};
 use crate::errors::StewardError;
 use crate::events::AutoAddValidatorEvent;
 use crate::state::{Config, StewardStateAccount};
-use crate::utils::{deserialize_stake_pool, get_stake_pool_address};
+use crate::utils::{deserialize_stake_pool, get_stake_pool_address, get_validator_list_length};
 use anchor_lang::prelude::*;
 use anchor_lang::solana_program::{program::invoke_signed, stake, sysvar, vote};
 use spl_stake_pool::find_stake_program_address;
@@ -117,6 +117,16 @@ pub fn handler(ctx: Context<AutoAddValidator>) -> Result<()> {
             state_account.state.validators_for_immediate_removal.count() == 0,
             StewardError::ValidatorsNeedToBeRemoved
         );
+
+        let validators_in_list = get_validator_list_length(&ctx.accounts.validator_list)?;
+
+        // Cannot call auto remove if there is a validator mismatch
+        require!(
+            state_account.state.num_pool_validators as usize
+                + state_account.state.validators_added as usize
+                == validators_in_list,
+            StewardError::ListStateMismatch
+        );
     }
 
     let validator_list_len = {

programs/steward/src/instructions/auto_remove_validator_from_pool.rs

@@ -5,7 +5,8 @@ use crate::errors::StewardError;
 use crate::events::AutoRemoveValidatorEvent;
 use crate::state::Config;
 use crate::utils::{
-    deserialize_stake_pool, get_stake_pool_address, get_validator_stake_info_at_index,
+    deserialize_stake_pool, get_stake_pool_address, get_validator_list_length,
+    get_validator_stake_info_at_index,
 };
 use crate::StewardStateAccount;
 use anchor_lang::solana_program::{program::invoke_signed, stake, sysvar, vote};
@@ -168,6 +169,16 @@ pub fn handler(ctx: Context<AutoRemoveValidator>, validator_list_index: usize) -
             stake_account_deactivated || vote_account_closed,
             StewardError::ValidatorNotRemovable
         );
+
+        let validators_in_list = get_validator_list_length(&validator_list)?;
+
+        // Cannot call auto remove if there is a validator mismatch
+        require!(
+            state_account.state.num_pool_validators as usize
+                + state_account.state.validators_added as usize
+                == validators_in_list,
+            StewardError::ListStateMismatch
+        );
     }
 
     {

programs/steward/src/instructions/epoch_maintenance.rs

@@ -55,11 +55,6 @@ pub fn handler(
         StewardError::EpochMaintenanceAlreadyComplete
     );
 
-    require!(
-        state_account.state.validators_for_immediate_removal.count() == 0,
-        StewardError::ValidatorsNeedToBeRemoved
-    );
-
     // Ensure there are no validators in the list that have not been removed, that should be
     require!(
         !check_validator_list_has_stake_status_other_than(
@@ -75,7 +70,8 @@ pub fn handler(
         // 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();
+        let validators_to_remove = state_account.state.validators_to_remove.count()
+            + state_account.state.validators_for_immediate_removal.count();
 
         // Ensure we have a 1-1 mapping between the number of validators in the list and the number of validators in the state
         // If we don't have this mapping, everything needs to be removed

programs/steward/src/instructions/instant_remove_validator.rs

@@ -45,17 +45,6 @@ pub fn handler(
     let validators_to_remove = state_account.state.validators_for_immediate_removal.count();
     let validators_in_list = get_validator_list_length(&ctx.accounts.validator_list)?;
 
-    require!(
-        validators_to_remove > 0,
-        StewardError::NoValidatorsNeedToBeRemoved
-    );
-
-    // Has to be done on the same epoch they are marked. If not, state must be reset
-    require!(
-        state_account.state.current_epoch == clock.epoch,
-        StewardError::EpochMaintenanceAlreadyComplete
-    );
-
     require!(
         clock.epoch == stake_pool.last_update_epoch,
         StewardError::StakePoolNotUpdated

programs/steward/src/state/steward_state.rs

@@ -475,10 +475,19 @@ impl StewardState {
             StewardError::ValidatorNotMarkedForRemoval
         );
 
-        self.num_pool_validators = self
-            .num_pool_validators
-            .checked_sub(1)
-            .ok_or(StewardError::ArithmeticError)?;
+        // If the validator was marked for removal in the current cycle, decrement validators_added
+        if index >= self.num_pool_validators as usize {
+            self.validators_added = self
+                .validators_added
+                .checked_sub(1)
+                .ok_or(StewardError::ArithmeticError)?;
+        } else {
+            self.num_pool_validators = self
+                .num_pool_validators
+                .checked_sub(1)
+                .ok_or(StewardError::ArithmeticError)?;
+        }
+
         let num_pool_validators = self.num_pool_validators as usize;
 
         // Shift all validator state to the left