Fixes: rebalance decrease amount calculation, validator_lamport_balance initialization, tweaks to keeper

.github/workflows/build.yaml

@@ -129,6 +129,8 @@ jobs:
         with:
           submodules: recursive
       - uses: actions-rust-lang/setup-rust-toolchain@v1
+        with:
+          toolchain: nightly-2024-02-04
       - uses: actions/download-artifact@v4
         with:
           name: validator_history.so
@@ -138,18 +140,12 @@ jobs:
           name: jito_steward.so
           path: target/deploy/
       - name: cargo test
-        run: cargo test --package tests --all-features --color auto -- --skip steward::test_state_methods
+        run: cargo test --package tests --all-features --color auto
         shell: bash
         env:
           RUST_LOG: trace
           SBF_OUT_DIR: ${{ github.workspace }}/target/deploy
           RUST_MIN_STACK: 5000000
-      - name: cargo test steward::test_state_methods
-        run: cargo test --package tests --test mod steward::test_state_methods
-        shell: bash
-        env:
-          RUST_LOG: trace
-          RUST_MIN_STACK: 5000000
 
   # release only runs on tagged commits
   # it should wait for all the other steps to finish, to ensure releases are the highest quality

programs/steward/idl/steward.json

@@ -1270,6 +1270,36 @@
       ],
       "args": []
     },
+    {
+      "name": "reset_validator_lamport_balances",
+      "docs": [
+        "Reset validator_lamport_balances to default"
+      ],
+      "discriminator": [
+        69,
+        111,
+        124,
+        69,
+        69,
+        29,
+        96,
+        181
+      ],
+      "accounts": [
+        {
+          "name": "steward_state",
+          "writable": true
+        },
+        {
+          "name": "config"
+        },
+        {
+          "name": "authority",
+          "signer": true
+        }
+      ],
+      "args": []
+    },
     {
       "name": "resume_steward",
       "docs": [

programs/steward/src/constants.rs

@@ -8,6 +8,7 @@ pub const MAX_VALIDATORS: usize = 5_000;
 pub const BASIS_POINTS_MAX: u16 = 10_000;
 pub const COMMISSION_MAX: u8 = 100;
 pub const SORTED_INDEX_DEFAULT: u16 = u16::MAX;
+pub const LAMPORT_BALANCE_DEFAULT: u64 = u64::MAX;
 // Need at least 1% of slots remaining (4320 slots) to execute steps in state machine
 pub const EPOCH_PROGRESS_MAX: f64 = 0.99;
 // Cannot go more than 100 epochs without scoring

programs/steward/src/delegation.rs

@@ -1,6 +1,7 @@
 use anchor_lang::prelude::*;
 use spl_stake_pool::big_vec::BigVec;
 
+use crate::constants::LAMPORT_BALANCE_DEFAULT;
 use crate::events::DecreaseComponents;
 use crate::{
     errors::StewardError,
@@ -22,10 +23,12 @@ pub enum RebalanceType {
 ///
 /// Unstaking is calculated this way because stake account balances can change at any time from users' stake withdrawals and deposits,
 /// and this ensures that unstaking is done fairly at the time of the rebalance. In addition, these instructions can run in any order.
+#[allow(clippy::too_many_arguments)]
 pub fn decrease_stake_calculation(
     state: &StewardState,
     target_index: usize,
     mut unstake_state: UnstakeState,
+    current_lamports: u64, // active lamports in target stake account adjusted for minimum delegation
     stake_pool_lamports: u64,
     validator_list: &BigVec<'_>,
     minimum_delegation: u64,
@@ -56,6 +59,10 @@ pub fn decrease_stake_calculation(
         // ValidatorList includes base lamports in active_stake_lamports
         temp_current_lamports = temp_current_lamports.saturating_sub(base_lamport_balance);
 
+        if temp_index == target_index {
+            temp_current_lamports = current_lamports;
+        }
+
         // For the current `temp` validator, calculate how much we can remove and what category it's coming from
         let unstake_amounts =
             if !some_transient_stake && temp_target_lamports < temp_current_lamports {
@@ -245,6 +252,7 @@ impl UnstakeState {
         // either to the target or to the previous balance before the deposit, whichever is lower in terms of total lamports unstaked
         if current_lamports > state.validator_lamport_balances[index]
             && self.stake_deposit_unstake_total < self.stake_deposit_unstake_cap
+            && state.validator_lamport_balances[index] != LAMPORT_BALANCE_DEFAULT
         {
             let lamports_above_target = current_lamports
                 .checked_sub(target_lamports)

programs/steward/src/instructions/mod.rs

@@ -16,6 +16,7 @@ pub mod realloc_state;
 pub mod rebalance;
 pub mod remove_validator_from_blacklist;
 pub mod reset_steward_state;
+pub mod reset_validator_lamport_balances;
 pub mod resume_steward;
 pub mod set_new_authority;
 pub mod spl_passthrough;
@@ -38,6 +39,7 @@ pub use realloc_state::*;
 pub use rebalance::*;
 pub use remove_validator_from_blacklist::*;
 pub use reset_steward_state::*;
+pub use reset_validator_lamport_balances::*;
 pub use resume_steward::*;
 pub use set_new_authority::*;
 pub use spl_passthrough::*;

programs/steward/src/instructions/realloc_state.rs

@@ -1,6 +1,6 @@
 use crate::{
     bitmask::BitMask,
-    constants::{MAX_ALLOC_BYTES, MAX_VALIDATORS, SORTED_INDEX_DEFAULT},
+    constants::{LAMPORT_BALANCE_DEFAULT, MAX_ALLOC_BYTES, MAX_VALIDATORS, SORTED_INDEX_DEFAULT},
     errors::StewardError,
     state::{Config, StewardStateAccount},
     utils::get_validator_list,
@@ -74,6 +74,7 @@ pub fn handler(ctx: Context<ReallocState>) -> Result<()> {
 
         state_account.state.state_tag = StewardStateEnum::ComputeScores;
         state_account.state.num_pool_validators = validator_list.len() as u64;
+        state_account.state.validator_lamport_balances = [LAMPORT_BALANCE_DEFAULT; MAX_VALIDATORS];
         state_account.state.scores = [0; MAX_VALIDATORS];
         state_account.state.sorted_score_indices = [SORTED_INDEX_DEFAULT; MAX_VALIDATORS];
         state_account.state.yield_scores = [0; MAX_VALIDATORS];

programs/steward/src/instructions/reset_steward_state.rs

@@ -1,5 +1,5 @@
 use crate::{
-    constants::{MAX_VALIDATORS, SORTED_INDEX_DEFAULT},
+    constants::{LAMPORT_BALANCE_DEFAULT, MAX_VALIDATORS, SORTED_INDEX_DEFAULT},
     errors::StewardError,
     state::{Config, StewardStateAccount},
     utils::{deserialize_stake_pool, get_config_admin, get_stake_pool_address},
@@ -47,6 +47,7 @@ pub fn handler(ctx: Context<ResetStewardState>) -> Result<()> {
 
     state_account.state.state_tag = StewardStateEnum::ComputeScores;
     state_account.state.num_pool_validators = validator_list.len() as u64;
+    state_account.state.validator_lamport_balances = [LAMPORT_BALANCE_DEFAULT; MAX_VALIDATORS];
     state_account.state.scores = [0; MAX_VALIDATORS];
     state_account.state.sorted_score_indices = [SORTED_INDEX_DEFAULT; MAX_VALIDATORS];
     state_account.state.yield_scores = [0; MAX_VALIDATORS];

programs/steward/src/instructions/reset_validator_lamport_balances.rs

@@ -0,0 +1,29 @@
+use crate::{
+    constants::{LAMPORT_BALANCE_DEFAULT, MAX_VALIDATORS},
+    state::*,
+    utils::get_config_admin,
+};
+use anchor_lang::prelude::*;
+
+#[derive(Accounts)]
+pub struct ResetValidatorLamportBalances<'info> {
+    #[account(
+        mut,
+        seeds = [StewardStateAccount::SEED, config.key().as_ref()],
+        bump
+    )]
+    pub steward_state: AccountLoader<'info, StewardStateAccount>,
+
+    pub config: AccountLoader<'info, Config>,
+
+    #[account(address = get_config_admin(&config)?)]
+    pub authority: Signer<'info>,
+}
+
+pub fn handler(ctx: Context<ResetValidatorLamportBalances>) -> Result<()> {
+    let state_account = &mut ctx.accounts.steward_state.load_mut()?;
+
+    state_account.state.validator_lamport_balances = [LAMPORT_BALANCE_DEFAULT; MAX_VALIDATORS];
+
+    Ok(())
+}

programs/steward/src/instructions/spl_passthrough.rs

@@ -4,7 +4,7 @@
 // is that the config, stake pool address, staker, signer, and sometimes state account match up.
 // Otherwise these instructions are intented to be minimally restrictive.
 
-use crate::constants::MAX_VALIDATORS;
+use crate::constants::{LAMPORT_BALANCE_DEFAULT, MAX_VALIDATORS};
 use crate::errors::StewardError;
 use crate::state::Config;
 use crate::utils::{
@@ -400,9 +400,11 @@ pub fn increase_validator_stake_handler(
             .ok_or(StewardError::ValidatorIndexOutOfBounds)?;
 
         // Set the balance
-        *balance = balance
-            .checked_add(lamports)
-            .ok_or(StewardError::ArithmeticError)?;
+        if *balance != LAMPORT_BALANCE_DEFAULT {
+            *balance = balance
+                .checked_add(lamports)
+                .ok_or(StewardError::ArithmeticError)?;
+        }
     }
 
     invoke_signed(
@@ -521,9 +523,11 @@ pub fn decrease_validator_stake_handler(
             .ok_or(StewardError::ValidatorIndexOutOfBounds)?;
 
         // Set the balance
-        *balance = balance
-            .checked_sub(lamports)
-            .ok_or(StewardError::ArithmeticError)?;
+        if *balance != LAMPORT_BALANCE_DEFAULT {
+            *balance = balance
+                .checked_sub(lamports)
+                .ok_or(StewardError::ArithmeticError)?;
+        }
     }
 
     invoke_signed(
@@ -641,9 +645,11 @@ pub fn increase_additional_validator_stake_handler(
             .ok_or(StewardError::ValidatorIndexOutOfBounds)?;
 
         // Set the balance
-        *balance = balance
-            .checked_add(lamports)
-            .ok_or(StewardError::ArithmeticError)?;
+        if *balance != LAMPORT_BALANCE_DEFAULT {
+            *balance = balance
+                .checked_add(lamports)
+                .ok_or(StewardError::ArithmeticError)?;
+        }
     }
 
     invoke_signed(
@@ -765,9 +771,11 @@ pub fn decrease_additional_validator_stake_handler(
             .ok_or(StewardError::ValidatorIndexOutOfBounds)?;
 
         // Set the balance
-        *balance = balance
-            .checked_sub(lamports)
-            .ok_or(StewardError::ArithmeticError)?;
+        if *balance != LAMPORT_BALANCE_DEFAULT {
+            *balance = balance
+                .checked_sub(lamports)
+                .ok_or(StewardError::ArithmeticError)?;
+        }
     }
 
     invoke_signed(

programs/steward/src/lib.rs

@@ -207,6 +207,13 @@ pub mod steward {
         )
     }
 
+    /// Reset validator_lamport_balances to default
+    pub fn reset_validator_lamport_balances(
+        ctx: Context<ResetValidatorLamportBalances>,
+    ) -> Result<()> {
+        instructions::reset_validator_lamport_balances::handler(ctx)
+    }
+
     /// Closes Steward PDA accounts associated with a given Config (StewardStateAccount, and Staker).
     /// Config is not closed as it is a Keypair, so lamports can simply be withdrawn.
     /// Reclaims lamports to authority

programs/steward/src/state/steward_state.rs

@@ -3,7 +3,7 @@ use std::fmt::Display;
 
 use crate::{
     bitmask::BitMask,
-    constants::{MAX_VALIDATORS, SORTED_INDEX_DEFAULT},
+    constants::{LAMPORT_BALANCE_DEFAULT, MAX_VALIDATORS, SORTED_INDEX_DEFAULT},
     delegation::{
         decrease_stake_calculation, increase_stake_calculation, RebalanceType, UnstakeState,
     },
@@ -859,12 +859,6 @@ impl StewardState {
                 .checked_add(stake_rent)
                 .ok_or(StewardError::ArithmeticError)?;
 
-            msg!("Reserve lamports before adjustment: {}", reserve_lamports);
-            msg!(
-                "Stake pool lamports before adjustment: {}",
-                stake_pool_lamports
-            );
-
             // Maximum increase amount is the total lamports in the reserve stake account minus (num_validators + 1) * stake_rent, which covers rent for all validators plus the transient rent
             let all_accounts_needed_reserve_for_rent = validator_list
                 .len()
@@ -907,6 +901,11 @@ impl StewardState {
                 In all cases where the current_lamports is now below the target or internal balance, we update the internal balance.
                 Otherwise, keep the internal balance the same to ensure we still see the stake deposit delta, until it can be unstaked.
                 */
+
+                if self.validator_lamport_balances[index] == LAMPORT_BALANCE_DEFAULT {
+                    self.validator_lamport_balances[index] = current_lamports;
+                }
+
                 self.validator_lamport_balances[index] = match (
                     current_lamports < self.validator_lamport_balances[index],
                     current_lamports < target_lamports,
@@ -947,6 +946,7 @@ impl StewardState {
                     self,
                     index,
                     unstake_state,
+                    current_lamports,
                     stake_pool_lamports,
                     validator_list,
                     minimum_delegation,
@@ -967,15 +967,6 @@ impl StewardState {
                 RebalanceType::None
             };
 
-            msg!("Reserve lamports after adjustment: {}", reserve_lamports);
-            msg!(
-                "Stake pool lamports after adjustment: {}",
-                stake_pool_lamports
-            );
-            msg!("Rebalance Type: {:?}", rebalance);
-            msg!("Current Lamports: {}", current_lamports);
-            msg!("Target Lamports: {}", target_lamports);
-
             // Update internal state based on rebalance
             match rebalance {
                 RebalanceType::Decrease(DecreaseComponents {
@@ -984,8 +975,11 @@ impl StewardState {
                     stake_deposit_unstake_lamports,
                     total_unstake_lamports,
                 }) => {
-                    self.validator_lamport_balances[index] = self.validator_lamport_balances[index]
-                        .saturating_sub(total_unstake_lamports);
+                    if self.validator_lamport_balances[index] != LAMPORT_BALANCE_DEFAULT {
+                        self.validator_lamport_balances[index] = self.validator_lamport_balances
+                            [index]
+                            .saturating_sub(total_unstake_lamports);
+                    }
 
                     self.scoring_unstake_total = self
                         .scoring_unstake_total
@@ -1026,9 +1020,12 @@ impl StewardState {
                     }
                 }
                 RebalanceType::Increase(amount) => {
-                    self.validator_lamport_balances[index] = self.validator_lamport_balances[index]
-                        .checked_add(amount)
-                        .ok_or(StewardError::ArithmeticError)?;
+                    if self.validator_lamport_balances[index] != LAMPORT_BALANCE_DEFAULT {
+                        self.validator_lamport_balances[index] = self.validator_lamport_balances
+                            [index]
+                            .checked_add(amount)
+                            .ok_or(StewardError::ArithmeticError)?;
+                    }
                 }
                 RebalanceType::None => {}
             }