miraland_program_runtime/
compute_budget_processor.rs

1use {
2    crate::{
3        compute_budget::DEFAULT_HEAP_COST,
4        prioritization_fee::{PrioritizationFeeDetails, PrioritizationFeeType},
5    },
6    miraland_sdk::{
7        borsh1::try_from_slice_unchecked,
8        compute_budget::{self, ComputeBudgetInstruction},
9        entrypoint::HEAP_LENGTH as MIN_HEAP_FRAME_BYTES,
10        fee::FeeBudgetLimits,
11        instruction::{CompiledInstruction, InstructionError},
12        pubkey::Pubkey,
13        transaction::TransactionError,
14    },
15};
16
17const MAX_HEAP_FRAME_BYTES: u32 = 256 * 1024;
18pub const DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT: u32 = 200_000;
19pub const MAX_COMPUTE_UNIT_LIMIT: u32 = 1_400_000;
20
21/// The total accounts data a transaction can load is limited to 64MiB to not break
22/// anyone in Mainnet today. It can be set by set_loaded_accounts_data_size_limit instruction
23pub const MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES: u32 = 64 * 1024 * 1024;
24
25#[derive(Clone, Debug, PartialEq, Eq)]
26pub struct ComputeBudgetLimits {
27    pub updated_heap_bytes: u32,
28    pub compute_unit_limit: u32,
29    pub compute_unit_price: u64,
30    pub loaded_accounts_bytes: u32,
31}
32
33impl Default for ComputeBudgetLimits {
34    fn default() -> Self {
35        ComputeBudgetLimits {
36            updated_heap_bytes: u32::try_from(MIN_HEAP_FRAME_BYTES).unwrap(),
37            compute_unit_limit: MAX_COMPUTE_UNIT_LIMIT,
38            compute_unit_price: 0,
39            loaded_accounts_bytes: MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES,
40        }
41    }
42}
43
44impl From<ComputeBudgetLimits> for FeeBudgetLimits {
45    fn from(val: ComputeBudgetLimits) -> Self {
46        let prioritization_fee_details = PrioritizationFeeDetails::new(
47            PrioritizationFeeType::ComputeUnitPrice(val.compute_unit_price),
48            u64::from(val.compute_unit_limit),
49        );
50        let prioritization_fee = prioritization_fee_details.get_fee();
51
52        FeeBudgetLimits {
53            // NOTE - usize::from(u32).unwrap() may fail if target is 16-bit and
54            // `loaded_accounts_bytes` is greater than u16::MAX. In that case, panic is proper.
55            loaded_accounts_data_size_limit: usize::try_from(val.loaded_accounts_bytes).unwrap(),
56            heap_cost: DEFAULT_HEAP_COST,
57            compute_unit_limit: u64::from(val.compute_unit_limit),
58            prioritization_fee,
59        }
60    }
61}
62
63/// Processing compute_budget could be part of tx sanitizing, failed to process
64/// these instructions will drop the transaction eventually without execution,
65/// may as well fail it early.
66/// If succeeded, the transaction's specific limits/requests (could be default)
67/// are retrieved and returned,
68pub fn process_compute_budget_instructions<'a>(
69    instructions: impl Iterator<Item = (&'a Pubkey, &'a CompiledInstruction)>,
70) -> Result<ComputeBudgetLimits, TransactionError> {
71    let mut num_non_compute_budget_instructions: u32 = 0;
72    let mut updated_compute_unit_limit = None;
73    let mut updated_compute_unit_price = None;
74    let mut requested_heap_size = None;
75    let mut updated_loaded_accounts_data_size_limit = None;
76
77    for (i, (program_id, instruction)) in instructions.enumerate() {
78        if compute_budget::check_id(program_id) {
79            let invalid_instruction_data_error = TransactionError::InstructionError(
80                i as u8,
81                InstructionError::InvalidInstructionData,
82            );
83            let duplicate_instruction_error = TransactionError::DuplicateInstruction(i as u8);
84
85            match try_from_slice_unchecked(&instruction.data) {
86                Ok(ComputeBudgetInstruction::RequestHeapFrame(bytes)) => {
87                    if requested_heap_size.is_some() {
88                        return Err(duplicate_instruction_error);
89                    }
90                    if sanitize_requested_heap_size(bytes) {
91                        requested_heap_size = Some(bytes);
92                    } else {
93                        return Err(invalid_instruction_data_error);
94                    }
95                }
96                Ok(ComputeBudgetInstruction::SetComputeUnitLimit(compute_unit_limit)) => {
97                    if updated_compute_unit_limit.is_some() {
98                        return Err(duplicate_instruction_error);
99                    }
100                    updated_compute_unit_limit = Some(compute_unit_limit);
101                }
102                Ok(ComputeBudgetInstruction::SetComputeUnitPrice(micro_lamports)) => {
103                    if updated_compute_unit_price.is_some() {
104                        return Err(duplicate_instruction_error);
105                    }
106                    updated_compute_unit_price = Some(micro_lamports);
107                }
108                Ok(ComputeBudgetInstruction::SetLoadedAccountsDataSizeLimit(bytes)) => {
109                    if updated_loaded_accounts_data_size_limit.is_some() {
110                        return Err(duplicate_instruction_error);
111                    }
112                    updated_loaded_accounts_data_size_limit = Some(bytes);
113                }
114                _ => return Err(invalid_instruction_data_error),
115            }
116        } else {
117            // only include non-request instructions in default max calc
118            num_non_compute_budget_instructions =
119                num_non_compute_budget_instructions.saturating_add(1);
120        }
121    }
122
123    // sanitize limits
124    let updated_heap_bytes = requested_heap_size
125        .unwrap_or(u32::try_from(MIN_HEAP_FRAME_BYTES).unwrap()) // loader's default heap_size
126        .min(MAX_HEAP_FRAME_BYTES);
127
128    let compute_unit_limit = updated_compute_unit_limit
129        .unwrap_or_else(|| {
130            num_non_compute_budget_instructions
131                .saturating_mul(DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT)
132        })
133        .min(MAX_COMPUTE_UNIT_LIMIT);
134
135    let compute_unit_price = updated_compute_unit_price.unwrap_or(0);
136
137    let loaded_accounts_bytes = updated_loaded_accounts_data_size_limit
138        .unwrap_or(MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES)
139        .min(MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES);
140
141    Ok(ComputeBudgetLimits {
142        updated_heap_bytes,
143        compute_unit_limit,
144        compute_unit_price,
145        loaded_accounts_bytes,
146    })
147}
148
149fn sanitize_requested_heap_size(bytes: u32) -> bool {
150    (u32::try_from(MIN_HEAP_FRAME_BYTES).unwrap()..=MAX_HEAP_FRAME_BYTES).contains(&bytes)
151        && bytes % 1024 == 0
152}
153
154#[cfg(test)]
155mod tests {
156    use {
157        super::*,
158        miraland_sdk::{
159            hash::Hash,
160            instruction::Instruction,
161            message::Message,
162            pubkey::Pubkey,
163            signature::Keypair,
164            signer::Signer,
165            system_instruction::{self},
166            transaction::{SanitizedTransaction, Transaction},
167        },
168    };
169
170    macro_rules! test {
171        ( $instructions: expr, $expected_result: expr) => {
172            let payer_keypair = Keypair::new();
173            let tx = SanitizedTransaction::from_transaction_for_tests(Transaction::new(
174                &[&payer_keypair],
175                Message::new($instructions, Some(&payer_keypair.pubkey())),
176                Hash::default(),
177            ));
178            let result =
179                process_compute_budget_instructions(tx.message().program_instructions_iter());
180            assert_eq!($expected_result, result);
181        };
182    }
183
184    #[test]
185    fn test_process_instructions() {
186        // Units
187        test!(
188            &[],
189            Ok(ComputeBudgetLimits {
190                compute_unit_limit: 0,
191                ..ComputeBudgetLimits::default()
192            })
193        );
194        test!(
195            &[
196                ComputeBudgetInstruction::set_compute_unit_limit(1),
197                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
198            ],
199            Ok(ComputeBudgetLimits {
200                compute_unit_limit: 1,
201                ..ComputeBudgetLimits::default()
202            })
203        );
204        test!(
205            &[
206                ComputeBudgetInstruction::set_compute_unit_limit(MAX_COMPUTE_UNIT_LIMIT + 1),
207                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
208            ],
209            Ok(ComputeBudgetLimits {
210                compute_unit_limit: MAX_COMPUTE_UNIT_LIMIT,
211                ..ComputeBudgetLimits::default()
212            })
213        );
214        test!(
215            &[
216                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
217                ComputeBudgetInstruction::set_compute_unit_limit(MAX_COMPUTE_UNIT_LIMIT),
218            ],
219            Ok(ComputeBudgetLimits {
220                compute_unit_limit: MAX_COMPUTE_UNIT_LIMIT,
221                ..ComputeBudgetLimits::default()
222            })
223        );
224        test!(
225            &[
226                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
227                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
228                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
229                ComputeBudgetInstruction::set_compute_unit_limit(1),
230            ],
231            Ok(ComputeBudgetLimits {
232                compute_unit_limit: 1,
233                ..ComputeBudgetLimits::default()
234            })
235        );
236        test!(
237            &[
238                ComputeBudgetInstruction::set_compute_unit_limit(1),
239                ComputeBudgetInstruction::set_compute_unit_price(42)
240            ],
241            Ok(ComputeBudgetLimits {
242                compute_unit_limit: 1,
243                compute_unit_price: 42,
244                ..ComputeBudgetLimits::default()
245            })
246        );
247
248        // HeapFrame
249        test!(
250            &[],
251            Ok(ComputeBudgetLimits {
252                compute_unit_limit: 0,
253                ..ComputeBudgetLimits::default()
254            })
255        );
256        test!(
257            &[
258                ComputeBudgetInstruction::request_heap_frame(40 * 1024),
259                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
260            ],
261            Ok(ComputeBudgetLimits {
262                compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT,
263                updated_heap_bytes: 40 * 1024,
264                ..ComputeBudgetLimits::default()
265            })
266        );
267        test!(
268            &[
269                ComputeBudgetInstruction::request_heap_frame(40 * 1024 + 1),
270                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
271            ],
272            Err(TransactionError::InstructionError(
273                0,
274                InstructionError::InvalidInstructionData,
275            ))
276        );
277        test!(
278            &[
279                ComputeBudgetInstruction::request_heap_frame(31 * 1024),
280                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
281            ],
282            Err(TransactionError::InstructionError(
283                0,
284                InstructionError::InvalidInstructionData,
285            ))
286        );
287        test!(
288            &[
289                ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES + 1),
290                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
291            ],
292            Err(TransactionError::InstructionError(
293                0,
294                InstructionError::InvalidInstructionData,
295            ))
296        );
297        test!(
298            &[
299                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
300                ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES),
301            ],
302            Ok(ComputeBudgetLimits {
303                compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT,
304                updated_heap_bytes: MAX_HEAP_FRAME_BYTES,
305                ..ComputeBudgetLimits::default()
306            })
307        );
308        test!(
309            &[
310                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
311                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
312                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
313                ComputeBudgetInstruction::request_heap_frame(1),
314            ],
315            Err(TransactionError::InstructionError(
316                3,
317                InstructionError::InvalidInstructionData,
318            ))
319        );
320        test!(
321            &[
322                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
323                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
324                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
325                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
326                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
327                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
328                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
329                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
330            ],
331            Ok(ComputeBudgetLimits {
332                compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT * 7,
333                ..ComputeBudgetLimits::default()
334            })
335        );
336
337        // Combined
338        test!(
339            &[
340                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
341                ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES),
342                ComputeBudgetInstruction::set_compute_unit_limit(MAX_COMPUTE_UNIT_LIMIT),
343                ComputeBudgetInstruction::set_compute_unit_price(u64::MAX),
344            ],
345            Ok(ComputeBudgetLimits {
346                compute_unit_price: u64::MAX,
347                compute_unit_limit: MAX_COMPUTE_UNIT_LIMIT,
348                updated_heap_bytes: MAX_HEAP_FRAME_BYTES,
349                ..ComputeBudgetLimits::default()
350            })
351        );
352        test!(
353            &[
354                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
355                ComputeBudgetInstruction::set_compute_unit_limit(1),
356                ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES),
357                ComputeBudgetInstruction::set_compute_unit_price(u64::MAX),
358            ],
359            Ok(ComputeBudgetLimits {
360                compute_unit_price: u64::MAX,
361                compute_unit_limit: 1,
362                updated_heap_bytes: MAX_HEAP_FRAME_BYTES,
363                ..ComputeBudgetLimits::default()
364            })
365        );
366
367        // Duplicates
368        test!(
369            &[
370                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
371                ComputeBudgetInstruction::set_compute_unit_limit(MAX_COMPUTE_UNIT_LIMIT),
372                ComputeBudgetInstruction::set_compute_unit_limit(MAX_COMPUTE_UNIT_LIMIT - 1),
373            ],
374            Err(TransactionError::DuplicateInstruction(2))
375        );
376
377        test!(
378            &[
379                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
380                ComputeBudgetInstruction::request_heap_frame(MIN_HEAP_FRAME_BYTES as u32),
381                ComputeBudgetInstruction::request_heap_frame(MAX_HEAP_FRAME_BYTES),
382            ],
383            Err(TransactionError::DuplicateInstruction(2))
384        );
385        test!(
386            &[
387                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
388                ComputeBudgetInstruction::set_compute_unit_price(0),
389                ComputeBudgetInstruction::set_compute_unit_price(u64::MAX),
390            ],
391            Err(TransactionError::DuplicateInstruction(2))
392        );
393    }
394
395    #[test]
396    fn test_process_loaded_accounts_data_size_limit_instruction() {
397        test!(
398            &[],
399            Ok(ComputeBudgetLimits {
400                compute_unit_limit: 0,
401                ..ComputeBudgetLimits::default()
402            })
403        );
404
405        // Assert when set_loaded_accounts_data_size_limit presents,
406        // budget is set with data_size
407        let data_size = 1;
408        let expected_result = Ok(ComputeBudgetLimits {
409            compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT,
410            loaded_accounts_bytes: data_size,
411            ..ComputeBudgetLimits::default()
412        });
413
414        test!(
415            &[
416                ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(data_size),
417                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
418            ],
419            expected_result
420        );
421
422        // Assert when set_loaded_accounts_data_size_limit presents, with greater than max value
423        // budget is set to max data size
424        let data_size = MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES + 1;
425        let expected_result = Ok(ComputeBudgetLimits {
426            compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT,
427            loaded_accounts_bytes: MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES,
428            ..ComputeBudgetLimits::default()
429        });
430
431        test!(
432            &[
433                ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(data_size),
434                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
435            ],
436            expected_result
437        );
438
439        // Assert when set_loaded_accounts_data_size_limit is not presented
440        // budget is set to default data size
441        let expected_result = Ok(ComputeBudgetLimits {
442            compute_unit_limit: DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT,
443            loaded_accounts_bytes: MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES,
444            ..ComputeBudgetLimits::default()
445        });
446
447        test!(
448            &[Instruction::new_with_bincode(
449                Pubkey::new_unique(),
450                &0_u8,
451                vec![]
452            ),],
453            expected_result
454        );
455
456        // Assert when set_loaded_accounts_data_size_limit presents more than once,
457        // return DuplicateInstruction
458        let data_size = MAX_LOADED_ACCOUNTS_DATA_SIZE_BYTES;
459        let expected_result = Err(TransactionError::DuplicateInstruction(2));
460
461        test!(
462            &[
463                Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
464                ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(data_size),
465                ComputeBudgetInstruction::set_loaded_accounts_data_size_limit(data_size),
466            ],
467            expected_result
468        );
469    }
470
471    #[test]
472    fn test_process_mixed_instructions_without_compute_budget() {
473        let payer_keypair = Keypair::new();
474
475        let transaction =
476            SanitizedTransaction::from_transaction_for_tests(Transaction::new_signed_with_payer(
477                &[
478                    Instruction::new_with_bincode(Pubkey::new_unique(), &0_u8, vec![]),
479                    system_instruction::transfer(&payer_keypair.pubkey(), &Pubkey::new_unique(), 2),
480                ],
481                Some(&payer_keypair.pubkey()),
482                &[&payer_keypair],
483                Hash::default(),
484            ));
485
486        let result =
487            process_compute_budget_instructions(transaction.message().program_instructions_iter());
488
489        // assert process_instructions will be successful with default,
490        // and the default compute_unit_limit is 2 times default: one for bpf ix, one for
491        // builtin ix.
492        assert_eq!(
493            result,
494            Ok(ComputeBudgetLimits {
495                compute_unit_limit: 2 * DEFAULT_INSTRUCTION_COMPUTE_UNIT_LIMIT,
496                ..ComputeBudgetLimits::default()
497            })
498        );
499    }
500}