stake-pool: Fixup accounting on increase transient account (#1985)

stake-pool/cli/src/main.rs

@@ -105,9 +105,10 @@ fn command_create_pool(
     max_validators: u32,
     stake_pool_keypair: Option<Keypair>,
     mint_keypair: Option<Keypair>,
+    reserve_keypair: Option<Keypair>,
 ) -> CommandResult {
-    let reserve_stake = Keypair::new();
+    let reserve_keypair = reserve_keypair.unwrap_or_else(Keypair::new);
-    println!("Creating reserve stake {}", reserve_stake.pubkey());
+    println!("Creating reserve stake {}", reserve_keypair.pubkey());
 
     let mint_keypair = mint_keypair.unwrap_or_else(Keypair::new);
     println!("Creating mint {}", mint_keypair.pubkey());
@@ -163,13 +164,13 @@ fn command_create_pool(
             // Account for the stake pool reserve
             system_instruction::create_account(
                 &config.fee_payer.pubkey(),
-                &reserve_stake.pubkey(),
+                &reserve_keypair.pubkey(),
                 reserve_stake_balance,
                 STAKE_STATE_LEN as u64,
                 &stake_program::id(),
             ),
             stake_program::initialize(
-                &reserve_stake.pubkey(),
+                &reserve_keypair.pubkey(),
                 &stake_program::Authorized {
                     staker: withdraw_authority,
                     withdrawer: withdraw_authority,
@@ -236,7 +237,7 @@ fn command_create_pool(
                 &config.manager.pubkey(),
                 &config.staker.pubkey(),
                 &validator_list.pubkey(),
-                &reserve_stake.pubkey(),
+                &reserve_keypair.pubkey(),
                 &mint_keypair.pubkey(),
                 &pool_fee_account.pubkey(),
                 &spl_token::id(),
@@ -259,7 +260,7 @@ fn command_create_pool(
         config.fee_payer.as_ref(),
         &mint_keypair,
         &pool_fee_account,
-        &reserve_stake,
+        &reserve_keypair,
     ];
     unique_signers!(setup_signers);
     setup_transaction.sign(&setup_signers, recent_blockhash);
@@ -640,6 +641,13 @@ fn command_list(config: &Config, stake_pool_address: &Pubkey) -> CommandResult {
     let pool_mint = get_token_mint(&config.rpc_client, &stake_pool.pool_mint)?;
     let epoch_info = config.rpc_client.get_epoch_info()?;
 
+    let reserve_stake = config.rpc_client.get_account(&stake_pool.reserve_stake)?;
+    println!(
+        "Reserve Account: {}\tBalance: {}",
+        stake_pool.reserve_stake,
+        Sol(reserve_stake.lamports),
+    );
+
     for validator in validator_list.validators {
         println!(
             "Validator Vote Account: {}\tBalance: {}\tLast Update Epoch: {}{}",
@@ -1241,6 +1249,14 @@ fn main() {
                     .takes_value(true)
                     .help("Stake pool mint keypair [default: new keypair]"),
             )
+            .arg(
+                Arg::with_name("reserve_keypair")
+                    .long("reserve-keypair")
+                    .validator(is_keypair_or_ask_keyword)
+                    .value_name("PATH")
+                    .takes_value(true)
+                    .help("Stake pool reserve keypair [default: new keypair]"),
+            )
         )
         .subcommand(SubCommand::with_name("create-validator-stake")
             .about("Create a new stake account to use with the pool. Must be signed by the pool staker.")
@@ -1710,6 +1726,7 @@ fn main() {
             let max_validators = value_t_or_exit!(arg_matches, "max_validators", u32);
             let pool_keypair = keypair_of(arg_matches, "pool_keypair");
             let mint_keypair = keypair_of(arg_matches, "mint_keypair");
+            let reserve_keypair = keypair_of(arg_matches, "reserve_keypair");
             command_create_pool(
                 &config,
                 deposit_authority,
@@ -1720,6 +1737,7 @@ fn main() {
                 max_validators,
                 pool_keypair,
                 mint_keypair,
+                reserve_keypair,
             )
         }
         ("create-validator-stake", Some(arg_matches)) => {

stake-pool/program/src/instruction.rs

@@ -162,7 +162,11 @@ pub enum StakePoolInstruction {
     /// 10. `[]` Stake Config sysvar
     /// 11. `[]` System program
     /// 12. `[]` Stake program
-    ///  userdata: amount of lamports to split into the transient stake account
+    ///  userdata: amount of lamports to increase on the given validator.
+    ///  The actual amount split into the transient stake account is:
+    ///  `lamports + stake_rent_exemption`
+    ///  The rent-exemption of the stake account is withdrawn back to the reserve
+    ///  after it is merged.
     IncreaseValidatorStake(u64),
 
     /// (Staker only) Set the preferred deposit or withdraw stake account for the

stake-pool/program/src/lib.rs

@@ -33,7 +33,7 @@ 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;
+pub const MAX_VALIDATORS_TO_UPDATE: usize = 9;
 
 /// Get the stake amount under consideration when calculating pool token
 /// conversions

stake-pool/program/src/processor.rs

@@ -1150,19 +1150,21 @@ impl Processor {
         }
 
         let stake_rent = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
-        let minimum_lamports = MINIMUM_ACTIVE_STAKE + stake_rent;
+        if lamports < MINIMUM_ACTIVE_STAKE {
-        if lamports < minimum_lamports {
             msg!(
                 "Need more than {} lamports for transient stake to be rent-exempt and mergeable, {} provided",
-                minimum_lamports,
+                MINIMUM_ACTIVE_STAKE,
                 lamports
             );
             return Err(ProgramError::AccountNotRentExempt);
         }
 
+        // the stake account rent exemption is withdrawn after the merge, so
+        let total_lamports = lamports.saturating_add(stake_rent);
+
         if reserve_stake_account_info
             .lamports()
-            .saturating_sub(lamports)
+            .saturating_sub(total_lamports)
             <= stake_rent
         {
             let max_split_amount = reserve_stake_account_info
@@ -1189,7 +1191,7 @@ impl Processor {
             withdraw_authority_info.clone(),
             AUTHORITY_WITHDRAW,
             stake_pool.withdraw_bump_seed,
-            lamports,
+            total_lamports,
             transient_stake_account_info.clone(),
         )?;
 
@@ -1206,7 +1208,7 @@ impl Processor {
             stake_pool.withdraw_bump_seed,
         )?;
 
-        validator_list_entry.transient_stake_lamports = lamports;
+        validator_list_entry.transient_stake_lamports = total_lamports;
         validator_list.serialize(&mut *validator_list_info.data.borrow_mut())?;
 
         Ok(())
@@ -1404,6 +1406,9 @@ impl Processor {
                             if stake_program::active_stakes_can_merge(&stake, &validator_stake)
                                 .is_ok()
                             {
+                                let additional_lamports = transient_stake_info
+                                    .lamports()
+                                    .saturating_sub(stake.delegation.stake);
                                 Self::stake_merge(
                                     stake_pool_info.key,
                                     transient_stake_info.clone(),
@@ -1415,6 +1420,23 @@ impl Processor {
                                     stake_history_info.clone(),
                                     stake_program_info.clone(),
                                 )?;
+
+                                // post merge of two active stakes, withdraw
+                                // the extra back to the reserve
+                                if additional_lamports > 0 {
+                                    Self::stake_withdraw(
+                                        stake_pool_info.key,
+                                        validator_stake_info.clone(),
+                                        withdraw_authority_info.clone(),
+                                        AUTHORITY_WITHDRAW,
+                                        stake_pool.withdraw_bump_seed,
+                                        reserve_stake_info.clone(),
+                                        clock_info.clone(),
+                                        stake_history_info.clone(),
+                                        stake_program_info.clone(),
+                                        additional_lamports,
+                                    )?;
+                                }
                             } else {
                                 msg!("Stake activating or just active, not ready to merge");
                                 transient_stake_lamports = account_stake;
@@ -1436,6 +1458,10 @@ impl Processor {
             // Status for validator stake
             //  * active -> do everything
             //  * any other state / not a stake -> error state, but account for transient stake
+            let validator_stake_state = try_from_slice_unchecked::<stake_program::StakeState>(
+                &validator_stake_info.data.borrow(),
+            )
+            .ok();
             match validator_stake_state {
                 Some(stake_program::StakeState::Stake(_, stake)) => {
                     if validator_stake_record.status == StakeStatus::Active {
@@ -1782,8 +1808,8 @@ impl Processor {
     /// Processes [Withdraw](enum.Instruction.html).
     fn process_withdraw(
         program_id: &Pubkey,
-        pool_tokens: u64,
         accounts: &[AccountInfo],
+        pool_tokens: u64,
     ) -> ProgramResult {
         let account_info_iter = &mut accounts.iter();
         let stake_pool_info = next_account_info(account_info_iter)?;
@@ -2126,7 +2152,7 @@ impl Processor {
             }
             StakePoolInstruction::Withdraw(amount) => {
                 msg!("Instruction: Withdraw");
-                Self::process_withdraw(program_id, amount, accounts)
+                Self::process_withdraw(program_id, accounts, amount)
             }
             StakePoolInstruction::SetManager => {
                 msg!("Instruction: SetManager");

stake-pool/program/tests/increase.rs

@@ -93,8 +93,8 @@ async fn success() {
     assert!(transient_account.is_none());
 
     let rent = banks_client.get_rent().await.unwrap();
-    let lamports = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
+    let stake_rent = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
-    let reserve_lamports = reserve_lamports - lamports;
+    let increase_amount = reserve_lamports - stake_rent - 1;
     let error = stake_pool_accounts
         .increase_validator_stake(
             &mut banks_client,
@@ -102,7 +102,7 @@ async fn success() {
             &recent_blockhash,
             &validator_stake.transient_stake_account,
             &validator_stake.vote.pubkey(),
-            reserve_lamports,
+            increase_amount,
         )
         .await;
     assert!(error.is_none());
@@ -116,7 +116,7 @@ async fn success() {
     let reserve_stake_state =
         deserialize::<stake_program::StakeState>(&reserve_stake_account.data).unwrap();
     assert_eq!(
-        pre_reserve_stake_account.lamports - reserve_lamports,
+        pre_reserve_stake_account.lamports - increase_amount - stake_rent,
         reserve_stake_account.lamports
     );
     assert!(reserve_stake_state.delegation().is_none());
@@ -126,7 +126,10 @@ async fn success() {
         get_account(&mut banks_client, &validator_stake.transient_stake_account).await;
     let transient_stake_state =
         deserialize::<stake_program::StakeState>(&transient_stake_account.data).unwrap();
-    assert_eq!(transient_stake_account.lamports, reserve_lamports);
+    assert_eq!(
+        transient_stake_account.lamports,
+        increase_amount + stake_rent
+    );
     assert_ne!(
         transient_stake_state.delegation().unwrap().activation_epoch,
         Epoch::MAX
@@ -332,7 +335,7 @@ async fn fail_with_small_lamport_amount() {
     ) = setup().await;
 
     let rent = banks_client.get_rent().await.unwrap();
-    let lamports = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
+    let stake_rent = rent.minimum_balance(std::mem::size_of::<stake_program::StakeState>());
 
     let error = stake_pool_accounts
         .increase_validator_stake(
@@ -341,7 +344,7 @@ async fn fail_with_small_lamport_amount() {
             &recent_blockhash,
             &validator_stake.transient_stake_account,
             &validator_stake.vote.pubkey(),
-            lamports,
+            stake_rent,
         )
         .await
         .unwrap()

stake-pool/program/tests/update_validator_list_balance.rs

@@ -381,14 +381,6 @@ async fn merge_into_validator_stake() {
     let stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool_info.data).unwrap();
     assert_eq!(expected_lamports, stake_pool.total_stake_lamports);
 
-    let stake_pool_info = get_account(
-        &mut context.banks_client,
-        &stake_pool_accounts.stake_pool.pubkey(),
-    )
-    .await;
-    let stake_pool = try_from_slice_unchecked::<StakePool>(&stake_pool_info.data).unwrap();
-    assert_eq!(expected_lamports, stake_pool.total_stake_lamports);
-
     // Warp one more epoch so the stakes activate, ready to merge
     let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
     slot += slots_per_epoch;