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::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.
42fn 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        .value;
57
58    // Bail if the simulated transaction failed
59    if let Some(err) = simulate_result.err {
60        return Err(err.into());
61    }
62
63    let units_consumed = simulate_result
64        .units_consumed
65        .expect("compute units unavailable");
66
67    u32::try_from(units_consumed).map_err(Into::into)
68}
69
70/// Returns the compute unit limit used during simulation
71///
72/// Returns an error if the message does not contain a compute unit limit
73/// instruction or if the simulation fails.
74pub(crate) fn simulate_for_compute_unit_limit(
75    rpc_client: &RpcClient,
76    message: &Message,
77) -> Result<u32, Box<dyn std::error::Error>> {
78    if get_compute_unit_limit_instruction_index(message).is_none() {
79        return Err("No compute unit limit instruction found".into());
80    }
81    simulate_for_compute_unit_limit_unchecked(rpc_client, message)
82}
83
84/// Simulates a message and returns the index of the compute unit limit
85/// instruction
86///
87/// If the message does not contain a compute unit limit instruction, or if
88/// simulation was not configured, then the function will not simulate the
89/// message.
90pub(crate) fn simulate_and_update_compute_unit_limit(
91    compute_unit_limit: &ComputeUnitLimit,
92    rpc_client: &RpcClient,
93    message: &mut Message,
94) -> Result<UpdateComputeUnitLimitResult, Box<dyn std::error::Error>> {
95    let Some(compute_unit_limit_ix_index) = get_compute_unit_limit_instruction_index(message)
96    else {
97        return Ok(UpdateComputeUnitLimitResult::NoInstructionFound);
98    };
99
100    match compute_unit_limit {
101        ComputeUnitLimit::Simulated | ComputeUnitLimit::SimulatedWithExtraPercentage(_) => {
102            let base_compute_unit_limit =
103                simulate_for_compute_unit_limit_unchecked(rpc_client, message)?;
104
105            let compute_unit_limit =
106                if let ComputeUnitLimit::SimulatedWithExtraPercentage(n) = compute_unit_limit {
107                    (base_compute_unit_limit as u64)
108                        .saturating_mul(100_u64.saturating_add(*n as u64))
109                        .saturating_div(100) as u32
110                } else {
111                    base_compute_unit_limit
112                };
113
114            // Overwrite the compute unit limit instruction with the actual units consumed
115            message.instructions[compute_unit_limit_ix_index].data =
116                ComputeBudgetInstruction::set_compute_unit_limit(compute_unit_limit).data;
117
118            Ok(UpdateComputeUnitLimitResult::UpdatedInstructionIndex(
119                compute_unit_limit_ix_index,
120            ))
121        }
122        ComputeUnitLimit::Static(_) | ComputeUnitLimit::Default => {
123            Ok(UpdateComputeUnitLimitResult::SimulationNotConfigured)
124        }
125    }
126}
127
128pub(crate) struct ComputeUnitConfig {
129    pub(crate) compute_unit_price: Option<u64>,
130    pub(crate) compute_unit_limit: ComputeUnitLimit,
131}
132
133pub(crate) trait WithComputeUnitConfig {
134    fn with_compute_unit_config(self, config: &ComputeUnitConfig) -> Self;
135}
136
137impl WithComputeUnitConfig for Vec<Instruction> {
138    fn with_compute_unit_config(mut self, config: &ComputeUnitConfig) -> Self {
139        if let Some(compute_unit_price) = config.compute_unit_price {
140            self.push(ComputeBudgetInstruction::set_compute_unit_price(
141                compute_unit_price,
142            ));
143            match config.compute_unit_limit {
144                ComputeUnitLimit::Default => {}
145                ComputeUnitLimit::Static(compute_unit_limit) => {
146                    self.push(ComputeBudgetInstruction::set_compute_unit_limit(
147                        compute_unit_limit,
148                    ));
149                }
150                ComputeUnitLimit::Simulated | ComputeUnitLimit::SimulatedWithExtraPercentage(_) => {
151                    // Default to the max compute unit limit because later transactions will be
152                    // simulated to get the exact compute units consumed.
153                    self.push(ComputeBudgetInstruction::set_compute_unit_limit(
154                        MAX_COMPUTE_UNIT_LIMIT,
155                    ));
156                }
157            }
158        }
159        self
160    }
161}
solana_cli/compute_budget.rs

solana_cli/
compute_budget.rs
