token-client: Specify compute unit limit more clearly

Author: joncinque

token/cli/src/clap_app.rs

@@ -15,6 +15,7 @@ use {
     },
     solana_sdk::{instruction::AccountMeta, pubkey::Pubkey},
     spl_token_2022::instruction::{AuthorityType, MAX_SIGNERS, MIN_SIGNERS},
+    spl_token_client::token::ComputeUnitLimit,
     std::{fmt, str::FromStr},
     strum::IntoEnumIterator,
     strum_macros::{EnumIter, EnumString, IntoStaticStr},
@@ -73,7 +74,10 @@ pub const COMPUTE_UNIT_PRICE_ARG: ArgConstant<'static> = ArgConstant {
 pub const COMPUTE_UNIT_LIMIT_ARG: ArgConstant<'static> = ArgConstant {
     name: "compute_unit_limit",
     long: "--with-compute-unit-limit",
-    help: "Set compute unit limit for transaction, in compute units.",
+    help: "Set compute unit limit for transaction, in compute units; also accepts \
+        keyword SIMULATED to use compute units from transaction simulation prior \
+        to sending. Note that SIMULATED may fail if accounts are modified by another \
+        transaction between simulation and execution.",
 };
 
 pub static VALID_TOKEN_PROGRAM_IDS: [Pubkey; 2] = [spl_token_2022::ID, spl_token::ID];
@@ -350,6 +354,32 @@ where
         Err(e) => Err(e),
     }
 }
+
+fn is_compute_unit_limit_or_simulated<T>(string: T) -> Result<(), String>
+where
+    T: AsRef<str> + fmt::Display,
+{
+    if string.as_ref().parse::<u32>().is_ok() || string.as_ref() == "SIMULATED" {
+        Ok(())
+    } else {
+        Err(format!(
+            "Unable to parse input compute unit limit as integer or SIMULATED, provided: {string}"
+        ))
+    }
+}
+pub(crate) fn parse_compute_unit_limit<T>(string: T) -> Result<ComputeUnitLimit, String>
+where
+    T: AsRef<str> + fmt::Display,
+{
+    match string.as_ref().parse::<u32>() {
+        Ok(compute_unit_limit) => Ok(ComputeUnitLimit::Static(compute_unit_limit)),
+        Err(_) if string.as_ref() == "SIMULATED" => Ok(ComputeUnitLimit::Simulated),
+        _ => Err(format!(
+            "Unable to parse compute unit limit, provided: {string}"
+        )),
+    }
+}
+
 struct SignOnlyNeedsFullMintSpec {}
 impl offline::ArgsConfig for SignOnlyNeedsFullMintSpec {
     fn sign_only_arg<'a, 'b>(&self, arg: Arg<'a, 'b>) -> Arg<'a, 'b> {
@@ -628,7 +658,7 @@ pub fn app<'a, 'b>(
                 .long(COMPUTE_UNIT_LIMIT_ARG.long)
                 .takes_value(true)
                 .value_name("COMPUTE-UNIT-LIMIT")
-                .validator(is_parsable::<u32>)
+                .validator(is_compute_unit_limit_or_simulated)
                 .help(COMPUTE_UNIT_LIMIT_ARG.help)
         )
         .arg(

token/cli/src/command.rs

@@ -68,7 +68,7 @@ use {
     },
     spl_token_client::{
         client::{ProgramRpcClientSendTransaction, RpcClientResponse},
-        token::{ExtensionInitializationParams, Token},
+        token::{ComputeUnitLimit, ExtensionInitializationParams, Token},
     },
     spl_token_group_interface::state::TokenGroup,
     spl_token_metadata_interface::state::{Field, TokenMetadata},
@@ -152,11 +152,7 @@ fn config_token_client(
     token: Token<ProgramRpcClientSendTransaction>,
     config: &Config<'_>,
 ) -> Result<Token<ProgramRpcClientSendTransaction>, Error> {
-    let token = if let Some(compute_unit_limit) = config.compute_unit_limit {
-        token.with_compute_unit_limit(compute_unit_limit)
-    } else {
-        token
-    };
+    let token = token.with_compute_unit_limit(config.compute_unit_limit.clone());
 
     let token = if let Some(compute_unit_price) = config.compute_unit_price {
         token.with_compute_unit_price(compute_unit_price)
@@ -434,7 +430,8 @@ async fn command_set_interest_rate(
 ) -> CommandResult {
     // Because set_interest_rate depends on the time, it can cost more between
     // simulation and execution. To help that, just set a static compute limit
-    let token = base_token_client(config, &token_pubkey, None)?.with_compute_unit_limit(2_500);
+    let token = base_token_client(config, &token_pubkey, None)?
+        .with_compute_unit_limit(ComputeUnitLimit::Static(2_500));
     let token = config_token_client(token, config)?;
 
     if !config.sign_only {

token/cli/src/config.rs

@@ -1,5 +1,8 @@
 use {
-    crate::clap_app::{Error, COMPUTE_UNIT_LIMIT_ARG, COMPUTE_UNIT_PRICE_ARG, MULTISIG_SIGNER_ARG},
+    crate::clap_app::{
+        parse_compute_unit_limit, Error, COMPUTE_UNIT_LIMIT_ARG, COMPUTE_UNIT_PRICE_ARG,
+        MULTISIG_SIGNER_ARG,
+    },
     clap::ArgMatches,
     solana_clap_utils::{
         input_parsers::{pubkey_of_signer, value_of},
@@ -20,8 +23,11 @@ use {
         extension::StateWithExtensionsOwned,
         state::{Account, Mint},
     },
-    spl_token_client::client::{
-        ProgramClient, ProgramOfflineClient, ProgramRpcClient, ProgramRpcClientSendTransaction,
+    spl_token_client::{
+        client::{
+            ProgramClient, ProgramOfflineClient, ProgramRpcClient, ProgramRpcClientSendTransaction,
+        },
+        token::ComputeUnitLimit,
     },
     std::{process::exit, rc::Rc, sync::Arc},
 };
@@ -68,7 +74,7 @@ pub struct Config<'a> {
     pub program_id: Pubkey,
     pub restrict_to_program_id: bool,
     pub compute_unit_price: Option<u64>,
-    pub compute_unit_limit: Option<u32>,
+    pub compute_unit_limit: ComputeUnitLimit,
 }
 
 impl<'a> Config<'a> {
@@ -282,7 +288,10 @@ impl<'a> Config<'a> {
 
         let nonce_blockhash = value_of(matches, BLOCKHASH_ARG.name);
         let compute_unit_price = value_of(matches, COMPUTE_UNIT_PRICE_ARG.name);
-        let compute_unit_limit = value_of(matches, COMPUTE_UNIT_LIMIT_ARG.name);
+        let compute_unit_limit = matches
+            .value_of(COMPUTE_UNIT_LIMIT_ARG.name)
+            .map(|x| parse_compute_unit_limit(x).unwrap())
+            .unwrap_or(ComputeUnitLimit::Default);
         Self {
             default_signer,
             rpc_client,

token/cli/tests/command.rs

@@ -44,7 +44,7 @@ use {
         client::{
             ProgramClient, ProgramOfflineClient, ProgramRpcClient, ProgramRpcClientSendTransaction,
         },
-        token::Token,
+        token::{ComputeUnitLimit, Token},
     },
     spl_token_group_interface::state::{TokenGroup, TokenGroupMember},
     spl_token_metadata_interface::state::TokenMetadata,
@@ -201,7 +201,7 @@ fn test_config_with_default_signer<'a>(
         program_id: *program_id,
         restrict_to_program_id: true,
         compute_unit_price: None,
-        compute_unit_limit: None,
+        compute_unit_limit: ComputeUnitLimit::Simulated,
     }
 }
 
@@ -230,7 +230,7 @@ fn test_config_without_default_signer<'a>(
         program_id: *program_id,
         restrict_to_program_id: true,
         compute_unit_price: None,
-        compute_unit_limit: None,
+        compute_unit_limit: ComputeUnitLimit::Simulated,
     }
 }
 
@@ -2991,6 +2991,7 @@ async fn offline_multisig_transfer_with_nonce(test_validator: &TestValidator, pa
         .unzip();
     for program_id in VALID_TOKEN_PROGRAM_IDS.iter() {
         let mut config = test_config_with_default_signer(test_validator, payer, program_id);
+        config.compute_unit_limit = ComputeUnitLimit::Default;
         let token = create_token(&config, payer).await;
         let nonce = create_nonce(&config, payer).await;
 
@@ -4081,7 +4082,7 @@ async fn compute_budget(test_validator: &TestValidator, payer: &Keypair) {
     for program_id in VALID_TOKEN_PROGRAM_IDS.iter() {
         let mut config = test_config_with_default_signer(test_validator, payer, program_id);
         config.compute_unit_price = Some(42);
-        config.compute_unit_limit = Some(30_000);
+        config.compute_unit_limit = ComputeUnitLimit::Static(30_000);
         run_transfer_test(&config, payer).await;
     }
 }

token/client/src/token.rs

@@ -332,6 +332,13 @@ impl TokenMemo {
     }
 }
 
+#[derive(Debug, Clone)]
+pub enum ComputeUnitLimit {
+    Default,
+    Simulated,
+    Static(u32),
+}
+
 pub struct Token<T> {
     client: Arc<dyn ProgramClient<T>>,
     pubkey: Pubkey, /* token mint */
@@ -344,7 +351,7 @@ pub struct Token<T> {
     memo: Arc<RwLock<Option<TokenMemo>>>,
     transfer_hook_accounts: Option<Vec<AccountMeta>>,
     compute_unit_price: Option<u64>,
-    compute_unit_limit: Option<u32>,
+    compute_unit_limit: ComputeUnitLimit,
 }
 
 impl<T> fmt::Debug for Token<T> {
@@ -411,7 +418,7 @@ where
             memo: Arc::new(RwLock::new(None)),
             transfer_hook_accounts: None,
             compute_unit_price: None,
-            compute_unit_limit: None,
+            compute_unit_limit: ComputeUnitLimit::Default,
         }
     }
 
@@ -466,8 +473,8 @@ where
         self
     }
 
-    pub fn with_compute_unit_limit(mut self, compute_unit_limit: u32) -> Self {
-        self.compute_unit_limit = Some(compute_unit_limit);
+    pub fn with_compute_unit_limit(mut self, compute_unit_limit: ComputeUnitLimit) -> Self {
+        self.compute_unit_limit = compute_unit_limit;
         self
     }
 
@@ -548,19 +555,16 @@ where
             .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();
-        }
+        let units_consumed = simulation_result
+            .get_compute_units_consumed()
+            .map_err(TokenError::Client)?;
+        // 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;
         Ok(())
     }
 
@@ -619,13 +623,17 @@ where
         // 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?;
+        match self.compute_unit_limit {
+            ComputeUnitLimit::Default => {}
+            ComputeUnitLimit::Simulated => {
+                self.add_compute_unit_limit_from_simulation(&mut instructions, &blockhash)
+                    .await?;
+            }
+            ComputeUnitLimit::Static(compute_unit_limit) => {
+                instructions.push(ComputeBudgetInstruction::set_compute_unit_limit(
+                    compute_unit_limit,
+                ));
+            }
         }
 
         let message = Message::new_with_blockhash(&instructions, fee_payer, &blockhash);

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

@@ -39,7 +39,7 @@ use {
     },
     spl_token_client::{
         proof_generation::transfer_with_fee_split_proof_data,
-        token::{ExtensionInitializationParams, TokenError as TokenClientError},
+        token::{ComputeUnitLimit, ExtensionInitializationParams, TokenError as TokenClientError},
     },
     std::{convert::TryInto, mem::size_of},
 };
@@ -2526,7 +2526,7 @@ async fn confidential_transfer_transfer_with_split_proof_contexts_in_parallel()
     // 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);
+    let token = token.with_compute_unit_limit(ComputeUnitLimit::Static(500_000));
     token
         .confidential_transfer_transfer_with_split_proofs_in_parallel(
             &alice_meta.token_account,
@@ -2948,7 +2948,7 @@ async fn confidential_transfer_transfer_with_fee_and_split_proof_context_in_para
     // 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);
+    let token = token.with_compute_unit_limit(ComputeUnitLimit::Static(500_000));
     token
         .confidential_transfer_transfer_with_fee_and_split_proofs_in_parallel(
             &alice_meta.token_account,

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

@@ -20,7 +20,7 @@ use {
             ProgramBanksClient, ProgramBanksClientProcessTransaction, ProgramClient,
             SendTransaction, SimulateTransaction,
         },
-        token::{ExtensionInitializationParams, Token, TokenResult},
+        token::{ComputeUnitLimit, ExtensionInitializationParams, Token, TokenResult},
     },
     std::sync::Arc,
 };
@@ -113,15 +113,17 @@ impl TestContext {
             &mint_account.pubkey(),
             Some(decimals),
             Arc::new(keypair_clone(&payer)),
-        );
+        )
+        .with_compute_unit_limit(ComputeUnitLimit::Simulated);
 
         let token_unchecked = Token::new(
             Arc::clone(&client),
             &id(),
             &mint_account.pubkey(),
             None,
             Arc::new(payer),
-        );
+        )
+        .with_compute_unit_limit(ComputeUnitLimit::Simulated);
 
         token
             .create_mint(
@@ -157,7 +159,8 @@ impl TestContext {
             Token::create_native_mint(Arc::clone(&client), &id(), Arc::new(keypair_clone(&payer)))
                 .await?;
         // unchecked native is never needed because decimals is known statically
-        let token_unchecked = Token::new_native(Arc::clone(&client), &id(), Arc::new(payer));
+        let token_unchecked = Token::new_native(Arc::clone(&client), &id(), Arc::new(payer))
+            .with_compute_unit_limit(ComputeUnitLimit::Simulated);
         self.token_context = Some(TokenContext {
             decimals: native_mint::DECIMALS,
             mint_authority: Keypair::new(), /* bogus */