stake-pool: Add merging transient stakes in update (#1618)

* Add check for transient stake account activation on removal

* Add proper merging logic during update

* Format + clippy

* Add max possible validators

* Disallow removal for any transient stake state

* Reduce number of accounts for BPF instruction usage

Author: joncinque

stake-pool/cli/src/main.rs

@@ -29,9 +29,10 @@ use {
         self,
         borsh::get_instance_packed_len,
         find_deposit_authority_program_address, find_stake_program_address,
-        find_withdraw_authority_program_address,
+        find_transient_stake_program_address, find_withdraw_authority_program_address,
         stake_program::{self, StakeAuthorize, StakeState},
         state::{Fee, StakePool, ValidatorList},
+        MAX_VALIDATORS_TO_UPDATE,
     },
     std::process::exit,
 };
@@ -51,7 +52,6 @@ type Error = Box<dyn std::error::Error>;
 type CommandResult = Result<(), Error>;
 
 const STAKE_STATE_LEN: usize = 200;
-const MAX_ACCOUNTS_TO_UPDATE: usize = 10;
 lazy_static! {
     static ref MIN_STAKE_BALANCE: u64 = native_token::sol_to_lamports(1.0);
 }
@@ -347,7 +347,7 @@ fn command_vsa_add(config: &Config, stake_pool_address: &Pubkey, stake: &Pubkey)
 fn command_vsa_remove(
     config: &Config,
     stake_pool_address: &Pubkey,
-    stake: &Pubkey,
+    vote_account: &Pubkey,
     new_authority: &Option<Pubkey>,
 ) -> CommandResult {
     if !config.no_update {
@@ -357,6 +357,13 @@ fn command_vsa_remove(
     let stake_pool = get_stake_pool(&config.rpc_client, stake_pool_address)?;
     let pool_withdraw_authority =
         find_withdraw_authority_program_address(&spl_stake_pool::id(), stake_pool_address).0;
+    let (validator_stake_account, _) =
+        find_stake_program_address(&spl_stake_pool::id(), &vote_account, stake_pool_address);
+    let (transient_stake_account, _) = find_transient_stake_program_address(
+        &spl_stake_pool::id(),
+        &vote_account,
+        stake_pool_address,
+    );
 
     let staker_pubkey = config.staker.pubkey();
     let new_authority = new_authority.as_ref().unwrap_or(&staker_pubkey);
@@ -371,7 +378,8 @@ fn command_vsa_remove(
                 &pool_withdraw_authority,
                 &new_authority,
                 &stake_pool.validator_list,
-                &stake,
+                &validator_stake_account,
+                &transient_stake_account,
             )?,
         ],
         Some(&config.fee_payer.pubkey()),
@@ -640,27 +648,34 @@ fn command_update(config: &Config, stake_pool_address: &Pubkey) -> CommandResult
         .collect();
 
     println!("Updating stake pool...");
+    let (withdraw_authority, _) =
+        find_withdraw_authority_program_address(&spl_stake_pool::id(), &stake_pool_address);
 
     let mut instructions: Vec<Instruction> = vec![];
-    for accounts_chunk in accounts_to_update.chunks(MAX_ACCOUNTS_TO_UPDATE) {
+    let mut start_index = 0;
+    for accounts_chunk in accounts_to_update.chunks(MAX_VALIDATORS_TO_UPDATE) {
         instructions.push(spl_stake_pool::instruction::update_validator_list_balance(
             &spl_stake_pool::id(),
+            stake_pool_address,
+            &withdraw_authority,
             &stake_pool.validator_list,
+            &stake_pool.reserve_stake,
             &accounts_chunk,
-        )?);
+            start_index,
+            false,
+        ));
+        start_index += MAX_VALIDATORS_TO_UPDATE as u32;
     }
 
-    let (withdraw_authority, _) =
-        find_withdraw_authority_program_address(&spl_stake_pool::id(), &stake_pool_address);
-
     instructions.push(spl_stake_pool::instruction::update_stake_pool_balance(
         &spl_stake_pool::id(),
         stake_pool_address,
-        &stake_pool.validator_list,
         &withdraw_authority,
+        &stake_pool.validator_list,
+        &stake_pool.reserve_stake,
         &stake_pool.manager_fee_account,
         &stake_pool.pool_mint,
-    )?);
+    ));
 
     // TODO: A faster solution would be to send all the `update_validator_list_balance` instructions concurrently
     for instruction in instructions {
@@ -1135,15 +1150,6 @@ fn main() {
                     .required(true)
                     .help("Stake account to add to the pool"),
             )
-            .arg(
-                Arg::with_name("token_receiver")
-                    .long("token-receiver")
-                    .validator(is_pubkey)
-                    .value_name("ADDRESS")
-                    .takes_value(true)
-                    .help("Account to receive pool token. Must be initialized account of the stake pool token. \
-                          Defaults to the new pool token account."),
-            )
         )
         .subcommand(SubCommand::with_name("remove-validator")
             .about("Remove validator account from the stake pool. Must be signed by the pool staker.")
@@ -1157,23 +1163,13 @@ fn main() {
                     .help("Stake pool address"),
             )
             .arg(
-                Arg::with_name("stake_account")
+                Arg::with_name("vote_account")
                     .index(2)
                     .validator(is_pubkey)
-                    .value_name("STAKE_ACCOUNT_ADDRESS")
-                    .takes_value(true)
-                    .required(true)
-                    .help("Stake account to remove from the pool"),
-            )
-            .arg(
-                Arg::with_name("withdraw_from")
-                    .long("withdraw-from")
-                    .validator(is_pubkey)
-                    .value_name("ADDRESS")
+                    .value_name("VOTE_ACCOUNT_ADDRESS")
                     .takes_value(true)
                     .required(true)
-                    .help("Token account to withdraw pool token from. \
-                          Must have enough tokens for the full stake address balance."),
+                    .help("Vote account for the validator to remove from the pool"),
             )
             .arg(
                 Arg::with_name("new_authority")
@@ -1455,9 +1451,9 @@ fn main() {
         }
         ("remove-validator", Some(arg_matches)) => {
             let stake_pool_address = pubkey_of(arg_matches, "pool").unwrap();
-            let stake_account = pubkey_of(arg_matches, "stake_account").unwrap();
+            let vote_account = pubkey_of(arg_matches, "vote_account").unwrap();
             let new_authority: Option<Pubkey> = pubkey_of(arg_matches, "new_authority");
-            command_vsa_remove(&config, &stake_pool_address, &stake_account, &new_authority)
+            command_vsa_remove(&config, &stake_pool_address, &vote_account, &new_authority)
         }
         ("deposit", Some(arg_matches)) => {
             let stake_pool_address = pubkey_of(arg_matches, "pool").unwrap();

stake-pool/program/src/instruction.rs

@@ -3,7 +3,9 @@
 #![allow(clippy::too_many_arguments)]
 
 use {
-    crate::{stake_program, state::Fee},
+    crate::{
+        find_stake_program_address, find_transient_stake_program_address, stake_program, state::Fee,
+    },
     borsh::{BorshDeserialize, BorshSchema, BorshSerialize},
     solana_program::{
         instruction::{AccountMeta, Instruction},
@@ -88,8 +90,9 @@ pub enum StakePoolInstruction {
     ///   3. `[]` New withdraw/staker authority to set in the stake account
     ///   4. `[w]` Validator stake list storage account
     ///   5. `[w]` Stake account to remove from the pool
-    ///   8. '[]' Sysvar clock
-    ///  10. `[]` Stake program id,
+    ///   6. `[]` Transient stake account, to check that that we're not trying to activate
+    ///   7. '[]' Sysvar clock
+    ///   8. `[]` Stake program id,
     RemoveValidatorFromPool,
 
     /// (Staker only) Decrease active stake on a validator, eventually moving it to the reserve
@@ -132,7 +135,7 @@ pub enum StakePoolInstruction {
     ///  0. `[]` Stake pool
     ///  1. `[s]` Stake pool staker
     ///  2. `[]` Stake pool withdraw authority
-    ///  3. `[]` Validator list
+    ///  3. `[w]` Validator list
     ///  4. `[w]` Stake pool reserve stake
     ///  5. `[w]` Transient stake account
     ///  6. `[]` Validator vote account to delegate to
@@ -149,26 +152,36 @@ pub enum StakePoolInstruction {
     ///
     ///  While going through the pairs of validator and transient stake accounts,
     ///  if the transient stake is inactive, it is merged into the reserve stake
-    ///  account.  If the transient stake is active and has matching credits
+    ///  account. If the transient stake is active and has matching credits
     ///  observed, it is merged into the canonical validator stake account. In
     ///  all other states, nothing is done, and the balance is simply added to
     ///  the canonical stake account balance.
     ///
     ///  0. `[]` Stake pool
-    ///  1. `[w]` Validator stake list storage account
-    ///  2. `[w]` Reserve stake account
-    ///  3. `[]` Stake pool withdraw authority
-    ///  4. `[]` Sysvar clock account
-    ///  5. `[]` Stake program
-    ///  6. ..6+N ` [] N pairs of validator and transient stake accounts
-    UpdateValidatorListBalance,
+    ///  1. `[]` Stake pool withdraw authority
+    ///  2. `[w]` Validator stake list storage account
+    ///  3. `[w]` Reserve stake account
+    ///  4. `[]` Sysvar clock
+    ///  5. `[]` Sysvar stake history
+    ///  6. `[]` Stake program
+    ///  7. ..7+N ` [] N pairs of validator and transient stake accounts
+    UpdateValidatorListBalance {
+        /// Index to start updating on the validator list
+        #[allow(dead_code)] // but it's not
+        start_index: u32,
+        /// If true, don't try merging transient stake accounts into the reserve or
+        /// validator stake account.  Useful for testing or if a particular stake
+        /// account is in a bad state, but we still want to update
+        #[allow(dead_code)] // but it's not
+        no_merge: bool,
+    },
 
     ///   Updates total pool balance based on balances in the reserve and validator list
     ///
     ///   0. `[w]` Stake pool
-    ///   1. `[]` Validator stake list storage account
-    ///   2. `[]` Reserve stake account
-    ///   3. `[]` Stake pool withdraw authority
+    ///   1. `[]` Stake pool withdraw authority
+    ///   2. `[]` Validator stake list storage account
+    ///   3. `[]` Reserve stake account
     ///   4. `[w]` Account to receive pool fee tokens
     ///   5. `[w]` Pool mint account
     ///   6. `[]` Sysvar clock account
@@ -347,6 +360,7 @@ pub fn remove_validator_from_pool(
     new_stake_authority: &Pubkey,
     validator_list: &Pubkey,
     stake_account: &Pubkey,
+    transient_stake_account: &Pubkey,
 ) -> Result<Instruction, ProgramError> {
     let accounts = vec![
         AccountMeta::new(*stake_pool, false),
@@ -355,6 +369,7 @@ pub fn remove_validator_from_pool(
         AccountMeta::new_readonly(*new_stake_authority, false),
         AccountMeta::new(*validator_list, false),
         AccountMeta::new(*stake_account, false),
+        AccountMeta::new_readonly(*transient_stake_account, false),
         AccountMeta::new_readonly(sysvar::clock::id(), false),
         AccountMeta::new_readonly(stake_program::id(), false),
     ];
@@ -413,7 +428,7 @@ pub fn increase_validator_stake(
         AccountMeta::new_readonly(*stake_pool, false),
         AccountMeta::new_readonly(*staker, true),
         AccountMeta::new_readonly(*stake_pool_withdraw_authority, false),
-        AccountMeta::new_readonly(*validator_list, false),
+        AccountMeta::new(*validator_list, false),
         AccountMeta::new(*reserve_stake, false),
         AccountMeta::new(*transient_stake, false),
         AccountMeta::new_readonly(*validator, false),
@@ -434,45 +449,80 @@ pub fn increase_validator_stake(
 /// Creates `UpdateValidatorListBalance` instruction (update validator stake account balances)
 pub fn update_validator_list_balance(
     program_id: &Pubkey,
-    validator_list_storage: &Pubkey,
-    validator_list: &[Pubkey],
-) -> Result<Instruction, ProgramError> {
-    let mut accounts: Vec<AccountMeta> = validator_list
-        .iter()
-        .map(|pubkey| AccountMeta::new_readonly(*pubkey, false))
-        .collect();
-    accounts.insert(0, AccountMeta::new(*validator_list_storage, false));
-    accounts.insert(1, AccountMeta::new_readonly(sysvar::clock::id(), false));
-    Ok(Instruction {
+    stake_pool: &Pubkey,
+    stake_pool_withdraw_authority: &Pubkey,
+    validator_list: &Pubkey,
+    reserve_stake: &Pubkey,
+    validator_vote_accounts: &[Pubkey],
+    start_index: u32,
+    no_merge: bool,
+) -> Instruction {
+    let mut accounts = vec![
+        AccountMeta::new_readonly(*stake_pool, false),
+        AccountMeta::new_readonly(*stake_pool_withdraw_authority, false),
+        AccountMeta::new(*validator_list, false),
+        AccountMeta::new(*reserve_stake, false),
+        AccountMeta::new_readonly(sysvar::clock::id(), false),
+        AccountMeta::new_readonly(sysvar::stake_history::id(), false),
+        AccountMeta::new_readonly(stake_program::id(), false),
+    ];
+    accounts.append(
+        &mut validator_vote_accounts
+            .iter()
+            .flat_map(|vote_account_address| {
+                let (validator_stake_account, _) =
+                    find_stake_program_address(program_id, vote_account_address, stake_pool);
+                let (transient_stake_account, _) = find_transient_stake_program_address(
+                    program_id,
+                    vote_account_address,
+                    stake_pool,
+                );
+                vec![
+                    AccountMeta::new(validator_stake_account, false),
+                    AccountMeta::new(transient_stake_account, false),
+                ]
+            })
+            .collect::<Vec<AccountMeta>>(),
+    );
+    Instruction {
         program_id: *program_id,
         accounts,
-        data: StakePoolInstruction::UpdateValidatorListBalance.try_to_vec()?,
-    })
+        data: StakePoolInstruction::UpdateValidatorListBalance {
+            start_index,
+            no_merge,
+        }
+        .try_to_vec()
+        .unwrap(),
+    }
 }
 
 /// Creates `UpdateStakePoolBalance` instruction (pool balance from the stake account list balances)
 pub fn update_stake_pool_balance(
     program_id: &Pubkey,
     stake_pool: &Pubkey,
-    validator_list_storage: &Pubkey,
     withdraw_authority: &Pubkey,
+    validator_list_storage: &Pubkey,
+    reserve_stake: &Pubkey,
     manager_fee_account: &Pubkey,
     stake_pool_mint: &Pubkey,
-) -> Result<Instruction, ProgramError> {
+) -> Instruction {
     let accounts = vec![
         AccountMeta::new(*stake_pool, false),
-        AccountMeta::new(*validator_list_storage, false),
         AccountMeta::new_readonly(*withdraw_authority, false),
+        AccountMeta::new_readonly(*validator_list_storage, false),
+        AccountMeta::new_readonly(*reserve_stake, false),
         AccountMeta::new(*manager_fee_account, false),
         AccountMeta::new(*stake_pool_mint, false),
         AccountMeta::new_readonly(sysvar::clock::id(), false),
         AccountMeta::new_readonly(spl_token::id(), false),
     ];
-    Ok(Instruction {
+    Instruction {
         program_id: *program_id,
         accounts,
-        data: StakePoolInstruction::UpdateStakePoolBalance.try_to_vec()?,
-    })
+        data: StakePoolInstruction::UpdateStakePoolBalance
+            .try_to_vec()
+            .unwrap(),
+    }
 }
 
 /// Creates a 'Deposit' instruction.

stake-pool/program/src/lib.rs

@@ -32,13 +32,23 @@ const TRANSIENT_STAKE_SEED: &[u8] = b"transient";
 /// for merges without a mismatch on credits observed
 pub const MINIMUM_ACTIVE_STAKE: u64 = LAMPORTS_PER_SOL;
 
+/// Maximum amount of validator stake accounts to update per
+/// `UpdateValidatorListBalance` instruction, based on compute limits
+pub const MAX_VALIDATORS_TO_UPDATE: usize = 10;
+
 /// Get the stake amount under consideration when calculating pool token
 /// conversions
 pub fn minimum_stake_lamports(meta: &Meta) -> u64 {
     meta.rent_exempt_reserve
         .saturating_add(MINIMUM_ACTIVE_STAKE)
 }
 
+/// Get the stake amount under consideration when calculating pool token
+/// conversions
+pub fn minimum_reserve_lamports(meta: &Meta) -> u64 {
+    meta.rent_exempt_reserve.saturating_add(1)
+}
+
 /// Generates the deposit authority program address for the stake pool
 pub fn find_deposit_authority_program_address(
     program_id: &Pubkey,