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::compute_budget_limits::MAX_COMPUTE_UNIT_LIMIT,
5    solana_compute_budget_interface::{self as compute_budget, ComputeBudgetInstruction},
6    solana_instruction::Instruction,
7    solana_message::Message,
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 => {
102            let compute_unit_limit =
103                simulate_for_compute_unit_limit_unchecked(rpc_client, message)?;
104
105            // Overwrite the compute unit limit instruction with the actual units consumed
106            message.instructions[compute_unit_limit_ix_index].data =
107                ComputeBudgetInstruction::set_compute_unit_limit(compute_unit_limit).data;
108
109            Ok(UpdateComputeUnitLimitResult::UpdatedInstructionIndex(
110                compute_unit_limit_ix_index,
111            ))
112        }
113        ComputeUnitLimit::Static(_) | ComputeUnitLimit::Default => {
114            Ok(UpdateComputeUnitLimitResult::SimulationNotConfigured)
115        }
116    }
117}
118
119pub(crate) struct ComputeUnitConfig {
120    pub(crate) compute_unit_price: Option<u64>,
121    pub(crate) compute_unit_limit: ComputeUnitLimit,
122}
123
124pub(crate) trait WithComputeUnitConfig {
125    fn with_compute_unit_config(self, config: &ComputeUnitConfig) -> Self;
126}
127
128impl WithComputeUnitConfig for Vec<Instruction> {
129    fn with_compute_unit_config(mut self, config: &ComputeUnitConfig) -> Self {
130        if let Some(compute_unit_price) = config.compute_unit_price {
131            self.push(ComputeBudgetInstruction::set_compute_unit_price(
132                compute_unit_price,
133            ));
134            match config.compute_unit_limit {
135                ComputeUnitLimit::Default => {}
136                ComputeUnitLimit::Static(compute_unit_limit) => {
137                    self.push(ComputeBudgetInstruction::set_compute_unit_limit(
138                        compute_unit_limit,
139                    ));
140                }
141                ComputeUnitLimit::Simulated => {
142                    // Default to the max compute unit limit because later transactions will be
143                    // simulated to get the exact compute units consumed.
144                    self.push(ComputeBudgetInstruction::set_compute_unit_limit(
145                        MAX_COMPUTE_UNIT_LIMIT,
146                    ));
147                }
148            }
149        }
150        self
151    }
152}
solana_cli/compute_budget.rs

solana_cli/
compute_budget.rs
