From 5817309b65639a816fff0ca4839ae380838973dc Mon Sep 17 00:00:00 2001 From: Jon C Date: Tue, 26 Mar 2024 01:58:31 +0100 Subject: [PATCH] token-client: Add compute unit price and limit fields --- token/client/src/token.rs | 68 ++++++++----------- .../tests/confidential_transfer.rs | 3 + .../tests/confidential_transfer_fee.rs | 5 ++ 3 files changed, 38 insertions(+), 38 deletions(-) diff --git a/token/client/src/token.rs b/token/client/src/token.rs index b77044d4616..a418173ea6c 100644 --- a/token/client/src/token.rs +++ b/token/client/src/token.rs @@ -8,6 +8,7 @@ use { solana_program_test::tokio::time, solana_sdk::{ account::Account as BaseAccount, + compute_budget::ComputeBudgetInstruction, hash::Hash, instruction::{AccountMeta, Instruction}, message::Message, @@ -339,6 +340,8 @@ pub struct Token { nonce_blockhash: Option, memo: Arc>>, transfer_hook_accounts: Option>, + compute_unit_price: Option, + compute_unit_limit: Option, } impl fmt::Debug for Token { @@ -356,6 +359,8 @@ impl fmt::Debug for Token { .field("nonce_blockhash", &self.nonce_blockhash) .field("memo", &self.memo.read().unwrap()) .field("transfer_hook_accounts", &self.transfer_hook_accounts) + .field("compute_unit_price", &self.compute_unit_price) + .field("compute_unit_limit", &self.compute_unit_limit) .finish() } } @@ -402,6 +407,8 @@ where nonce_blockhash: None, memo: Arc::new(RwLock::new(None)), transfer_hook_accounts: None, + compute_unit_price: None, + compute_unit_limit: None, } } @@ -451,6 +458,16 @@ where self } + pub fn with_compute_unit_price(mut self, compute_unit_price: u64) -> Self { + self.compute_unit_price = Some(compute_unit_price); + self + } + + pub fn with_compute_unit_limit(mut self, compute_unit_limit: u32) -> Self { + self.compute_unit_limit = Some(compute_unit_limit); + self + } + pub fn with_memo>(&self, memo: M, signers: Vec) -> &Self { let mut w_memo = self.memo.write().unwrap(); *w_memo = Some(TokenMemo { @@ -508,7 +525,6 @@ where async fn construct_tx( &self, token_instructions: &[Instruction], - additional_compute_budget: Option, signing_keypairs: &S, ) -> TokenResult { let mut instructions = vec![]; @@ -533,12 +549,16 @@ where instructions.extend_from_slice(token_instructions); - if let Some(additional_compute_budget) = additional_compute_budget { - instructions.push( - solana_sdk::compute_budget::ComputeBudgetInstruction::set_compute_unit_limit( - additional_compute_budget, - ), - ); + 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 (message, blockhash) = @@ -598,7 +618,7 @@ where signing_keypairs: &S, ) -> TokenResult { let transaction = self - .construct_tx(token_instructions, None, signing_keypairs) + .construct_tx(token_instructions, signing_keypairs) .await?; self.client @@ -613,27 +633,7 @@ where signing_keypairs: &S, ) -> TokenResult { let transaction = self - .construct_tx(token_instructions, None, signing_keypairs) - .await?; - - self.client - .send_transaction(&transaction) - .await - .map_err(TokenError::Client) - } - - pub async fn process_ixs_with_additional_compute_budget( - &self, - token_instructions: &[Instruction], - additional_compute_budget: u32, - signing_keypairs: &S, - ) -> TokenResult { - let transaction = self - .construct_tx( - token_instructions, - Some(additional_compute_budget), - signing_keypairs, - ) + .construct_tx(token_instructions, signing_keypairs) .await?; self.client @@ -2728,9 +2728,6 @@ where .new_decryptable_available_balance(transfer_amount, source_aes_key) .map_err(|_| TokenError::AccountDecryption)?; - // additional compute budget required for `VerifyTransferWithFee` - const TRANSFER_WITH_FEE_COMPUTE_BUDGET: u32 = 500_000; - let mut instructions = confidential_transfer::instruction::transfer_with_fee( &self.program_id, source_account, @@ -2756,12 +2753,7 @@ where ) .await .map_err(|_| TokenError::AccountNotFound)?; - self.process_ixs_with_additional_compute_budget( - &instructions, - TRANSFER_WITH_FEE_COMPUTE_BUDGET, - signing_keypairs, - ) - .await + self.process_ixs(&instructions, signing_keypairs).await } /// Transfer tokens confidentially with fee using split proofs. diff --git a/token/program-2022-test/tests/confidential_transfer.rs b/token/program-2022-test/tests/confidential_transfer.rs index 91a5e7fafac..a395b6acf55 100644 --- a/token/program-2022-test/tests/confidential_transfer.rs +++ b/token/program-2022-test/tests/confidential_transfer.rs @@ -1110,6 +1110,7 @@ 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, @@ -2523,6 +2524,7 @@ async fn confidential_transfer_transfer_with_split_proof_contexts_in_parallel() &range_proof_context_state_account, &context_state_authority, ]; + let token = token.with_compute_unit_limit(500_000); token .confidential_transfer_transfer_with_split_proofs_in_parallel( &alice_meta.token_account, @@ -2941,6 +2943,7 @@ async fn confidential_transfer_transfer_with_fee_and_split_proof_context_in_para &range_proof_context_state_account, &context_state_authority, ]; + let token = token.with_compute_unit_limit(500_000); token .confidential_transfer_transfer_with_fee_and_split_proofs_in_parallel( &alice_meta.token_account, diff --git a/token/program-2022-test/tests/confidential_transfer_fee.rs b/token/program-2022-test/tests/confidential_transfer_fee.rs index f2eb4086eb7..3fd2c8b0278 100644 --- a/token/program-2022-test/tests/confidential_transfer_fee.rs +++ b/token/program-2022-test/tests/confidential_transfer_fee.rs @@ -513,6 +513,7 @@ 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, @@ -671,6 +672,7 @@ 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, @@ -801,6 +803,7 @@ 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, @@ -969,6 +972,7 @@ 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, @@ -1159,6 +1163,7 @@ 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,