token-client: Use transaction simulation results for the compute unit limit

token/client/src/client.rs

@@ -20,7 +20,12 @@ pub trait SendTransaction {
 
 /// Basic trait for simulating transactions in a validator.
 pub trait SimulateTransaction {
-    type SimulationOutput;
+    type SimulationOutput: SimulationResult;
+}
+
+/// Trait for the output of a simulation
+pub trait SimulationResult {
+    fn get_compute_units_consumed(&self) -> ProgramClientResult<u64>;
 }
 
 /// Extends basic `SendTransaction` trait with function `send` where client is
@@ -66,6 +71,15 @@ impl SendTransactionBanksClient for ProgramBanksClientProcessTransaction {
     }
 }
 
+impl SimulationResult for BanksTransactionResultWithSimulation {
+    fn get_compute_units_consumed(&self) -> ProgramClientResult<u64> {
+        self.simulation_details
+            .as_ref()
+            .map(|x| x.units_consumed)
+            .ok_or("No simulation results found".into())
+    }
+}
+
 impl SimulateTransaction for ProgramBanksClientProcessTransaction {
     type SimulationOutput = BanksTransactionResultWithSimulation;
 }
@@ -139,6 +153,20 @@ impl SendTransactionRpc for ProgramRpcClientSendTransaction {
     }
 }
 
+impl SimulationResult for RpcClientResponse {
+    fn get_compute_units_consumed(&self) -> ProgramClientResult<u64> {
+        match self {
+            // `Transaction` is the result of an offline simulation. The error
+            // should be properly handled by a caller that supports offline
+            // signing
+            Self::Signature(_) | Self::Transaction(_) => Err("Not a simulation result".into()),
+            Self::Simulation(simulation_result) => simulation_result
+                .units_consumed
+                .ok_or("No simulation results found".into()),
+        }
+    }
+}
+
 impl SimulateTransaction for ProgramRpcClientSendTransaction {
     type SimulationOutput = RpcClientResponse;
 }

token/client/src/token.rs

@@ -1,6 +1,9 @@
 use {
     crate::{
-        client::{ProgramClient, ProgramClientError, SendTransaction, SimulateTransaction},
+        client::{
+            ProgramClient, ProgramClientError, SendTransaction, SimulateTransaction,
+            SimulationResult,
+        },
         proof_generation::transfer_with_fee_split_proof_data,
     },
     futures::{future::join_all, try_join},
@@ -522,6 +525,45 @@ where
         }
     }
 
+    /// Helper function to add a compute unit limit instruction to a given set
+    /// of instructions
+    async fn add_compute_unit_limit_from_simulation(
+        &self,
+        instructions: &mut Vec<Instruction>,
+        blockhash: &Hash,
+    ) -> TokenResult<()> {
+        // add a max compute unit limit instruction for the simulation
+        const MAX_COMPUTE_UNIT_LIMIT: u32 = 1_400_000;
+        instructions.push(ComputeBudgetInstruction::set_compute_unit_limit(
+            MAX_COMPUTE_UNIT_LIMIT,
+        ));
+
+        let transaction = Transaction::new_unsigned(Message::new_with_blockhash(
+            instructions,
+            Some(&self.payer.pubkey()),
+            blockhash,
+        ));
+        let simulation_result = self
+            .client
+            .simulate_transaction(&transaction)
+            .await
+            .map_err(TokenError::Client)?;
+        if let Ok(units_consumed) = simulation_result.get_compute_units_consumed() {
+            // Overwrite the compute unit limit instruction with the actual units consumed
+            let compute_unit_limit =
+                u32::try_from(units_consumed).map_err(|x| TokenError::Client(x.into()))?;
+            instructions
+                .last_mut()
+                .expect("Compute budget instruction was added earlier")
+                .data = ComputeBudgetInstruction::set_compute_unit_limit(compute_unit_limit).data;
+        } else {
+            // `get_compute_units_consumed()` fails for offline signing, so we
+            // catch that error and remove the instruction that was added
+            instructions.pop();
+        }
+        Ok(())
+    }
+
     async fn construct_tx<S: Signers>(
         &self,
         token_instructions: &[Instruction],
@@ -549,18 +591,6 @@ where
 
         instructions.extend_from_slice(token_instructions);
 
-        if let Some(compute_unit_limit) = self.compute_unit_limit {
-            instructions.push(ComputeBudgetInstruction::set_compute_unit_limit(
-                compute_unit_limit,
-            ));
-        }
-
-        if let Some(compute_unit_price) = self.compute_unit_price {
-            instructions.push(ComputeBudgetInstruction::set_compute_unit_price(
-                compute_unit_price,
-            ));
-        }
-
         let blockhash = if let (Some(nonce_account), Some(nonce_authority), Some(nonce_blockhash)) = (
             self.nonce_account,
             &self.nonce_authority,
@@ -579,6 +609,25 @@ where
                 .map_err(TokenError::Client)?
         };
 
+        if let Some(compute_unit_price) = self.compute_unit_price {
+            instructions.push(ComputeBudgetInstruction::set_compute_unit_price(
+                compute_unit_price,
+            ));
+        }
+
+        // The simulation to find out the compute unit usage must be run after
+        // all instructions have been added to the transaction, so be sure to
+        // keep this instruction as the last one before creating and sending the
+        // transaction.
+        if let Some(compute_unit_limit) = self.compute_unit_limit {
+            instructions.push(ComputeBudgetInstruction::set_compute_unit_limit(
+                compute_unit_limit,
+            ));
+        } else {
+            self.add_compute_unit_limit_from_simulation(&mut instructions, &blockhash)
+                .await?;
+        }
+
         let message = Message::new_with_blockhash(&instructions, fee_payer, &blockhash);
         let mut transaction = Transaction::new_unsigned(message);
         let signing_pubkeys = signing_keypairs.pubkeys();
@@ -2087,12 +2136,26 @@ where
         )
         .await?;
 
-        self.process_ixs(
+        // This instruction is right at the transaction size limit, so we cannot
+        // add any other instructions to it
+        let blockhash = self
+            .client
+            .get_latest_blockhash()
+            .await
+            .map_err(TokenError::Client)?;
+
+        let transaction = Transaction::new_signed_with_payer(
             &[instruction_type
                 .encode_verify_proof(Some(withdraw_proof_context_state_info), withdraw_proof_data)],
-            &[] as &[&dyn Signer; 0],
-        )
-        .await
+            Some(&self.payer.pubkey()),
+            &[self.payer.as_ref()],
+            blockhash,
+        );
+
+        self.client
+            .send_transaction(&transaction)
+            .await
+            .map_err(TokenError::Client)
     }
 
     /// Transfer tokens confidentially
@@ -2571,12 +2634,24 @@ where
 
         // This instruction is right at the transaction size limit, but in the
         // future it might be able to support the transfer too
-        self.process_ixs(
+        let blockhash = self
+            .client
+            .get_latest_blockhash()
+            .await
+            .map_err(TokenError::Client)?;
+
+        let transaction = Transaction::new_signed_with_payer(
             &[instruction_type
                 .encode_verify_proof(Some(range_proof_context_state_info), range_proof_data)],
-            &[] as &[&dyn Signer; 0],
-        )
-        .await
+            Some(&self.payer.pubkey()),
+            &[self.payer.as_ref()],
+            blockhash,
+        );
+
+        self.client
+            .send_transaction(&transaction)
+            .await
+            .map_err(TokenError::Client)
     }
 
     /// Create a range proof context state account with a confidential transfer

token/program-2022-test/tests/confidential_transfer.rs

@@ -1110,7 +1110,6 @@ async fn confidential_transfer_transfer_with_fee() {
     let bob_meta = ConfidentialTokenAccountMeta::new(&token, &bob, None, false, true).await;
 
     // Self-transfer of 0 tokens
-    let token = token.with_compute_unit_limit(500_000);
     token
         .confidential_transfer_transfer_with_fee(
             &alice_meta.token_account,
@@ -2524,6 +2523,9 @@ async fn confidential_transfer_transfer_with_split_proof_contexts_in_parallel()
         &range_proof_context_state_account,
         &context_state_authority,
     ];
+    // With split proofs in parallel, one of the transactions does more work
+    // than the other, which isn't caught during the simulation to discover the
+    // compute unit limit.
     let token = token.with_compute_unit_limit(500_000);
     token
         .confidential_transfer_transfer_with_split_proofs_in_parallel(
@@ -2943,6 +2945,9 @@ async fn confidential_transfer_transfer_with_fee_and_split_proof_context_in_para
         &range_proof_context_state_account,
         &context_state_authority,
     ];
+    // With split proofs in parallel, one of the transactions does more work
+    // than the other, which isn't caught during the simulation to discover the
+    // compute unit limit.
     let token = token.with_compute_unit_limit(500_000);
     token
         .confidential_transfer_transfer_with_fee_and_split_proofs_in_parallel(

token/program-2022-test/tests/confidential_transfer_fee.rs

@@ -513,7 +513,6 @@ async fn confidential_transfer_withdraw_withheld_tokens_from_mint() {
     };
 
     // Test fee is 2.5% so the withheld fees should be 3
-    let token = token.with_compute_unit_limit(500_000);
     token
         .confidential_transfer_transfer_with_fee(
             &alice_meta.token_account,
@@ -672,7 +671,6 @@ async fn confidential_transfer_withdraw_withheld_tokens_from_accounts() {
     };
 
     // Test fee is 2.5% so the withheld fees should be 3
-    let token = token.with_compute_unit_limit(500_000);
     token
         .confidential_transfer_transfer_with_fee(
             &alice_meta.token_account,
@@ -803,7 +801,6 @@ async fn confidential_transfer_withdraw_withheld_tokens_from_mint_with_proof_con
     };
 
     // Test fee is 2.5% so the withheld fees should be 3
-    let token = token.with_compute_unit_limit(500_000);
     token
         .confidential_transfer_transfer_with_fee(
             &alice_meta.token_account,
@@ -972,7 +969,6 @@ async fn confidential_transfer_withdraw_withheld_tokens_from_accounts_with_proof
     };
 
     // Test fee is 2.5% so the withheld fees should be 3
-    let token = token.with_compute_unit_limit(500_000);
     token
         .confidential_transfer_transfer_with_fee(
             &alice_meta.token_account,
@@ -1163,7 +1159,6 @@ async fn confidential_transfer_harvest_withheld_tokens_to_mint() {
         .unwrap();
 
     // Test fee is 2.5% so the withheld fees should be 3
-    let token = token.with_compute_unit_limit(500_000);
     token
         .confidential_transfer_transfer_with_fee(
             &alice_meta.token_account,