#![cfg(feature = "full")]
use crate::{
entrypoint::HEAP_LENGTH as MIN_HEAP_FRAME_BYTES,
feature_set::{requestable_heap_size, FeatureSet},
process_instruction::BpfComputeBudget,
transaction::{Transaction, TransactionError},
};
use borsh::{BorshDeserialize, BorshSchema, BorshSerialize};
use solana_sdk::{
borsh::try_from_slice_unchecked,
instruction::{Instruction, InstructionError},
};
use std::sync::Arc;
crate::declare_id!("ComputeBudget111111111111111111111111111111");
const MAX_UNITS: u32 = 1_000_000;
const MAX_HEAP_FRAME_BYTES: u32 = 256 * 1024;
#[derive(
Serialize,
Deserialize,
BorshSerialize,
BorshDeserialize,
BorshSchema,
Debug,
Clone,
PartialEq,
AbiExample,
AbiEnumVisitor,
)]
pub enum ComputeBudgetInstruction {
RequestUnits(u32),
RequestHeapFrame(u32),
}
pub fn request_units(units: u32) -> Instruction {
Instruction::new_with_borsh(id(), &ComputeBudgetInstruction::RequestUnits(units), vec![])
}
pub fn request_heap_frame(bytes: u32) -> Instruction {
Instruction::new_with_borsh(
id(),
&ComputeBudgetInstruction::RequestHeapFrame(bytes),
vec![],
)
}
pub fn process_request(
compute_budget: &mut BpfComputeBudget,
tx: &Transaction,
feature_set: Arc<FeatureSet>,
) -> Result<(), TransactionError> {
let error = TransactionError::InstructionError(0, InstructionError::InvalidInstructionData);
for instruction in tx.message().instructions.iter().take(3) {
if check_id(instruction.program_id(&tx.message().account_keys)) {
match try_from_slice_unchecked(&instruction.data) {
Ok(ComputeBudgetInstruction::RequestUnits(units)) => {
if units > MAX_UNITS {
return Err(error);
}
compute_budget.max_units = units as u64;
}
Ok(ComputeBudgetInstruction::RequestHeapFrame(bytes)) => {
if !feature_set.is_active(&requestable_heap_size::id())
|| bytes > MAX_HEAP_FRAME_BYTES
|| bytes < MIN_HEAP_FRAME_BYTES as u32
|| bytes % 1024 != 0
{
return Err(error);
}
compute_budget.heap_size = Some(bytes as usize);
}
_ => return Err(error),
}
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{hash::Hash, message::Message, pubkey::Pubkey, signature::Keypair, signer::Signer};
macro_rules! test {
( $instructions: expr, $expected_error: expr, $expected_budget: expr ) => {
let payer_keypair = Keypair::new();
let tx = Transaction::new(
&[&payer_keypair],
Message::new($instructions, Some(&payer_keypair.pubkey())),
Hash::default(),
);
let feature_set = Arc::new(FeatureSet::all_enabled());
let mut compute_budget = BpfComputeBudget::default();
let result = process_request(&mut compute_budget, &tx, feature_set);
assert_eq!($expected_error as Result<(), TransactionError>, result);
assert_eq!(compute_budget, $expected_budget);
};
}
#[test]
fn test_process_request() {
test!(&[], Ok(()), BpfComputeBudget::default());
test!(
&[
request_units(1),
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
],
Ok(()),
BpfComputeBudget {
max_units: 1,
..BpfComputeBudget::default()
}
);
test!(
&[
request_units(MAX_UNITS + 1),
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
],
Err(TransactionError::InstructionError(
0,
InstructionError::InvalidInstructionData,
)),
BpfComputeBudget::default()
);
test!(
&[
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
request_units(MAX_UNITS),
],
Ok(()),
BpfComputeBudget {
max_units: MAX_UNITS as u64,
..BpfComputeBudget::default()
}
);
test!(
&[
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
request_units(1),
],
Ok(()),
BpfComputeBudget::default()
);
test!(&[], Ok(()), BpfComputeBudget::default());
test!(
&[
request_heap_frame(40 * 1024),
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
],
Ok(()),
BpfComputeBudget {
heap_size: Some(40 * 1024),
..BpfComputeBudget::default()
}
);
test!(
&[
request_heap_frame(40 * 1024 + 1),
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
],
Err(TransactionError::InstructionError(
0,
InstructionError::InvalidInstructionData,
)),
BpfComputeBudget::default()
);
test!(
&[
request_heap_frame(31 * 1024),
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
],
Err(TransactionError::InstructionError(
0,
InstructionError::InvalidInstructionData,
)),
BpfComputeBudget::default()
);
test!(
&[
request_heap_frame(MAX_HEAP_FRAME_BYTES + 1),
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
],
Err(TransactionError::InstructionError(
0,
InstructionError::InvalidInstructionData,
)),
BpfComputeBudget::default()
);
test!(
&[
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
request_heap_frame(MAX_HEAP_FRAME_BYTES),
],
Ok(()),
BpfComputeBudget {
heap_size: Some(MAX_HEAP_FRAME_BYTES as usize),
..BpfComputeBudget::default()
}
);
test!(
&[
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
request_heap_frame(1), ],
Ok(()),
BpfComputeBudget::default()
);
test!(
&[
Instruction::new_with_bincode(Pubkey::new_unique(), &0, vec![]),
request_heap_frame(MAX_HEAP_FRAME_BYTES),
request_units(MAX_UNITS),
],
Ok(()),
BpfComputeBudget {
max_units: MAX_UNITS as u64,
heap_size: Some(MAX_HEAP_FRAME_BYTES as usize),
..BpfComputeBudget::default()
}
);
}
}