FIX: Fixed instant removal validator stuck scenario

programs/steward/idl/steward.json

@@ -1810,6 +1810,11 @@
       "code": 6029,
       "name": "ValidatorNeedsToBeMarkedForRemoval",
       "msg": "Validator needs to be marked for removal"
+    },
+    {
+      "code": 6030,
+      "name": "InvalidStakeState",
+      "msg": "Invalid stake state"
     }
   ],
   "types": [

programs/steward/src/errors.rs

@@ -64,4 +64,6 @@ pub enum StewardError {
     VoteAccountDoesNotMatch,
     #[msg("Validator needs to be marked for removal")]
     ValidatorNeedsToBeMarkedForRemoval,
+    #[msg("Invalid stake state")]
+    InvalidStakeState,
 }

programs/steward/src/instructions/instant_remove_validator.rs

@@ -1,13 +1,12 @@
 use crate::{
     errors::StewardError,
     utils::{
-        check_validator_list_has_stake_status_other_than, deserialize_stake_pool,
-        get_stake_pool_address, get_validator_list, get_validator_list_length,
+        deserialize_stake_pool, get_stake_pool_address, get_validator_list,
+        get_validator_list_length, tally_stake_status,
     },
     Config, StewardStateAccount,
 };
 use anchor_lang::prelude::*;
-use spl_stake_pool::state::StakeStatus;
 
 #[derive(Accounts)]
 pub struct InstantRemoveValidator<'info> {
@@ -40,7 +39,8 @@ pub fn handler(
     let mut state_account = ctx.accounts.state_account.load_mut()?;
 
     let clock = Clock::get()?;
-    let validators_to_remove = state_account.state.validators_for_immediate_removal.count();
+    let validators_for_immediate_removal =
+        state_account.state.validators_for_immediate_removal.count();
     let validators_in_list = get_validator_list_length(&ctx.accounts.validator_list)?;
 
     require!(
@@ -61,23 +61,27 @@ pub fn handler(
         StewardError::ValidatorNotInList
     );
 
-    // Ensure there are no validators in the list that have not been removed, that should be
+    let stake_status_tally = tally_stake_status(&ctx.accounts.validator_list)?;
+
+    let total_deactivating = stake_status_tally.deactivating_all
+        + stake_status_tally.deactivating_transient
+        + stake_status_tally.deactivating_validator
+        + stake_status_tally.ready_for_removal;
+
     require!(
-        !check_validator_list_has_stake_status_other_than(
-            &ctx.accounts.validator_list,
-            &[
-                StakeStatus::Active,
-                StakeStatus::DeactivatingAll,
-                StakeStatus::DeactivatingTransient
-            ]
-        )?,
+        total_deactivating == state_account.state.validators_to_remove.count() as u64,
         StewardError::ValidatorsHaveNotBeenRemoved
     );
 
+    require!(
+        stake_status_tally.ready_for_removal == 0,
+        StewardError::ValidatorsNeedToBeRemoved
+    );
+
     require!(
         state_account.state.num_pool_validators as usize
             + state_account.state.validators_added as usize
-            - validators_to_remove
+            - validators_for_immediate_removal
             == validators_in_list,
         StewardError::ListStateMismatch
     );

programs/steward/src/utils.rs

@@ -204,6 +204,78 @@ pub fn get_validator_stake_info_at_index(
     Ok(validator_stake_info)
 }
 
+pub struct StakeStatusTally {
+    pub active: u64,
+    pub deactivating_transient: u64,
+    pub ready_for_removal: u64,
+    pub deactivating_validator: u64,
+    pub deactivating_all: u64,
+}
+
+pub fn tally_stake_status(validator_list_account_info: &AccountInfo) -> Result<StakeStatusTally> {
+    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
+    );
+
+    let mut tally = StakeStatusTally {
+        active: 0,
+        deactivating_transient: 0,
+        ready_for_removal: 0,
+        deactivating_validator: 0,
+        deactivating_all: 0,
+    };
+
+    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(STAKE_STATUS_OFFSET)
+            .ok_or(StewardError::ArithmeticError)?;
+
+        let stake_status = validator_list.data[stake_status_index];
+
+        match stake_status {
+            x if x == StakeStatus::Active as u8 => {
+                tally.active = tally
+                    .active
+                    .checked_add(1)
+                    .ok_or(StewardError::ArithmeticError)?;
+            }
+            x if x == StakeStatus::DeactivatingTransient as u8 => {
+                tally.deactivating_transient = tally
+                    .deactivating_transient
+                    .checked_add(1)
+                    .ok_or(StewardError::ArithmeticError)?;
+            }
+            x if x == StakeStatus::ReadyForRemoval as u8 => {
+                tally.ready_for_removal = tally
+                    .ready_for_removal
+                    .checked_add(1)
+                    .ok_or(StewardError::ArithmeticError)?;
+            }
+            x if x == StakeStatus::DeactivatingValidator as u8 => {
+                tally.deactivating_validator = tally
+                    .deactivating_validator
+                    .checked_add(1)
+                    .ok_or(StewardError::ArithmeticError)?;
+            }
+            x if x == StakeStatus::DeactivatingAll as u8 => {
+                tally.deactivating_all = tally
+                    .deactivating_all
+                    .checked_add(1)
+                    .ok_or(StewardError::ArithmeticError)?;
+            }
+            _ => {
+                return Err(StewardError::InvalidStakeState.into());
+            }
+        }
+    }
+
+    Ok(tally)
+}
+
 pub fn check_validator_list_has_stake_status_other_than(
     validator_list_account_info: &AccountInfo,
     flags: &[StakeStatus],