directly use compute_budget MAX_UNITS and DEFAULT_UNITS

This commit is contained in:
Tao Zhu
2022-03-16 10:34:15 -05:00
committed by Tao Zhu
parent a4cacf3389
commit 0ed23899e7
7 changed files with 44 additions and 48 deletions

View File

@ -7,7 +7,8 @@ use solana_sdk::{
transaction::TransactionError,
};
const MAX_UNITS: u32 = 1_400_000;
pub const DEFAULT_UNITS: u32 = 200_000;
pub const MAX_UNITS: u32 = 1_400_000;
const MAX_HEAP_FRAME_BYTES: u32 = 256 * 1024;
#[cfg(RUSTC_WITH_SPECIALIZATION)]
@ -67,14 +68,14 @@ pub struct ComputeBudget {
impl Default for ComputeBudget {
fn default() -> Self {
Self::new(true)
Self::new(MAX_UNITS)
}
}
impl ComputeBudget {
pub fn new(use_max_units_default: bool) -> Self {
pub fn new(max_units: u32) -> Self {
ComputeBudget {
max_units: ComputeBudget::get_max_units(use_max_units_default),
max_units: max_units as u64,
log_64_units: 100,
create_program_address_units: 1500,
invoke_units: 1000,
@ -97,14 +98,6 @@ impl ComputeBudget {
}
}
pub fn get_max_units(use_max_units_default: bool) -> u64 {
if use_max_units_default {
MAX_UNITS as u64
} else {
200_000
}
}
pub fn process_message(
&mut self,
message: &SanitizedMessage,

View File

@ -1193,6 +1193,7 @@ pub fn visit_each_account_once<E>(
mod tests {
use {
super::*,
crate::compute_budget,
serde::{Deserialize, Serialize},
solana_sdk::account::{ReadableAccount, WritableAccount},
};
@ -1674,12 +1675,12 @@ mod tests {
let mut transaction_context = TransactionContext::new(accounts, 1, 3);
let mut invoke_context = InvokeContext::new_mock(&mut transaction_context, &[]);
invoke_context.feature_set = Arc::new(feature_set);
invoke_context.compute_budget = ComputeBudget::new(false);
invoke_context.compute_budget = ComputeBudget::new(compute_budget::DEFAULT_UNITS);
invoke_context.push(&[], &[0], &[]).unwrap();
assert_eq!(
*invoke_context.get_compute_budget(),
ComputeBudget::new(false)
ComputeBudget::new(compute_budget::DEFAULT_UNITS)
);
invoke_context.pop().unwrap();
@ -1687,7 +1688,7 @@ mod tests {
let expected_compute_budget = ComputeBudget {
max_units: 500_000,
heap_size: Some(256_usize.saturating_mul(1024)),
..ComputeBudget::new(false)
..ComputeBudget::new(compute_budget::DEFAULT_UNITS)
};
assert_eq!(
*invoke_context.get_compute_budget(),
@ -1698,7 +1699,7 @@ mod tests {
invoke_context.push(&[], &[0], &[]).unwrap();
assert_eq!(
*invoke_context.get_compute_budget(),
ComputeBudget::new(false)
ComputeBudget::new(compute_budget::DEFAULT_UNITS)
);
invoke_context.pop().unwrap();
}