cli: Use simulated compute unit limit in stake interactions (#2710)

* cli: Use simulated compute unit limit in stake interactions

#### Problem

The CLI has the ability to simulate transactions before sending to use
the correct number of compute units, but stake commands are still using
the default compute unit limit.

#### Summary of changes

Update tests to use a compute unit limit, and then update the stake
commands to use the simulated compute unit limit.

* Simulate for compute units in stake commands

* Pass compute unit limit enum to `simulate_and_update...`

cli/src/compute_budget.rs

@@ -12,12 +12,12 @@ use {
     },
 };
 
-// This enum is equivalent to an Option but was added to self-document
-// the ok variants and has the benefit of not forcing the caller to use
-// the result if they don't care about it.
+/// Enum capturing the possible results of updating a message based on the
+/// compute unit limits consumed during simulation.
 pub(crate) enum UpdateComputeUnitLimitResult {
     UpdatedInstructionIndex(usize),
     NoInstructionFound,
+    SimulationNotConfigured,
 }
 
 fn get_compute_unit_limit_instruction_index(message: &Message) -> Option<usize> {
@@ -83,8 +83,14 @@ pub(crate) fn simulate_for_compute_unit_limit(
     simulate_for_compute_unit_limit_unchecked(rpc_client, message)
 }
 
-// Returns the index of the compute unit limit instruction
+/// Simulates a message and returns the index of the compute unit limit
+/// instruction
+///
+/// If the message does not contain a compute unit limit instruction, or if
+/// simulation was not configured, then the function will not simulate the
+/// message.
 pub(crate) fn simulate_and_update_compute_unit_limit(
+    compute_unit_limit: &ComputeUnitLimit,
     rpc_client: &RpcClient,
     message: &mut Message,
 ) -> Result<UpdateComputeUnitLimitResult, Box<dyn std::error::Error>> {
@@ -93,15 +99,23 @@ pub(crate) fn simulate_and_update_compute_unit_limit(
         return Ok(UpdateComputeUnitLimitResult::NoInstructionFound);
     };
 
-    let compute_unit_limit = simulate_for_compute_unit_limit_unchecked(rpc_client, message)?;
+    match compute_unit_limit {
+        ComputeUnitLimit::Simulated => {
+            let compute_unit_limit =
+                simulate_for_compute_unit_limit_unchecked(rpc_client, message)?;
 
-    // Overwrite the compute unit limit instruction with the actual units consumed
-    message.instructions[compute_unit_limit_ix_index].data =
-        ComputeBudgetInstruction::set_compute_unit_limit(compute_unit_limit).data;
+            // Overwrite the compute unit limit instruction with the actual units consumed
+            message.instructions[compute_unit_limit_ix_index].data =
+                ComputeBudgetInstruction::set_compute_unit_limit(compute_unit_limit).data;
 
-    Ok(UpdateComputeUnitLimitResult::UpdatedInstructionIndex(
-        compute_unit_limit_ix_index,
-    ))
+            Ok(UpdateComputeUnitLimitResult::UpdatedInstructionIndex(
+                compute_unit_limit_ix_index,
+            ))
+        }
+        ComputeUnitLimit::Static(_) | ComputeUnitLimit::Default => {
+            Ok(UpdateComputeUnitLimitResult::SimulationNotConfigured)
+        }
+    }
 }
 
 pub(crate) struct ComputeUnitConfig {

cli/src/nonce.rs

@@ -415,6 +415,7 @@ pub fn process_authorize_nonce_account(
     let latest_blockhash = rpc_client.get_latest_blockhash()?;
 
     let nonce_authority = config.signers[nonce_authority];
+    let compute_unit_limit = ComputeUnitLimit::Simulated;
     let ixs = vec![authorize_nonce_account(
         nonce_account,
         &nonce_authority.pubkey(),
@@ -423,10 +424,10 @@ pub fn process_authorize_nonce_account(
     .with_memo(memo)
     .with_compute_unit_config(&ComputeUnitConfig {
         compute_unit_price,
-        compute_unit_limit: ComputeUnitLimit::Simulated,
+        compute_unit_limit,
     });
     let mut message = Message::new(&ixs, Some(&config.signers[0].pubkey()));
-    simulate_and_update_compute_unit_limit(rpc_client, &mut message)?;
+    simulate_and_update_compute_unit_limit(&compute_unit_limit, rpc_client, &mut message)?;
     let mut tx = Transaction::new_unsigned(message);
     tx.try_sign(&config.signers, latest_blockhash)?;
 
@@ -575,18 +576,19 @@ pub fn process_new_nonce(
     }
 
     let nonce_authority = config.signers[nonce_authority];
+    let compute_unit_limit = ComputeUnitLimit::Simulated;
     let ixs = vec![advance_nonce_account(
         nonce_account,
         &nonce_authority.pubkey(),
     )]
     .with_memo(memo)
     .with_compute_unit_config(&ComputeUnitConfig {
         compute_unit_price,
-        compute_unit_limit: ComputeUnitLimit::Simulated,
+        compute_unit_limit,
     });
     let latest_blockhash = rpc_client.get_latest_blockhash()?;
     let mut message = Message::new(&ixs, Some(&config.signers[0].pubkey()));
-    simulate_and_update_compute_unit_limit(rpc_client, &mut message)?;
+    simulate_and_update_compute_unit_limit(&compute_unit_limit, rpc_client, &mut message)?;
     let mut tx = Transaction::new_unsigned(message);
     tx.try_sign(&config.signers, latest_blockhash)?;
     check_account_for_fee_with_commitment(
@@ -643,6 +645,7 @@ pub fn process_withdraw_from_nonce_account(
     let latest_blockhash = rpc_client.get_latest_blockhash()?;
 
     let nonce_authority = config.signers[nonce_authority];
+    let compute_unit_limit = ComputeUnitLimit::Simulated;
     let ixs = vec![withdraw_nonce_account(
         nonce_account,
         &nonce_authority.pubkey(),
@@ -652,10 +655,10 @@ pub fn process_withdraw_from_nonce_account(
     .with_memo(memo)
     .with_compute_unit_config(&ComputeUnitConfig {
         compute_unit_price,
-        compute_unit_limit: ComputeUnitLimit::Simulated,
+        compute_unit_limit,
     });
     let mut message = Message::new(&ixs, Some(&config.signers[0].pubkey()));
-    simulate_and_update_compute_unit_limit(rpc_client, &mut message)?;
+    simulate_and_update_compute_unit_limit(&compute_unit_limit, rpc_client, &mut message)?;
     let mut tx = Transaction::new_unsigned(message);
     tx.try_sign(&config.signers, latest_blockhash)?;
     check_account_for_fee_with_commitment(
@@ -677,14 +680,15 @@ pub(crate) fn process_upgrade_nonce_account(
     compute_unit_price: Option<u64>,
 ) -> ProcessResult {
     let latest_blockhash = rpc_client.get_latest_blockhash()?;
+    let compute_unit_limit = ComputeUnitLimit::Simulated;
     let ixs = vec![upgrade_nonce_account(nonce_account)]
         .with_memo(memo)
         .with_compute_unit_config(&ComputeUnitConfig {
             compute_unit_price,
-            compute_unit_limit: ComputeUnitLimit::Simulated,
+            compute_unit_limit,
         });
     let mut message = Message::new(&ixs, Some(&config.signers[0].pubkey()));
-    simulate_and_update_compute_unit_limit(rpc_client, &mut message)?;
+    simulate_and_update_compute_unit_limit(&compute_unit_limit, rpc_client, &mut message)?;
     let mut tx = Transaction::new_unsigned(message);
     tx.try_sign(&config.signers, latest_blockhash)?;
     check_account_for_fee_with_commitment(

cli/src/program.rs

@@ -2420,6 +2420,7 @@ fn do_process_program_deploy(
     use_rpc: bool,
 ) -> ProcessResult {
     let blockhash = rpc_client.get_latest_blockhash()?;
+    let compute_unit_limit = ComputeUnitLimit::Simulated;
 
     let (initial_instructions, balance_needed, buffer_program_data) =
         if let Some(buffer_program_data) = buffer_program_data {
@@ -2442,7 +2443,7 @@ fn do_process_program_deploy(
         Some(Message::new_with_blockhash(
             &initial_instructions.with_compute_unit_config(&ComputeUnitConfig {
                 compute_unit_price,
-                compute_unit_limit: ComputeUnitLimit::Simulated,
+                compute_unit_limit,
             }),
             Some(&fee_payer_signer.pubkey()),
             &blockhash,
@@ -2462,7 +2463,7 @@ fn do_process_program_deploy(
 
         let instructions = vec![instruction].with_compute_unit_config(&ComputeUnitConfig {
             compute_unit_price,
-            compute_unit_limit: ComputeUnitLimit::Simulated,
+            compute_unit_limit,
         });
         Message::new_with_blockhash(&instructions, Some(&fee_payer_signer.pubkey()), &blockhash)
     };
@@ -2489,7 +2490,7 @@ fn do_process_program_deploy(
         )?
         .with_compute_unit_config(&ComputeUnitConfig {
             compute_unit_price,
-            compute_unit_limit: ComputeUnitLimit::Simulated,
+            compute_unit_limit,
         });
 
         Some(Message::new_with_blockhash(
@@ -2523,6 +2524,7 @@ fn do_process_program_deploy(
         Some(program_signers),
         max_sign_attempts,
         use_rpc,
+        &compute_unit_limit,
     )?;
 
     let program_id = CliProgramId {
@@ -2550,6 +2552,7 @@ fn do_process_write_buffer(
     use_rpc: bool,
 ) -> ProcessResult {
     let blockhash = rpc_client.get_latest_blockhash()?;
+    let compute_unit_limit = ComputeUnitLimit::Simulated;
 
     let (initial_instructions, balance_needed, buffer_program_data) =
         if let Some(buffer_program_data) = buffer_program_data {
@@ -2572,7 +2575,7 @@ fn do_process_write_buffer(
         Some(Message::new_with_blockhash(
             &initial_instructions.with_compute_unit_config(&ComputeUnitConfig {
                 compute_unit_price,
-                compute_unit_limit: ComputeUnitLimit::Simulated,
+                compute_unit_limit,
             }),
             Some(&fee_payer_signer.pubkey()),
             &blockhash,
@@ -2592,7 +2595,7 @@ fn do_process_write_buffer(
 
         let instructions = vec![instruction].with_compute_unit_config(&ComputeUnitConfig {
             compute_unit_price,
-            compute_unit_limit: ComputeUnitLimit::Simulated,
+            compute_unit_limit,
         });
         Message::new_with_blockhash(&instructions, Some(&fee_payer_signer.pubkey()), &blockhash)
     };
@@ -2630,6 +2633,7 @@ fn do_process_write_buffer(
         None,
         max_sign_attempts,
         use_rpc,
+        &compute_unit_limit,
     )?;
 
     let buffer = CliProgramBuffer {
@@ -2658,6 +2662,7 @@ fn do_process_program_upgrade(
     use_rpc: bool,
 ) -> ProcessResult {
     let blockhash = rpc_client.get_latest_blockhash()?;
+    let compute_unit_limit = ComputeUnitLimit::Simulated;
 
     let (initial_message, write_messages, balance_needed) = if let Some(buffer_signer) =
         buffer_signer
@@ -2714,7 +2719,7 @@ fn do_process_program_upgrade(
             )]
             .with_compute_unit_config(&ComputeUnitConfig {
                 compute_unit_price,
-                compute_unit_limit: ComputeUnitLimit::Simulated,
+                compute_unit_limit,
             });
             Message::new_with_blockhash(&instructions, Some(&fee_payer_signer.pubkey()), &blockhash)
         };
@@ -2743,7 +2748,7 @@ fn do_process_program_upgrade(
     )]
     .with_compute_unit_config(&ComputeUnitConfig {
         compute_unit_price,
-        compute_unit_limit: ComputeUnitLimit::Simulated,
+        compute_unit_limit,
     });
     let final_message = Message::new_with_blockhash(
         &final_instructions,
@@ -2776,6 +2781,7 @@ fn do_process_program_upgrade(
         Some(&[upgrade_authority]),
         max_sign_attempts,
         use_rpc,
+        &compute_unit_limit,
     )?;
 
     let program_id = CliProgramId {
@@ -2912,11 +2918,12 @@ fn send_deploy_messages(
     final_signers: Option<&[&dyn Signer]>,
     max_sign_attempts: usize,
     use_rpc: bool,
+    compute_unit_limit: &ComputeUnitLimit,
 ) -> Result<Option<Signature>, Box<dyn std::error::Error>> {
     if let Some(mut message) = initial_message {
         if let Some(initial_signer) = initial_signer {
             trace!("Preparing the required accounts");
-            simulate_and_update_compute_unit_limit(&rpc_client, &mut message)?;
+            simulate_and_update_compute_unit_limit(compute_unit_limit, &rpc_client, &mut message)?;
             let mut initial_transaction = Transaction::new_unsigned(message.clone());
             let blockhash = rpc_client.get_latest_blockhash()?;
 
@@ -2947,7 +2954,11 @@ fn send_deploy_messages(
             {
                 let mut message = write_messages[0].clone();
                 if let UpdateComputeUnitLimitResult::UpdatedInstructionIndex(ix_index) =
-                    simulate_and_update_compute_unit_limit(&rpc_client, &mut message)?
+                    simulate_and_update_compute_unit_limit(
+                        compute_unit_limit,
+                        &rpc_client,
+                        &mut message,
+                    )?
                 {
                     for msg in &mut write_messages {
                         // Write messages are all assumed to be identical except
@@ -3024,7 +3035,7 @@ fn send_deploy_messages(
         if let Some(final_signers) = final_signers {
             trace!("Deploying program");
 
-            simulate_and_update_compute_unit_limit(&rpc_client, &mut message)?;
+            simulate_and_update_compute_unit_limit(compute_unit_limit, &rpc_client, &mut message)?;
             let mut final_tx = Transaction::new_unsigned(message);
             let blockhash = rpc_client.get_latest_blockhash()?;
             let mut signers = final_signers.to_vec();

cli/src/spend_utils.rs

@@ -197,18 +197,18 @@ where
             let mut dummy_message = build_message(lamports);
 
             dummy_message.recent_blockhash = *blockhash;
-            let compute_unit_info = if compute_unit_limit == ComputeUnitLimit::Simulated {
-                // Simulate for correct compute units
+            let compute_unit_info =
                 if let UpdateComputeUnitLimitResult::UpdatedInstructionIndex(ix_index) =
-                    simulate_and_update_compute_unit_limit(rpc_client, &mut dummy_message)?
+                    simulate_and_update_compute_unit_limit(
+                        &compute_unit_limit,
+                        rpc_client,
+                        &mut dummy_message,
+                    )?
                 {
                     Some((ix_index, dummy_message.instructions[ix_index].data.clone()))
                 } else {
                     None
-                }
-            } else {
-                None
-            };
+                };
             (
                 get_fee_for_messages(rpc_client, &[&dummy_message])?,
                 compute_unit_info,