- reimplement prepare_withdraw_accounts according to the documentation (

solana-labs#2513)

* - reimplement prepare_withdraw_accounts according to the docu

* - add preferred withdraw validator

* - fix clippy issues

* - fix styling

stake-pool/cli/src/client.rs

@@ -77,43 +77,6 @@ pub(crate) fn get_stake_state(
     Ok(stake_state)
 }
 
-pub(crate) fn get_stake_accounts_by_withdraw_authority(
-    rpc_client: &RpcClient,
-    withdraw_authority: &Pubkey,
-) -> Result<Vec<(Pubkey, u64, stake_program::StakeState)>, ClientError> {
-    rpc_client
-        .get_program_accounts_with_config(
-            &stake_program::id(),
-            #[allow(clippy::needless_update)] // TODO: Remove after updating to solana >=1.6.10
-            RpcProgramAccountsConfig {
-                filters: Some(vec![RpcFilterType::Memcmp(Memcmp {
-                    offset: 44, // 44 is Withdrawer authority offset in stake account stake
-                    bytes: MemcmpEncodedBytes::Binary(format!("{}", withdraw_authority)),
-                    encoding: None,
-                })]),
-                account_config: RpcAccountInfoConfig {
-                    encoding: Some(UiAccountEncoding::Base64),
-                    ..RpcAccountInfoConfig::default()
-                },
-                ..RpcProgramAccountsConfig::default()
-            },
-        )
-        .map(|accounts| {
-            accounts
-                .into_iter()
-                .filter_map(
-                    |(address, account)| match deserialize(account.data.as_slice()) {
-                        Ok(stake_state) => Some((address, account.lamports, stake_state)),
-                        Err(err) => {
-                            eprintln!("Invalid stake account data for {}: {}", address, err);
-                            None
-                        }
-                    },
-                )
-                .collect()
-        })
-}
-
 pub(crate) fn get_stake_pools(
     rpc_client: &RpcClient,
 ) -> Result<Vec<(Pubkey, StakePool, ValidatorList)>, ClientError> {

stake-pool/cli/src/main.rs

@@ -31,6 +31,7 @@ use {
         transaction::Transaction,
     },
     spl_associated_token_account::{create_associated_token_account, get_associated_token_address},
+    spl_stake_pool::state::ValidatorStakeInfo,
     spl_stake_pool::{
         self, find_stake_program_address, find_transient_stake_program_address,
         find_withdraw_authority_program_address,
@@ -39,6 +40,7 @@ use {
         state::{Fee, FeeType, StakePool, ValidatorList},
         MINIMUM_ACTIVE_STAKE,
     },
+    std::cmp::Ordering,
     std::{process::exit, sync::Arc},
 };
 
@@ -1061,31 +1063,93 @@ struct WithdrawAccount {
     pool_amount: u64,
 }
 
+fn sorted_accounts<F>(
+    validator_list: &ValidatorList,
+    stake_pool: &StakePool,
+    get_pubkey: F,
+) -> Vec<(Pubkey, u64, Option<Pubkey>)>
+where
+    F: Fn(&ValidatorStakeInfo) -> Pubkey,
+{
+    let mut result: Vec<(Pubkey, u64, Option<Pubkey>)> = validator_list
+        .validators
+        .iter()
+        .map(|validator| {
+            (
+                get_pubkey(validator),
+                validator.active_stake_lamports,
+                Some(validator.vote_account_address),
+            )
+        })
+        .collect::<Vec<_>>();
+
+    result.sort_by(|left, right| {
+        if left.2 == stake_pool.preferred_withdraw_validator_vote_address {
+            Ordering::Less
+        } else if right.2 == stake_pool.preferred_withdraw_validator_vote_address {
+            Ordering::Greater
+        } else {
+            right.1.cmp(&left.1)
+        }
+    });
+
+    result
+}
+
 fn prepare_withdraw_accounts(
     rpc_client: &RpcClient,
     stake_pool: &StakePool,
-    pool_withdraw_authority: &Pubkey,
     pool_amount: u64,
+    stake_pool_address: &Pubkey,
 ) -> Result<Vec<WithdrawAccount>, Error> {
-    let mut accounts =
-        get_stake_accounts_by_withdraw_authority(rpc_client, pool_withdraw_authority)?;
-    if accounts.is_empty() {
-        return Err("No accounts found.".to_string().into());
-    }
     let min_balance = rpc_client
         .get_minimum_balance_for_rent_exemption(STAKE_STATE_LEN)?
         .saturating_add(MINIMUM_ACTIVE_STAKE);
     let pool_mint = get_token_mint(rpc_client, &stake_pool.pool_mint)?;
 
-    // Sort from highest to lowest balance
-    accounts.sort_by(|a, b| b.1.cmp(&a.1));
+    let validator_list: ValidatorList = get_validator_list(rpc_client, &stake_pool.validator_list)?;
+
+    let mut accounts: Vec<(Pubkey, u64, Option<Pubkey>)> = Vec::new();
+
+    accounts.append(&mut sorted_accounts(
+        &validator_list,
+        stake_pool,
+        |validator| {
+            let (stake_account_address, _) = find_stake_program_address(
+                &spl_stake_pool::id(),
+                &validator.vote_account_address,
+                stake_pool_address,
+            );
+
+            stake_account_address
+        },
+    ));
+
+    accounts.append(&mut sorted_accounts(
+        &validator_list,
+        stake_pool,
+        |validator| {
+            let (transient_stake_account_address, _) = find_transient_stake_program_address(
+                &spl_stake_pool::id(),
+                &validator.vote_account_address,
+                stake_pool_address,
+                validator.transient_seed_suffix_start,
+            );
+
+            transient_stake_account_address
+        },
+    ));
+
+    let reserve_stake = rpc_client.get_account(&stake_pool.reserve_stake)?;
+
+    accounts.push((stake_pool.reserve_stake, reserve_stake.lamports, None));
 
     // Prepare the list of accounts to withdraw from
     let mut withdraw_from: Vec<WithdrawAccount> = vec![];
     let mut remaining_amount = pool_amount;
 
     // Go through available accounts and withdraw from largest to smallest
-    for (stake_address, lamports, stake) in accounts {
+    for (stake_address, lamports, vote_address_opt) in accounts {
         if lamports <= min_balance {
             continue;
         }
@@ -1099,7 +1163,7 @@ fn prepare_withdraw_accounts(
         // Those accounts will be withdrawn completely with `claim` instruction
         withdraw_from.push(WithdrawAccount {
             stake_address,
-            vote_address: stake.delegation().map(|x| x.voter_pubkey),
+            vote_address: vote_address_opt,
             pool_amount,
         });
         remaining_amount -= pool_amount;
@@ -1204,8 +1268,8 @@ fn command_withdraw_stake(
         prepare_withdraw_accounts(
             &config.rpc_client,
             &stake_pool,
-            &pool_withdraw_authority,
             pool_amount,
+            stake_pool_address,
         )?
     };