Add wait for max stake command (solana-labs#13532)

cli/src/cli.rs

@@ -149,6 +149,9 @@ pub enum CliCommand {
         limit: usize,
         show_transactions: bool,
     },
+    WaitForMaxStake {
+        max_stake_percent: f32,
+    },
     // Nonce commands
     AuthorizeNonceAccount {
         nonce_account: Pubkey,
@@ -624,6 +627,13 @@ pub fn parse_command(
                 signers,
             })
         }
+        ("wait-for-max-stake", Some(matches)) => {
+            let max_stake_percent = value_t_or_exit!(matches, "max_percent", f32);
+            Ok(CliCommandInfo {
+                command: CliCommand::WaitForMaxStake { max_stake_percent },
+                signers: vec![],
+            })
+        }
         // Stake Commands
         ("create-stake-account", Some(matches)) => {
             parse_stake_create_account(matches, default_signer, wallet_manager)
@@ -1565,6 +1575,9 @@ pub fn process_command(config: &CliConfig) -> ProcessResult {
             *use_lamports_unit,
             vote_account_pubkeys.as_deref(),
         ),
+        CliCommand::WaitForMaxStake { max_stake_percent } => {
+            process_wait_for_max_stake(&rpc_client, config, *max_stake_percent)
+        }
         CliCommand::ShowValidators { use_lamports_unit } => {
             process_show_validators(&rpc_client, config, *use_lamports_unit)
         }

cli/src/cluster_query.rs

@@ -318,6 +318,17 @@ impl ClusterQuerySubCommands for App<'_, '_> {
                         .help("Display the full transactions"),
                 )
         )
+        .subcommand(
+            SubCommand::with_name("wait-for-max-stake")
+                .about("Wait for the max stake of any one node to drop below a percentage of total.")
+                .arg(
+                    Arg::with_name("max_percent")
+                        .long("max-percent")
+                        .value_name("PERCENT")
+                        .takes_value(true)
+                        .index(1),
+                ),
+        )
     }
 }
 
@@ -1400,6 +1411,16 @@ pub fn process_show_stakes(
         .formatted_string(&CliStakeVec::new(stake_accounts)))
 }
 
+pub fn process_wait_for_max_stake(
+    rpc_client: &RpcClient,
+    config: &CliConfig,
+    max_stake_percent: f32,
+) -> ProcessResult {
+    let now = std::time::Instant::now();
+    rpc_client.wait_for_max_stake(config.commitment, max_stake_percent)?;
+    Ok(format!("Done waiting, took: {}s", now.elapsed().as_secs()))
+}
+
 pub fn process_show_validators(
     rpc_client: &RpcClient,
     config: &CliConfig,

client/src/rpc_client.rs

@@ -353,6 +353,38 @@ impl RpcClient {
         self.send(RpcRequest::GetVoteAccounts, json!([commitment_config]))
     }
 
+    pub fn wait_for_max_stake(
+        &self,
+        commitment: CommitmentConfig,
+        max_stake_percent: f32,
+    ) -> ClientResult<()> {
+        let mut current_percent;
+        loop {
+            let vote_accounts = self.get_vote_accounts_with_commitment(commitment)?;
+
+            let mut max = 0;
+            let total_active_stake = vote_accounts
+                .current
+                .iter()
+                .chain(vote_accounts.delinquent.iter())
+                .map(|vote_account| {
+                    max = std::cmp::max(max, vote_account.activated_stake);
+                    vote_account.activated_stake
+                })
+                .sum::<u64>();
+            current_percent = 100f32 * max as f32 / total_active_stake as f32;
+            if current_percent < max_stake_percent {
+                break;
+            }
+            info!(
+                "Waiting for stake to drop below {} current: {:.1}",
+                max_stake_percent, current_percent
+            );
+            sleep(Duration::from_secs(10));
+        }
+        Ok(())
+    }
+
     pub fn get_cluster_nodes(&self) -> ClientResult<Vec<RpcContactInfo>> {
         self.send(RpcRequest::GetClusterNodes, Value::Null)
     }

local-cluster/tests/local_cluster.rs

@@ -1340,6 +1340,26 @@ fn test_faulty_node(faulty_node_type: BroadcastStageType) {
     cluster.check_for_new_roots(16, &"test_faulty_node");
 }
 
+#[test]
+fn test_wait_for_max_stake() {
+    solana_logger::setup();
+    let mut validator_config = ValidatorConfig::default();
+    validator_config.rpc_config.enable_validator_exit = true;
+    let config = ClusterConfig {
+        cluster_lamports: 10_000,
+        node_stakes: vec![100; 4],
+        validator_configs: vec![validator_config; 4],
+        ..ClusterConfig::default()
+    };
+    let cluster = LocalCluster::new(&config);
+    let client = RpcClient::new_socket(cluster.entry_point_info.rpc);
+
+    assert!(client
+        .wait_for_max_stake(CommitmentConfig::default(), 33.0f32)
+        .is_ok());
+    assert!(client.get_slot().unwrap() > 10);
+}
+
 #[test]
 // Test that when a leader is leader for banks B_i..B_{i+n}, and B_i is not
 // votable, then B_{i+1} still chains to B_i

system-test/automation_utils.sh

@@ -69,18 +69,8 @@ function wait_for_bootstrap_validator_stake_drop {
   source "${REPO_ROOT}"/net/common.sh
   loadConfigFile
 
-  while true; do
   # shellcheck disable=SC2154
-    bootstrap_validator_validator_info="$(ssh "${sshOptions[@]}" "${validatorIpList[0]}" '$HOME/.cargo/bin/solana --url http://127.0.0.1:8899 validators | grep "$($HOME/.cargo/bin/solana-keygen pubkey ~/solana/config/bootstrap-validator/identity.json)"')"
-    bootstrap_validator_stake_percentage="$(echo "$bootstrap_validator_validator_info" | awk '{gsub(/[\(,\),\%]/,""); print $9}')"
-
-    if [[ $(echo "$bootstrap_validator_stake_percentage < $max_stake" | bc) -ne 0 ]]; then
-      echo "Bootstrap validator stake has fallen below $max_stake to $bootstrap_validator_stake_percentage"
-      break
-    fi
-    echo "Max bootstrap validator stake: $max_stake.  Current stake: $bootstrap_validator_stake_percentage.  Sleeping 30s for stake to distribute."
-    sleep 30
-  done
+  ssh "${sshOptions[@]}" "${validatorIpList[0]}" '$HOME/.cargo/bin/solana wait-for-max-stake $max_stake --url http://127.0.0.1:8899'
 }
 
 function get_slot {