KEEPER: Update Vault (#72)

+ Calls `update_vault_balance` after we distribute to vault
+ Does a full update vault loop if a vault needs to be updated

cli/src/getters.rs

@@ -26,8 +26,9 @@ use jito_tip_router_core::{
     weight_table::WeightTable,
 };
 use jito_vault_core::{
-    vault::Vault, vault_ncn_ticket::VaultNcnTicket,
+    config::Config as VaultConfig, vault::Vault, vault_ncn_ticket::VaultNcnTicket,
     vault_operator_delegation::VaultOperatorDelegation,
+    vault_update_state_tracker::VaultUpdateStateTracker,
 };
 use solana_account_decoder::{UiAccountEncoding, UiDataSliceConfig};
 use solana_client::{
@@ -418,6 +419,19 @@ pub async fn get_restaking_config(handler: &CliHandler) -> Result<RestakingConfi
     Ok(*account)
 }
 
+pub async fn get_vault_config(handler: &CliHandler) -> Result<VaultConfig> {
+    let (address, _, _) = VaultConfig::find_program_address(&handler.vault_program_id);
+    let account = get_account(handler, &address).await?;
+
+    if account.is_none() {
+        return Err(anyhow::anyhow!("Account not found"));
+    }
+    let account = account.unwrap();
+
+    let account = VaultConfig::try_from_slice_unchecked(account.data.as_slice())?;
+    Ok(*account)
+}
+
 pub async fn get_ncn(handler: &CliHandler) -> Result<Ncn> {
     let account = get_account(handler, handler.ncn()?).await?;
 
@@ -438,6 +452,21 @@ pub async fn get_vault(handler: &CliHandler, vault: &Pubkey) -> Result<Vault> {
     Ok(*account)
 }
 
+pub async fn get_vault_update_state_tracker(
+    handler: &CliHandler,
+    vault: &Pubkey,
+    ncn_epoch: u64,
+) -> Result<VaultUpdateStateTracker> {
+    let (vault_update_state_tracker, _, _) =
+        VaultUpdateStateTracker::find_program_address(&handler.vault_program_id, vault, ncn_epoch);
+
+    let account = get_account(handler, &vault_update_state_tracker)
+        .await?
+        .expect("Account not found");
+    let account = VaultUpdateStateTracker::try_from_slice_unchecked(account.data.as_slice())?;
+    Ok(*account)
+}
+
 pub async fn get_operator(handler: &CliHandler, operator: &Pubkey) -> Result<Operator> {
     let account = get_account(handler, operator).await?;
 
@@ -584,6 +613,68 @@ pub async fn get_stake_pool_accounts(handler: &CliHandler) -> Result<StakePoolAc
     Ok(accounts)
 }
 
+pub async fn get_all_sorted_operators_for_vault(
+    handler: &CliHandler,
+    vault: &Pubkey,
+) -> Result<Vec<Pubkey>> {
+    let client = handler.rpc_client();
+
+    let vault_operator_delegation_size = size_of::<VaultOperatorDelegation>() + 8;
+
+    let size_filter = RpcFilterType::DataSize(vault_operator_delegation_size as u64);
+
+    let vault_filter = RpcFilterType::Memcmp(Memcmp::new(
+        8,                                                  // offset
+        MemcmpEncodedBytes::Bytes(vault.to_bytes().into()), // encoded bytes
+    ));
+
+    let config = RpcProgramAccountsConfig {
+        filters: Some(vec![size_filter, vault_filter]),
+        account_config: RpcAccountInfoConfig {
+            encoding: Some(UiAccountEncoding::Base64),
+            data_slice: Some(UiDataSliceConfig {
+                offset: 0,
+                length: vault_operator_delegation_size,
+            }),
+            commitment: Some(handler.commitment),
+            min_context_slot: None,
+        },
+        with_context: Some(false),
+    };
+
+    let results = client
+        .get_program_accounts_with_config(&handler.vault_program_id, config)
+        .await?;
+
+    let accounts: Vec<(Pubkey, VaultOperatorDelegation)> = results
+        .iter()
+        .filter_map(|result| {
+            VaultOperatorDelegation::try_from_slice_unchecked(result.1.data.as_slice())
+                .map(|account| (result.0, *account))
+                .ok()
+        })
+        .collect();
+
+    let mut index_and_operator = accounts
+        .iter()
+        .map(|(_, vault_operator_delegation)| {
+            (
+                vault_operator_delegation.index(),
+                vault_operator_delegation.operator,
+            )
+        })
+        .collect::<Vec<(u64, Pubkey)>>();
+
+    index_and_operator.sort_by_cached_key(|(index, _pubkey)| *index);
+
+    let sorted_operators = index_and_operator
+        .iter()
+        .map(|(_index, pubkey)| *pubkey)
+        .collect::<Vec<Pubkey>>();
+
+    Ok(sorted_operators)
+}
+
 pub async fn get_all_operators_in_ncn(handler: &CliHandler) -> Result<Vec<Pubkey>> {
     let client = handler.rpc_client();
 

cli/src/instructions.rs

@@ -2,11 +2,12 @@ use std::{str::FromStr, time::Duration};
 
 use crate::{
     getters::{
-        get_account, get_all_operators_in_ncn, get_all_vaults_in_ncn, get_ballot_box,
-        get_base_reward_receiver_rewards, get_base_reward_router, get_epoch_snapshot,
+        get_account, get_all_operators_in_ncn, get_all_sorted_operators_for_vault,
+        get_all_vaults_in_ncn, get_ballot_box, get_base_reward_receiver_rewards,
+        get_base_reward_router, get_current_slot, get_epoch_snapshot,
         get_ncn_reward_receiver_rewards, get_ncn_reward_router, get_operator,
         get_operator_snapshot, get_stake_pool_accounts, get_tip_router_config, get_vault,
-        get_vault_registry, get_weight_table,
+        get_vault_config, get_vault_registry, get_vault_update_state_tracker, get_weight_table,
     },
     handler::CliHandler,
     log::boring_progress_bar,
@@ -56,13 +57,20 @@ use jito_tip_router_core::{
     vault_registry::VaultRegistry,
     weight_table::WeightTable,
 };
-use jito_vault_client::instructions::{
-    AddDelegationBuilder, InitializeVaultBuilder, InitializeVaultNcnTicketBuilder,
-    InitializeVaultOperatorDelegationBuilder, MintToBuilder, WarmupVaultNcnTicketBuilder,
+use jito_vault_client::{
+    instructions::{
+        AddDelegationBuilder, CloseVaultUpdateStateTrackerBuilder,
+        CrankVaultUpdateStateTrackerBuilder, InitializeVaultBuilder,
+        InitializeVaultNcnTicketBuilder, InitializeVaultOperatorDelegationBuilder,
+        InitializeVaultUpdateStateTrackerBuilder, MintToBuilder, UpdateVaultBalanceBuilder,
+        WarmupVaultNcnTicketBuilder,
+    },
+    types::WithdrawalAllocationMethod,
 };
 use jito_vault_core::{
     config::Config as VaultConfig, vault::Vault, vault_ncn_ticket::VaultNcnTicket,
     vault_operator_delegation::VaultOperatorDelegation,
+    vault_update_state_tracker::VaultUpdateStateTracker,
 };
 use log::info;
 use solana_client::rpc_config::RpcSendTransactionConfig;
@@ -1036,6 +1044,7 @@ pub async fn create_operator_snapshot(
 
     Ok(())
 }
+
 pub async fn snapshot_vault_operator_delegation(
     handler: &CliHandler,
     vault: &Pubkey,
@@ -1766,9 +1775,33 @@ pub async fn distribute_ncn_vault_rewards(
         .epoch(epoch)
         .instruction();
 
+    let vault_account = get_vault(handler, &vault).await?;
+    let st_mint = vault_account.supported_mint;
+    let vrt_mint = vault_account.vrt_mint;
+    let vault_fee_wallet = vault_account.fee_wallet;
+
+    let vault_fee_token_account = get_associated_token_address(&vault_fee_wallet, &vrt_mint);
+    let vault_token_account = get_associated_token_address(&vault, &st_mint);
+
+    let (vault_config, _, _) = VaultConfig::find_program_address(&handler.vault_program_id);
+
+    let update_vault_balance_ix = UpdateVaultBalanceBuilder::new()
+        .config(vault_config)
+        .vault(vault)
+        .token_program(spl_token::id())
+        .vault_fee_token_account(vault_fee_token_account)
+        .vault_token_account(vault_token_account)
+        .vrt_mint(vrt_mint)
+        .instruction();
+
     send_and_log_transaction(
         handler,
-        &[create_vault_ata_ix, distribute_ncn_vault_rewards_ix],
+        &[
+            ComputeBudgetInstruction::set_compute_unit_limit(1_400_000),
+            create_vault_ata_ix,
+            distribute_ncn_vault_rewards_ix,
+            update_vault_balance_ix,
+        ],
         &[],
         "Distributed NCN Vault Rewards",
         &[
@@ -1964,6 +1997,163 @@ pub async fn check_created(handler: &CliHandler, address: &Pubkey) -> Result<()>
     Ok(())
 }
 
+pub async fn full_vault_update(handler: &CliHandler, vault: &Pubkey) -> Result<()> {
+    let payer = handler.keypair()?;
+
+    // Get Epoch Info
+    let current_slot = get_current_slot(handler).await?;
+    let (ncn_epoch, epoch_length) = {
+        let vault_config = get_vault_config(handler).await?;
+        let ncn_epoch = vault_config.get_epoch_from_slot(current_slot)?;
+        let epoch_length = vault_config.epoch_length();
+        (ncn_epoch, epoch_length)
+    };
+
+    // Check Vault
+    let vault_account = get_vault(handler, vault).await?;
+    let current_slot = get_current_slot(handler).await?;
+
+    if !vault_account.is_update_needed(current_slot, epoch_length)? {
+        return Ok(());
+    }
+
+    // Initialize Vault Update State Tracker
+    let (vault_config, _, _) = VaultConfig::find_program_address(&handler.vault_program_id);
+
+    let (vault_update_state_tracker, _, _) =
+        VaultUpdateStateTracker::find_program_address(&handler.vault_program_id, vault, ncn_epoch);
+
+    let vault_update_state_tracker_account =
+        get_account(handler, &vault_update_state_tracker).await?;
+
+    if vault_update_state_tracker_account.is_none() {
+        let initialize_vault_update_state_tracker_ix =
+            InitializeVaultUpdateStateTrackerBuilder::new()
+                .vault(*vault)
+                .vault_update_state_tracker(vault_update_state_tracker)
+                .system_program(system_program::id())
+                .withdrawal_allocation_method(WithdrawalAllocationMethod::Greedy)
+                .payer(payer.pubkey())
+                .config(vault_config)
+                .instruction();
+
+        let result = send_and_log_transaction(
+            handler,
+            &[initialize_vault_update_state_tracker_ix],
+            &[payer],
+            "Initialize Vault Update State Tracker",
+            &[
+                format!("VAULT: {:?}", vault),
+                format!("Vault Epoch: {:?}", ncn_epoch),
+            ],
+        )
+        .await;
+
+        if result.is_err() {
+            log::error!(
+                "Failed to initialize Vault Update State Tracker for Vault: {:?} at NCN Epoch: {:?} with error: {:?}",
+                vault,
+                ncn_epoch,
+                result.err().unwrap()
+            );
+        }
+    }
+
+    // Crank Vault Update State Tracker
+    let all_operators = get_all_sorted_operators_for_vault(handler, vault).await?;
+
+    let starting_index = {
+        let vault_update_state_tracker_account =
+            get_vault_update_state_tracker(handler, vault, ncn_epoch).await?;
+        let last_updated_index = vault_update_state_tracker_account.last_updated_index();
+
+        if last_updated_index == u64::MAX {
+            ncn_epoch % all_operators.len() as u64
+        } else {
+            (last_updated_index + 1) % all_operators.len() as u64
+        }
+    };
+
+    for index in 0..all_operators.len() {
+        let current_index = (starting_index as usize + index) % all_operators.len();
+        let operator = all_operators.get(current_index).unwrap();
+
+        let (vault_operator_delegation, _, _) = VaultOperatorDelegation::find_program_address(
+            &handler.vault_program_id,
+            vault,
+            operator,
+        );
+
+        let crank_vault_update_state_tracker_ix = CrankVaultUpdateStateTrackerBuilder::new()
+            .vault(*vault)
+            .operator(*operator)
+            .config(vault_config)
+            .vault_operator_delegation(vault_operator_delegation)
+            .vault_update_state_tracker(vault_update_state_tracker)
+            .instruction();
+
+        let result = send_and_log_transaction(
+            handler,
+            &[crank_vault_update_state_tracker_ix],
+            &[payer],
+            "Crank Vault Update State Tracker",
+            &[
+                format!("VAULT: {:?}", vault),
+                format!("Operator: {:?}", operator),
+                format!("Vault Epoch: {:?}", ncn_epoch),
+            ],
+        )
+        .await;
+
+        if result.is_err() {
+            log::error!(
+                "Failed to crank Vault Update State Tracker for Vault: {:?} and Operator: {:?} at NCN Epoch: {:?} with error: {:?}",
+                vault,
+                operator,
+                ncn_epoch,
+                result.err().unwrap()
+            );
+        }
+    }
+
+    // Close Update State Tracker
+    let vault_update_state_tracker_account =
+        get_account(handler, &vault_update_state_tracker).await?;
+
+    if vault_update_state_tracker_account.is_some() {
+        let close_vault_update_state_tracker_ix = CloseVaultUpdateStateTrackerBuilder::new()
+            .vault(*vault)
+            .vault_update_state_tracker(vault_update_state_tracker)
+            .payer(payer.pubkey())
+            .config(vault_config)
+            .ncn_epoch(ncn_epoch)
+            .instruction();
+
+        let result = send_and_log_transaction(
+            handler,
+            &[close_vault_update_state_tracker_ix],
+            &[payer],
+            "Close Vault Update State Tracker",
+            &[
+                format!("VAULT: {:?}", vault),
+                format!("Vault Epoch: {:?}", ncn_epoch),
+            ],
+        )
+        .await;
+
+        if result.is_err() {
+            log::error!(
+                "Failed to close Vault Update State Tracker for Vault: {:?} at NCN Epoch: {:?} with error: {:?}",
+                vault,
+                ncn_epoch,
+                result.err().unwrap()
+            );
+        }
+    }
+
+    Ok(())
+}
+
 pub async fn get_or_create_weight_table(handler: &CliHandler, epoch: u64) -> Result<WeightTable> {
     let ncn = *handler.ncn()?;
 
@@ -2185,6 +2375,16 @@ pub async fn crank_snapshot(handler: &CliHandler, epoch: u64) -> Result<()> {
             .collect();
 
         for vault in vaults_to_run.iter() {
+            let result = full_vault_update(handler, vault).await;
+
+            if let Err(err) = result {
+                log::error!(
+                    "Failed to update the vault: {:?} with error: {:?}",
+                    vault,
+                    err
+                );
+            }
+
             let result = snapshot_vault_operator_delegation(handler, vault, operator, epoch).await;
 
             if let Err(err) = result {