little tweaks

jito-foundation · Jul 22, 2024 · 41a4f22 · 41a4f22
1 parent 4bb062c
commit 41a4f22
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 11 deletions.
diff --git a/programs/steward/src/instructions/auto_remove_validator_from_pool.rs b/programs/steward/src/instructions/auto_remove_validator_from_pool.rs
@@ -129,7 +129,7 @@ pub fn handler(ctx: Context<AutoRemoveValidator>, validator_list_index: usize) -
     let vote_account_closed;
 
     {
-        let state_account = ctx.accounts.state_account.load().unwrap();
+        let state_account = ctx.accounts.state_account.load()?;
         let validator_list = &ctx.accounts.validator_list;
         let epoch = Clock::get()?.epoch;
 
@@ -212,7 +212,7 @@ pub fn handler(ctx: Context<AutoRemoveValidator>, validator_list_index: usize) -
         let validator_stake_info =
             get_validator_stake_info_at_index(validator_list, validator_list_index)?;
 
-        let stake_status = StakeStatus::try_from(validator_stake_info.status).unwrap();
+        let stake_status = StakeStatus::try_from(validator_stake_info.status)?;
         let marked_for_immediate_removal: bool;
 
         match stake_status {

diff --git a/programs/steward/src/instructions/epoch_maintenance.rs b/programs/steward/src/instructions/epoch_maintenance.rs
@@ -3,7 +3,7 @@ use crate::{
     events::EpochMaintenanceEvent,
     utils::{
         check_validator_list_has_stake_status_other_than, deserialize_stake_pool,
-        get_stake_pool_address, get_validator_list_length,
+        get_stake_pool_address, get_validator_list, get_validator_list_length,
     },
     Config, StewardStateAccount, CHECKED_VALIDATORS_REMOVED_FROM_LIST, COMPUTE_INSTANT_UNSTAKES,
     EPOCH_MAINTENANCE, POST_LOOP_IDLE, PRE_LOOP_IDLE, REBALANCE, RESET_TO_IDLE,
@@ -23,7 +23,7 @@ pub struct EpochMaintenance<'info> {
     pub state_account: AccountLoader<'info, StewardStateAccount>,
 
     /// CHECK: Correct account guaranteed if address is correct
-    #[account(address = deserialize_stake_pool(&stake_pool)?.validator_list)]
+    #[account(address = get_validator_list(&config)?)]
     pub validator_list: AccountInfo<'info>,
 
     /// CHECK: Correct account guaranteed if address is correct

diff --git a/programs/steward/src/instructions/instant_remove_validator.rs b/programs/steward/src/instructions/instant_remove_validator.rs
@@ -2,7 +2,7 @@ use crate::{
     errors::StewardError,
     utils::{
         check_validator_list_has_stake_status_other_than, deserialize_stake_pool,
-        get_stake_pool_address, get_validator_list_length,
+        get_stake_pool_address, get_validator_list, get_validator_list_length,
     },
     Config, StewardStateAccount,
 };
@@ -21,7 +21,7 @@ pub struct InstantRemoveValidator<'info> {
     pub state_account: AccountLoader<'info, StewardStateAccount>,
 
     /// CHECK: Correct account guaranteed if address is correct
-    #[account(address = deserialize_stake_pool(&stake_pool)?.validator_list)]
+    #[account(address = get_validator_list(&config)?)]
     pub validator_list: AccountInfo<'info>,
 
     /// CHECK: Correct account guaranteed if address is correct
@@ -31,9 +31,7 @@ pub struct InstantRemoveValidator<'info> {
     pub stake_pool: AccountInfo<'info>,
 }
 
-/// Runs maintenance tasks at the start of each epoch, needs to be run multiple times
-/// Routines:
-/// - Remove delinquent validators
+/// Removes validators from the pool that have been marked for immediate removal
 pub fn handler(
     ctx: Context<InstantRemoveValidator>,
     validator_index_to_remove: usize,
@@ -54,8 +52,7 @@ pub fn handler(
         state_account
             .state
             .validators_for_immediate_removal
-            .get(validator_index_to_remove)
-            .unwrap(),
+            .get(validator_index_to_remove)?,
         StewardError::ValidatorNotInList
     );