solana_cli/
compute_budget.rs

1use {
2    solana_borsh::v1::try_from_slice_unchecked,
3    solana_clap_utils::compute_budget::ComputeUnitLimit,
4    solana_compute_budget_interface::{self as compute_budget, ComputeBudgetInstruction},
5    solana_instruction::Instruction,
6    solana_message::Message,
7    solana_program_runtime::execution_budget::MAX_COMPUTE_UNIT_LIMIT,
8    solana_rpc_client::nonblocking::rpc_client::RpcClient,
9    solana_rpc_client_api::config::RpcSimulateTransactionConfig,
10    solana_transaction::Transaction,
11};
12
13/// Enum capturing the possible results of updating a message based on the
14/// compute unit limits consumed during simulation.
15pub(crate) enum UpdateComputeUnitLimitResult {
16    UpdatedInstructionIndex(usize),
17    NoInstructionFound,
18    SimulationNotConfigured,
19}
20
21fn get_compute_unit_limit_instruction_index(message: &Message) -> Option<usize> {
22    message
23        .instructions
24        .iter()
25        .enumerate()
26        .find_map(|(ix_index, instruction)| {
27            let ix_program_id = message.program_id(ix_index)?;
28            if ix_program_id != &compute_budget::id() {
29                return None;
30            }
31
32            matches!(
33                try_from_slice_unchecked(&instruction.data),
34                Ok(ComputeBudgetInstruction::SetComputeUnitLimit(_))
35            )
36            .then_some(ix_index)
37        })
38}
39
40/// Like `simulate_for_compute_unit_limit`, but does not check that the message
41/// contains a compute unit limit instruction.
42async fn simulate_for_compute_unit_limit_unchecked(
43    rpc_client: &RpcClient,
44    message: &Message,
45) -> Result<u32, Box<dyn std::error::Error>> {
46    let transaction = Transaction::new_unsigned(message.clone());
47    let simulate_result = rpc_client
48        .simulate_transaction_with_config(
49            &transaction,
50            RpcSimulateTransactionConfig {
51                replace_recent_blockhash: true,
52                commitment: Some(rpc_client.commitment()),
53                ..RpcSimulateTransactionConfig::default()
54            },
55        )
56        .await?
57        .value;
58
59    // Bail if the simulated transaction failed
60    if let Some(err) = simulate_result.err {
61        return Err(err.into());
62    }
63
64    let units_consumed = simulate_result
65        .units_consumed
66        .expect("compute units unavailable");
67
68    u32::try_from(units_consumed).map_err(Into::into)
69}
70
71/// Simulates a message and returns the index of the compute unit limit
72/// instruction
73///
74/// If the message does not contain a compute unit limit instruction, or if
75/// simulation was not configured, then the function will not simulate the
76/// message.
77pub(crate) async fn simulate_and_update_compute_unit_limit(
78    compute_unit_limit: &ComputeUnitLimit,
79    rpc_client: &RpcClient,
80    message: &mut Message,
81) -> Result<UpdateComputeUnitLimitResult, Box<dyn std::error::Error>> {
82    let Some(compute_unit_limit_ix_index) = get_compute_unit_limit_instruction_index(message)
83    else {
84        return Ok(UpdateComputeUnitLimitResult::NoInstructionFound);
85    };
86
87    match compute_unit_limit {
88        ComputeUnitLimit::Simulated | ComputeUnitLimit::SimulatedWithExtraPercentage(_) => {
89            let base_compute_unit_limit =
90                simulate_for_compute_unit_limit_unchecked(rpc_client, message).await?;
91
92            let compute_unit_limit =
93                if let ComputeUnitLimit::SimulatedWithExtraPercentage(n) = compute_unit_limit {
94                    (base_compute_unit_limit as u64)
95                        .saturating_mul(100_u64.saturating_add(*n as u64))
96                        .saturating_div(100) as u32
97                } else {
98                    base_compute_unit_limit
99                };
100
101            // Overwrite the compute unit limit instruction with the actual units consumed
102            message.instructions[compute_unit_limit_ix_index].data =
103                ComputeBudgetInstruction::set_compute_unit_limit(compute_unit_limit).data;
104
105            Ok(UpdateComputeUnitLimitResult::UpdatedInstructionIndex(
106                compute_unit_limit_ix_index,
107            ))
108        }
109        ComputeUnitLimit::Static(_) | ComputeUnitLimit::Default => {
110            Ok(UpdateComputeUnitLimitResult::SimulationNotConfigured)
111        }
112    }
113}
114
115pub(crate) struct ComputeUnitConfig {
116    pub(crate) compute_unit_price: Option<u64>,
117    pub(crate) compute_unit_limit: ComputeUnitLimit,
118}
119
120pub(crate) trait WithComputeUnitConfig {
121    fn with_compute_unit_config(self, config: &ComputeUnitConfig) -> Self;
122}
123
124impl WithComputeUnitConfig for Vec<Instruction> {
125    fn with_compute_unit_config(mut self, config: &ComputeUnitConfig) -> Self {
126        if let Some(compute_unit_price) = config.compute_unit_price {
127            self.push(ComputeBudgetInstruction::set_compute_unit_price(
128                compute_unit_price,
129            ));
130            match config.compute_unit_limit {
131                ComputeUnitLimit::Default => {}
132                ComputeUnitLimit::Static(compute_unit_limit) => {
133                    self.push(ComputeBudgetInstruction::set_compute_unit_limit(
134                        compute_unit_limit,
135                    ));
136                }
137                ComputeUnitLimit::Simulated | ComputeUnitLimit::SimulatedWithExtraPercentage(_) => {
138                    // Default to the max compute unit limit because later transactions will be
139                    // simulated to get the exact compute units consumed.
140                    self.push(ComputeBudgetInstruction::set_compute_unit_limit(
141                        MAX_COMPUTE_UNIT_LIMIT,
142                    ));
143                }
144            }
145        }
146        self
147    }
148}
solana_cli/compute_budget.rs

solana_cli/
compute_budget.rs
