solana-labs · joncinque · Mar 27, 2024 · Mar 26, 2024
diff --git a/token/cli/src/clap_app.rs b/token/cli/src/clap_app.rs
@@ -64,6 +64,18 @@ pub const MULTISIG_SIGNER_ARG: ArgConstant<'static> = ArgConstant {
     help: "Member signer of a multisig account",
 };
 
+pub const COMPUTE_UNIT_PRICE_ARG: ArgConstant<'static> = ArgConstant {
+    name: "compute_unit_price",
+    long: "--with-compute-unit-price",
+    help: "Set compute unit price for transaction, in increments of 0.000001 lamports per compute unit.",
+};
+
+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.",
+};
+
 pub static VALID_TOKEN_PROGRAM_IDS: [Pubkey; 2] = [spl_token_2022::ID, spl_token::ID];
 
 #[derive(Debug, Clone, Copy, PartialEq, EnumString, IntoStaticStr)]
@@ -611,6 +623,22 @@ pub fn app<'a, 'b>(
                 .hidden(true)
                 .help("Use unchecked instruction if appropriate. Supports transfer, burn, mint, and approve."),
         )
+        .arg(
+            Arg::with_name(COMPUTE_UNIT_LIMIT_ARG.name)
+                .long(COMPUTE_UNIT_LIMIT_ARG.long)
+                .takes_value(true)
+                .value_name("COMPUTE-UNIT-LIMIT")
+                .validator(is_parsable::<u32>)
+                .help(COMPUTE_UNIT_LIMIT_ARG.help)
+        )
+        .arg(
+            Arg::with_name(COMPUTE_UNIT_PRICE_ARG.name)
+                .long(COMPUTE_UNIT_PRICE_ARG.long)
+                .takes_value(true)
+                .value_name("COMPUTE-UNIT-PRICE")
+                .validator(is_parsable::<u64>)
+                .help(COMPUTE_UNIT_PRICE_ARG.help)
+        )
         .bench_subcommand()
         .subcommand(SubCommand::with_name(CommandName::CreateToken.into()).about("Create a new token")
                 .arg(

diff --git a/token/cli/src/command.rs b/token/cli/src/command.rs
@@ -147,6 +147,18 @@ fn token_client_from_config(
         config.fee_payer()?.clone(),
     );
 
+    let token = if let Some(compute_unit_limit) = config.compute_unit_limit {
+        token.with_compute_unit_limit(compute_unit_limit)
+    } else {
+        token
+    };
+
+    let token = if let Some(compute_unit_price) = config.compute_unit_price {
+        token.with_compute_unit_price(compute_unit_price)
+    } else {
+        token
+    };
+
     if let (Some(nonce_account), Some(nonce_authority), Some(nonce_blockhash)) = (
         config.nonce_account,
         &config.nonce_authority,

diff --git a/token/cli/src/config.rs b/token/cli/src/config.rs
@@ -1,5 +1,5 @@
 use {
-    crate::clap_app::{Error, MULTISIG_SIGNER_ARG},
+    crate::clap_app::{Error, COMPUTE_UNIT_LIMIT_ARG, COMPUTE_UNIT_PRICE_ARG, MULTISIG_SIGNER_ARG},
     clap::ArgMatches,
     solana_clap_utils::{
         input_parsers::{pubkey_of_signer, value_of},
@@ -67,6 +67,8 @@ pub struct Config<'a> {
     pub multisigner_pubkeys: Vec<&'a Pubkey>,
     pub program_id: Pubkey,
     pub restrict_to_program_id: bool,
+    pub compute_unit_price: Option<u64>,
+    pub compute_unit_limit: Option<u32>,
 }
 
 impl<'a> Config<'a> {
@@ -279,6 +281,8 @@ 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);
         Self {
             default_signer,
             rpc_client,
@@ -294,6 +298,8 @@ impl<'a> Config<'a> {
             multisigner_pubkeys,
             program_id,
             restrict_to_program_id,
+            compute_unit_price,
+            compute_unit_limit,
         }
     }
 

diff --git a/token/cli/tests/command.rs b/token/cli/tests/command.rs
@@ -135,6 +135,7 @@ async fn main() {
         async_trial!(metadata, test_validator, payer),
         async_trial!(group, test_validator, payer),
         async_trial!(confidential_transfer_with_fee, test_validator, payer),
+        async_trial!(compute_budget, test_validator, payer),
         // GC messes with every other test, so have it on its own test validator
         async_trial!(gc, gc_test_validator, gc_payer),
     ];
@@ -199,6 +200,8 @@ fn test_config_with_default_signer<'a>(
         multisigner_pubkeys: vec![],
         program_id: *program_id,
         restrict_to_program_id: true,
+        compute_unit_price: None,
+        compute_unit_limit: None,
     }
 }
 
@@ -226,6 +229,8 @@ fn test_config_without_default_signer<'a>(
         multisigner_pubkeys: vec![],
         program_id: *program_id,
         restrict_to_program_id: true,
+        compute_unit_price: None,
+        compute_unit_limit: None,
     }
 }
 
@@ -380,6 +385,38 @@ async fn mint_tokens(
     .await
 }
 
+async fn run_transfer_test(config: &Config<'_>, payer: &Keypair) {
+    let token = create_token(config, payer).await;
+    let source = create_associated_account(config, payer, &token, &payer.pubkey()).await;
+    let destination = create_auxiliary_account(config, payer, token).await;
+    let ui_amount = 100.0;
+    mint_tokens(config, payer, token, ui_amount, source)
+        .await
+        .unwrap();
+    let result = process_test_command(
+        config,
+        payer,
+        &[
+            "spl-token",
+            CommandName::Transfer.into(),
+            &token.to_string(),
+            "10",
+            &destination.to_string(),
+        ],
+    )
+    .await;
+    result.unwrap();
+
+    let account = config.rpc_client.get_account(&source).await.unwrap();
+    let token_account = StateWithExtensionsOwned::<Account>::unpack(account.data).unwrap();
+    let amount = spl_token::ui_amount_to_amount(90.0, TEST_DECIMALS);
+    assert_eq!(token_account.base.amount, amount);
+    let account = config.rpc_client.get_account(&destination).await.unwrap();
+    let token_account = StateWithExtensionsOwned::<Account>::unpack(account.data).unwrap();
+    let amount = spl_token::ui_amount_to_amount(10.0, TEST_DECIMALS);
+    assert_eq!(token_account.base.amount, amount);
+}
+
 async fn process_test_command<I, T>(config: &Config<'_>, payer: &Keypair, args: I) -> CommandResult
 where
     I: IntoIterator<Item = T>,
@@ -855,35 +892,7 @@ async fn wrapped_sol(test_validator: &TestValidator, payer: &Keypair) {
 async fn transfer(test_validator: &TestValidator, payer: &Keypair) {
     for program_id in VALID_TOKEN_PROGRAM_IDS.iter() {
         let config = test_config_with_default_signer(test_validator, payer, program_id);
-        let token = create_token(&config, payer).await;
-        let source = create_associated_account(&config, payer, &token, &payer.pubkey()).await;
-        let destination = create_auxiliary_account(&config, payer, token).await;
-        let ui_amount = 100.0;
-        mint_tokens(&config, payer, token, ui_amount, source)
-            .await
-            .unwrap();
-        let result = process_test_command(
-            &config,
-            payer,
-            &[
-                "spl-token",
-                CommandName::Transfer.into(),
-                &token.to_string(),
-                "10",
-                &destination.to_string(),
-            ],
-        )
-        .await;
-        result.unwrap();
-
-        let account = config.rpc_client.get_account(&source).await.unwrap();
-        let token_account = StateWithExtensionsOwned::<Account>::unpack(account.data).unwrap();
-        let amount = spl_token::ui_amount_to_amount(90.0, TEST_DECIMALS);
-        assert_eq!(token_account.base.amount, amount);
-        let account = config.rpc_client.get_account(&destination).await.unwrap();
-        let token_account = StateWithExtensionsOwned::<Account>::unpack(account.data).unwrap();
-        let amount = spl_token::ui_amount_to_amount(10.0, TEST_DECIMALS);
-        assert_eq!(token_account.base.amount, amount);
+        run_transfer_test(&config, payer).await;
     }
 }
 
@@ -4010,3 +4019,12 @@ async fn group(test_validator: &TestValidator, payer: &Keypair) {
     let extension = mint_state.get_extension::<TokenGroup>().unwrap();
     assert_eq!(extension.update_authority, Some(mint).try_into().unwrap());
 }
+
+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);
+        run_transfer_test(&config, payer).await;
+    }
+}