cli: Use simulated compute units in ping (solana-labs#2693)

* cli: Use simulated compute units in ping

#### Problem

The CLI has the ability to simulate transactions before sending to use
the correct number of compute units, but `solana ping` is still using
the default compute unit limit.

#### Summary of changes

Simulate once to get the compute unit limit and then re-use the
simulated number for every ping.

* Refactor per review

* Only get compute unit limit if simulation needed, add test

cli/src/cluster_query.rs

@@ -1,7 +1,9 @@
 use {
     crate::{
         cli::{CliCommand, CliCommandInfo, CliConfig, CliError, ProcessResult},
-        compute_budget::{ComputeUnitConfig, WithComputeUnitConfig},
+        compute_budget::{
+            simulate_for_compute_unit_limit, ComputeUnitConfig, WithComputeUnitConfig,
+        },
         feature::get_feature_activation_epoch,
         spend_utils::{resolve_spend_tx_and_check_account_balance, SpendAmount},
     },
@@ -1435,10 +1437,14 @@ pub fn process_ping(
     rpc_client: &RpcClient,
 ) -> ProcessResult {
     let (signal_sender, signal_receiver) = unbounded();
-    ctrlc::set_handler(move || {
+    match ctrlc::try_set_handler(move || {
         let _ = signal_sender.send(());
-    })
-    .expect("Error setting Ctrl-C handler");
+    }) {
+        // It's possible to set the ctrl-c handler more than once in testing
+        // situations, so let that case through
+        Err(ctrlc::Error::MultipleHandlers) => {}
+        result => result.expect("Error setting Ctrl-C handler"),
+    }
 
     let mut cli_pings = vec![];
 
@@ -1458,6 +1464,23 @@ pub fn process_ping(
         }
     }
 
+    let to = config.signers[0].pubkey();
+    let compute_unit_limit = if compute_unit_price.is_some() {
+        let ixs = vec![system_instruction::transfer(
+            &config.signers[0].pubkey(),
+            &to,
+            lamports,
+        )]
+        .with_compute_unit_config(&ComputeUnitConfig {
+            compute_unit_price,
+            compute_unit_limit: ComputeUnitLimit::Simulated,
+        });
+        let message = Message::new(&ixs, Some(&config.signers[0].pubkey()));
+        ComputeUnitLimit::Static(simulate_for_compute_unit_limit(rpc_client, &message)?)
+    } else {
+        ComputeUnitLimit::Default
+    };
+
     'mainloop: for seq in 0..count.unwrap_or(u64::MAX) {
         let now = Instant::now();
         if fixed_blockhash.is_none() && now.duration_since(blockhash_acquired).as_secs() > 60 {
@@ -1468,10 +1491,8 @@ pub fn process_ping(
             blockhash_acquired = Instant::now();
         }
 
-        let to = config.signers[0].pubkey();
         lamports = lamports.saturating_add(1);
 
-        let compute_unit_limit = ComputeUnitLimit::Default;
         let build_message = |lamports| {
             let ixs = vec![system_instruction::transfer(
                 &config.signers[0].pubkey(),

cli/src/compute_budget.rs

@@ -20,32 +20,31 @@ pub(crate) enum UpdateComputeUnitLimitResult {
     NoInstructionFound,
 }
 
-// Returns the index of the compute unit limit instruction
-pub(crate) fn simulate_and_update_compute_unit_limit(
-    rpc_client: &RpcClient,
-    message: &mut Message,
-) -> Result<UpdateComputeUnitLimitResult, Box<dyn std::error::Error>> {
-    let Some(compute_unit_limit_ix_index) =
-        message
-            .instructions
-            .iter()
-            .enumerate()
-            .find_map(|(ix_index, instruction)| {
-                let ix_program_id = message.program_id(ix_index)?;
-                if ix_program_id != &compute_budget::id() {
-                    return None;
-                }
+fn get_compute_unit_limit_instruction_index(message: &Message) -> Option<usize> {
+    message
+        .instructions
+        .iter()
+        .enumerate()
+        .find_map(|(ix_index, instruction)| {
+            let ix_program_id = message.program_id(ix_index)?;
+            if ix_program_id != &compute_budget::id() {
+                return None;
+            }
 
-                matches!(
-                    try_from_slice_unchecked(&instruction.data),
-                    Ok(ComputeBudgetInstruction::SetComputeUnitLimit(_))
-                )
-                .then_some(ix_index)
-            })
-    else {
-        return Ok(UpdateComputeUnitLimitResult::NoInstructionFound);
-    };
+            matches!(
+                try_from_slice_unchecked(&instruction.data),
+                Ok(ComputeBudgetInstruction::SetComputeUnitLimit(_))
+            )
+            .then_some(ix_index)
+        })
+}
 
+/// Like `simulate_for_compute_unit_limit`, but does not check that the message
+/// contains a compute unit limit instruction.
+fn simulate_for_compute_unit_limit_unchecked(
+    rpc_client: &RpcClient,
+    message: &Message,
+) -> Result<u32, Box<dyn std::error::Error>> {
     let transaction = Transaction::new_unsigned(message.clone());
     let simulate_result = rpc_client
         .simulate_transaction_with_config(
@@ -67,8 +66,36 @@ pub(crate) fn simulate_and_update_compute_unit_limit(
         .units_consumed
         .expect("compute units unavailable");
 
+    u32::try_from(units_consumed).map_err(Into::into)
+}
+
+/// Returns the compute unit limit used during simulation
+///
+/// Returns an error if the message does not contain a compute unit limit
+/// instruction or if the simulation fails.
+pub(crate) fn simulate_for_compute_unit_limit(
+    rpc_client: &RpcClient,
+    message: &Message,
+) -> Result<u32, Box<dyn std::error::Error>> {
+    if get_compute_unit_limit_instruction_index(message).is_none() {
+        return Err("No compute unit limit instruction found".into());
+    }
+    simulate_for_compute_unit_limit_unchecked(rpc_client, message)
+}
+
+// Returns the index of the compute unit limit instruction
+pub(crate) fn simulate_and_update_compute_unit_limit(
+    rpc_client: &RpcClient,
+    message: &mut Message,
+) -> Result<UpdateComputeUnitLimitResult, Box<dyn std::error::Error>> {
+    let Some(compute_unit_limit_ix_index) = get_compute_unit_limit_instruction_index(message)
+    else {
+        return Ok(UpdateComputeUnitLimitResult::NoInstructionFound);
+    };
+
+    let compute_unit_limit = simulate_for_compute_unit_limit_unchecked(rpc_client, message)?;
+
     // Overwrite the compute unit limit instruction with the actual units consumed
-    let compute_unit_limit = u32::try_from(units_consumed)?;
     message.instructions[compute_unit_limit_ix_index].data =
         ComputeBudgetInstruction::set_compute_unit_limit(compute_unit_limit).data;
 

cli/tests/cluster_query.rs

@@ -0,0 +1,61 @@
+use {
+    solana_cli::{
+        check_balance,
+        cli::{process_command, request_and_confirm_airdrop, CliCommand, CliConfig},
+        test_utils::check_ready,
+    },
+    solana_faucet::faucet::run_local_faucet,
+    solana_rpc_client::rpc_client::RpcClient,
+    solana_sdk::{
+        commitment_config::CommitmentConfig,
+        fee::FeeStructure,
+        native_token::sol_to_lamports,
+        signature::{Keypair, Signer},
+    },
+    solana_streamer::socket::SocketAddrSpace,
+    solana_test_validator::TestValidator,
+    std::time::Duration,
+    test_case::test_case,
+};
+
+#[test_case(None; "base")]
+#[test_case(Some(1_000_000); "with_compute_unit_price")]
+fn test_ping(compute_unit_price: Option<u64>) {
+    solana_logger::setup();
+    let fee = FeeStructure::default().get_max_fee(1, 0);
+    let mint_keypair = Keypair::new();
+    let mint_pubkey = mint_keypair.pubkey();
+    let faucet_addr = run_local_faucet(mint_keypair, None);
+    let test_validator = TestValidator::with_custom_fees(
+        mint_pubkey,
+        fee,
+        Some(faucet_addr),
+        SocketAddrSpace::Unspecified,
+    );
+
+    let rpc_client =
+        RpcClient::new_with_commitment(test_validator.rpc_url(), CommitmentConfig::processed());
+
+    let default_signer = Keypair::new();
+    let signer_pubkey = default_signer.pubkey();
+
+    let mut config = CliConfig::recent_for_tests();
+    config.json_rpc_url = test_validator.rpc_url();
+    config.signers = vec![&default_signer];
+
+    request_and_confirm_airdrop(&rpc_client, &config, &signer_pubkey, sol_to_lamports(1.0))
+        .unwrap();
+    check_balance!(sol_to_lamports(1.0), &rpc_client, &signer_pubkey);
+    check_ready(&rpc_client);
+
+    let count = 5;
+    config.command = CliCommand::Ping {
+        interval: Duration::from_secs(0),
+        count: Some(count),
+        timeout: Duration::from_secs(15),
+        blockhash: None,
+        print_timestamp: false,
+        compute_unit_price,
+    };
+    process_command(&config).unwrap();
+}