diff --git a/programs/bpf_loader/src/lib.rs b/programs/bpf_loader/src/lib.rs index 94477c80962aee..91b244c4d0a816 100644 --- a/programs/bpf_loader/src/lib.rs +++ b/programs/bpf_loader/src/lib.rs @@ -45,7 +45,7 @@ use { check_slice_translation_size, delay_visibility_of_program_deployment, disable_deploy_of_alloc_free_syscall, enable_bpf_loader_extend_program_ix, enable_bpf_loader_set_authority_checked_ix, enable_program_redeployment_cooldown, - limit_max_instruction_trace_length, FeatureSet, + limit_max_instruction_trace_length, round_up_heap_size, FeatureSet, }, instruction::{AccountMeta, InstructionError}, loader_instruction::LoaderInstruction, @@ -298,6 +298,20 @@ fn check_loader_id(id: &Pubkey) -> bool { || bpf_loader_upgradeable::check_id(id) } +fn calculate_heap_cost(heap_size: u64, heap_cost: u64, enable_rounding_fix: bool) -> u64 { + const KIBIBYTE: u64 = 1024; + const PAGE_SIZE_KB: u64 = 32; + let mut rounded_heap_size = heap_size; + if enable_rounding_fix { + rounded_heap_size = rounded_heap_size + .saturating_add(PAGE_SIZE_KB.saturating_mul(KIBIBYTE).saturating_sub(1)); + } + rounded_heap_size + .saturating_div(PAGE_SIZE_KB.saturating_mul(KIBIBYTE)) + .saturating_sub(1) + .saturating_mul(heap_cost) +} + /// Create the SBF virtual machine pub fn create_ebpf_vm<'a, 'b>( program: &'a VerifiedExecutable>, @@ -307,11 +321,17 @@ pub fn create_ebpf_vm<'a, 'b>( orig_account_lengths: Vec, invoke_context: &'a mut InvokeContext<'b>, ) -> Result>, EbpfError> { - let _ = invoke_context.consume_checked( - ((heap.len() as u64).saturating_div(32_u64.saturating_mul(1024))) - .saturating_sub(1) - .saturating_mul(invoke_context.get_compute_budget().heap_cost), - ); + let round_up_heap_size = invoke_context + .feature_set + .is_active(&round_up_heap_size::id()); + let heap_cost_result = invoke_context.consume_checked(calculate_heap_cost( + heap.len() as u64, + invoke_context.get_compute_budget().heap_cost, + round_up_heap_size, + )); + if round_up_heap_size { + heap_cost_result.map_err(SyscallError::InstructionError)?; + } let check_aligned = bpf_loader_deprecated::id() != invoke_context .transaction_context @@ -3921,4 +3941,50 @@ mod tests { }, ); } + + #[test] + fn test_calculate_heap_cost() { + let heap_cost = 8_u64; + + // heap allocations are in 32K block, `heap_cost` of CU is consumed per additional 32k + + // when `enable_heap_size_round_up` not enabled: + { + // assert less than 32K heap should cost zero unit + assert_eq!(0, calculate_heap_cost(31_u64 * 1024, heap_cost, false)); + + // assert exact 32K heap should be cost zero unit + assert_eq!(0, calculate_heap_cost(32_u64 * 1024, heap_cost, false)); + + // assert slightly more than 32K heap is mistakenly cost zero unit + assert_eq!(0, calculate_heap_cost(33_u64 * 1024, heap_cost, false)); + + // assert exact 64K heap should cost 1 * heap_cost + assert_eq!( + heap_cost, + calculate_heap_cost(64_u64 * 1024, heap_cost, false) + ); + } + + // when `enable_heap_size_round_up` is enabled: + { + // assert less than 32K heap should cost zero unit + assert_eq!(0, calculate_heap_cost(31_u64 * 1024, heap_cost, true)); + + // assert exact 32K heap should be cost zero unit + assert_eq!(0, calculate_heap_cost(32_u64 * 1024, heap_cost, true)); + + // assert slightly more than 32K heap should cost 1 * heap_cost + assert_eq!( + heap_cost, + calculate_heap_cost(33_u64 * 1024, heap_cost, true) + ); + + // assert exact 64K heap should cost 1 * heap_cost + assert_eq!( + heap_cost, + calculate_heap_cost(64_u64 * 1024, heap_cost, true) + ); + } + } } diff --git a/sdk/src/feature_set.rs b/sdk/src/feature_set.rs index 4814c4494fe070..484fea3c5e575b 100644 --- a/sdk/src/feature_set.rs +++ b/sdk/src/feature_set.rs @@ -626,6 +626,10 @@ pub mod switch_to_new_elf_parser { solana_sdk::declare_id!("Cdkc8PPTeTNUPoZEfCY5AyetUrEdkZtNPMgz58nqyaHD"); } +pub mod round_up_heap_size { + solana_sdk::declare_id!("CE2et8pqgyQMP2mQRg3CgvX8nJBKUArMu3wfiQiQKY1y"); +} + lazy_static! { /// Map of feature identifiers to user-visible description pub static ref FEATURE_NAMES: HashMap = [ @@ -777,6 +781,7 @@ lazy_static! { (apply_cost_tracker_during_replay::id(), "apply cost tracker to blocks during replay #29595"), (add_set_tx_loaded_accounts_data_size_instruction::id(), "add compute budget instruction for setting account data size per transaction #30366"), (switch_to_new_elf_parser::id(), "switch to new ELF parser #30497"), + (round_up_heap_size::id(), "round up heap size when calculating heap cost #30679"), /*************** ADD NEW FEATURES HERE ***************/ ] .iter()