Skip to content

Commit

Permalink
add test for refaction heap size, fix size truncating by div op
Browse files Browse the repository at this point in the history
  • Loading branch information
tao-stones committed Mar 10, 2023
1 parent 13107b4 commit 443dbf3
Showing 1 changed file with 29 additions and 5 deletions.
34 changes: 29 additions & 5 deletions programs/bpf_loader/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -295,6 +295,15 @@ fn check_loader_id(id: &Pubkey) -> bool {
|| bpf_loader_upgradeable::check_id(id)
}

fn calculate_heap_cost(heap_size: usize, heap_cost: u64) -> u64 {
const KILOBYTE: u64 = 1024;
((heap_size as u64)
.saturating_add((32_u64 - 1).saturating_mul(KILOBYTE))
.saturating_div(32_u64.saturating_mul(KILOBYTE)))
.saturating_sub(1)
.saturating_mul(heap_cost)
}

/// Create the SBF virtual machine
pub fn create_vm<'a, 'b>(
program: &'a VerifiedExecutable<RequisiteVerifier, InvokeContext<'b>>,
Expand All @@ -304,11 +313,7 @@ pub fn create_vm<'a, 'b>(
) -> Result<EbpfVm<'a, RequisiteVerifier, InvokeContext<'b>>, EbpfError> {
let compute_budget = invoke_context.get_compute_budget();
let heap_size = compute_budget.heap_size.unwrap_or(HEAP_LENGTH);
let _ = invoke_context.consume_checked(
((heap_size as u64).saturating_div(32_u64.saturating_mul(1024)))
.saturating_sub(1)
.saturating_mul(compute_budget.heap_cost),
);
let _ = invoke_context.consume_checked(calculate_heap_cost(heap_size, compute_budget.heap_cost));
let heap =
AlignedMemory::<HOST_ALIGN>::zero_filled(compute_budget.heap_size.unwrap_or(HEAP_LENGTH));
let check_aligned = bpf_loader_deprecated::id()
Expand Down Expand Up @@ -3832,4 +3837,23 @@ 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

// assert less than 32K heap should cost zero unit
assert_eq!(0, calculate_heap_cost(31_usize * 1024, heap_cost));

// assert exact 32K heap should be cost zero unit
assert_eq!(0, calculate_heap_cost(32_usize * 1024, heap_cost));

// assert slightly more than 32K heap should cost 1 * heap_cost
assert_eq!(heap_cost, calculate_heap_cost(33_usize * 1024, heap_cost));

// assert exact 64K heap should cost 1 * heap_cost
assert_eq!(heap_cost, calculate_heap_cost(64_usize * 1024, heap_cost));
}
}

0 comments on commit 443dbf3

Please sign in to comment.