ic_testkit_native/
lib.rs

1//! Thread-local native allocation measurements without runtime dependencies.
2//! Optional `stable-memory` provides shared backing-memory IO observations.
3#![cfg(not(target_arch = "wasm32"))]
4
5#[cfg(feature = "stable-memory")]
6pub mod memory;
7
8use std::{
9    alloc::{GlobalAlloc, Layout, System},
10    cell::Cell,
11    time::{Duration, Instant},
12};
13
14///
15/// Measurement
16///
17/// Counters for one thread's scoped workload, excluding setup and reporting.
18///
19
20#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
21pub struct Measurement {
22    /// Successful allocation calls, excluding reallocations.
23    pub allocations: usize,
24    /// Successful reallocation calls.
25    pub reallocations: usize,
26    /// Allocated minus freed bytes; may be negative for preexisting allocations.
27    pub net_bytes: i128,
28    /// Peak positive net bytes relative to scope entry, not process heap or RSS.
29    pub peak_bytes: usize,
30    /// Native elapsed time for the workload.
31    pub elapsed: Duration,
32}
33
34#[derive(Clone, Copy, Default)]
35struct State {
36    active: bool,
37    counts: Measurement,
38}
39thread_local! { static STATE: Cell<State> = const { Cell::new(State {
40    active: false,
41    counts: Measurement { allocations: 0, reallocations: 0, net_bytes: 0, peak_bytes: 0, elapsed: Duration::ZERO },
42}) }; }
43
44fn account(old: usize, new: usize, allocation: bool, reallocation: bool) {
45    let _ = STATE.try_with(|cell| {
46        let mut state = cell.get();
47        if state.active {
48            state.counts.net_bytes += new as i128 - old as i128;
49            state.counts.peak_bytes = state
50                .counts
51                .peak_bytes
52                .max(usize::try_from(state.counts.net_bytes).unwrap_or(0));
53            state.counts.allocations += usize::from(allocation);
54            state.counts.reallocations += usize::from(reallocation);
55            cell.set(state);
56        }
57    });
58}
59
60///
61/// `MeasuredAllocator`
62///
63/// Instrument an existing global allocator while preserving its allocation contract.
64/// Counters are opt-in and scoped to the calling thread.
65pub struct MeasuredAllocator<A = System> {
66    allocator: A,
67}
68impl<A> MeasuredAllocator<A> {
69    /// Wrap the allocator selected by this test binary.
70    #[must_use]
71    pub const fn new(allocator: A) -> Self {
72        Self { allocator }
73    }
74}
75
76#[expect(
77    unsafe_code,
78    reason = "forward the GlobalAlloc contract to the caller-selected allocator"
79)]
80unsafe impl<A: GlobalAlloc> GlobalAlloc for MeasuredAllocator<A> {
81    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
82        // SAFETY: The caller supplies GlobalAlloc's valid layout unchanged.
83        let pointer = unsafe { self.allocator.alloc(layout) };
84        if !pointer.is_null() {
85            account(0, layout.size(), true, false);
86        }
87        pointer
88    }
89    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
90        // SAFETY: Forward the valid layout and preserve zero-initialization.
91        let pointer = unsafe { self.allocator.alloc_zeroed(layout) };
92        if !pointer.is_null() {
93            account(0, layout.size(), true, false);
94        }
95        pointer
96    }
97    unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
98        account(layout.size(), 0, false, false);
99        // SAFETY: Forward the caller's allocation and matching layout unchanged.
100        unsafe { self.allocator.dealloc(pointer, layout) };
101    }
102    unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, size: usize) -> *mut u8 {
103        // SAFETY: Forward the valid pointer/layout and requested nonzero size.
104        let result = unsafe { self.allocator.realloc(pointer, layout, size) };
105        if !result.is_null() {
106            account(layout.size(), size, false, true);
107        }
108        result
109    }
110}
111
112///
113/// `AlreadyMeasuring`
114///
115/// A nested measurement was refused without executing its workload.
116///
117
118#[derive(Clone, Copy, Debug, Eq, PartialEq)]
119pub struct AlreadyMeasuring;
120impl std::fmt::Display for AlreadyMeasuring {
121    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
122        f.write_str("allocation measurement already active on this thread")
123    }
124}
125impl std::error::Error for AlreadyMeasuring {}
126
127struct Guard;
128impl Drop for Guard {
129    fn drop(&mut self) {
130        STATE.with(|cell| cell.set(State::default()));
131    }
132}
133
134/// Measure a workload on this thread; nested scopes refuse before executing.
135///
136/// Setup, result formatting and other threads are excluded. Panics propagate,
137/// and an unwind guard disables accounting before returning control to the caller.
138///
139/// # Errors
140/// Returns `AlreadyMeasuring` when called from an active measurement.
141pub fn measure(
142    iterations: usize,
143    mut operation: impl FnMut(),
144) -> Result<Measurement, AlreadyMeasuring> {
145    STATE.with(|cell| {
146        if cell.get().active {
147            return Err(AlreadyMeasuring);
148        }
149        cell.set(State {
150            active: true,
151            ..State::default()
152        });
153        Ok(())
154    })?;
155    let guard = Guard;
156    let started = Instant::now();
157    for _ in 0..iterations {
158        operation();
159    }
160    let elapsed = started.elapsed();
161    let mut counts = STATE.with(|cell| cell.get().counts);
162    drop(guard);
163    counts.elapsed = elapsed;
164    Ok(counts)
165}
ic_testkit_native/lib.rs

ic_testkit_native/
lib.rs
