Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

token-client: Add compute unit price and limit fields #6493

Merged
merged 1 commit into from
Mar 26, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 30 additions & 38 deletions token/client/src/token.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -339,6 +340,8 @@ pub struct Token<T> {
nonce_blockhash: Option<Hash>,
memo: Arc<RwLock<Option<TokenMemo>>>,
transfer_hook_accounts: Option<Vec<AccountMeta>>,
compute_unit_price: Option<u64>,
compute_unit_limit: Option<u32>,
}

impl<T> fmt::Debug for Token<T> {
Expand All @@ -356,6 +359,8 @@ impl<T> fmt::Debug for Token<T> {
.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()
}
}
Expand Down Expand Up @@ -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,
}
}

Expand Down Expand Up @@ -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<M: AsRef<str>>(&self, memo: M, signers: Vec<Pubkey>) -> &Self {
let mut w_memo = self.memo.write().unwrap();
*w_memo = Some(TokenMemo {
Expand Down Expand Up @@ -508,7 +525,6 @@ where
async fn construct_tx<S: Signers>(
&self,
token_instructions: &[Instruction],
additional_compute_budget: Option<u32>,
signing_keypairs: &S,
) -> TokenResult<Transaction> {
let mut instructions = vec![];
Expand All @@ -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) =
Expand Down Expand Up @@ -598,7 +618,7 @@ where
signing_keypairs: &S,
) -> TokenResult<T::SimulationOutput> {
let transaction = self
.construct_tx(token_instructions, None, signing_keypairs)
.construct_tx(token_instructions, signing_keypairs)
.await?;

self.client
Expand All @@ -613,27 +633,7 @@ where
signing_keypairs: &S,
) -> TokenResult<T::Output> {
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<S: Signers>(
&self,
token_instructions: &[Instruction],
additional_compute_budget: u32,
signing_keypairs: &S,
) -> TokenResult<T::Output> {
let transaction = self
.construct_tx(
token_instructions,
Some(additional_compute_budget),
signing_keypairs,
)
.construct_tx(token_instructions, signing_keypairs)
.await?;

self.client
Expand Down Expand Up @@ -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,
Expand All @@ -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.
Expand Down
3 changes: 3 additions & 0 deletions token/program-2022-test/tests/confidential_transfer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
5 changes: 5 additions & 0 deletions token/program-2022-test/tests/confidential_transfer_fee.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
Loading