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.
65///
66
67pub struct MeasuredAllocator<A = System> {
68    allocator: A,
69}
70impl<A> MeasuredAllocator<A> {
71    /// Wrap the allocator selected by this test binary.
72    #[must_use]
73    pub const fn new(allocator: A) -> Self {
74        Self { allocator }
75    }
76}
77
78#[expect(
79    unsafe_code,
80    reason = "forward the GlobalAlloc contract to the caller-selected allocator"
81)]
82unsafe impl<A: GlobalAlloc> GlobalAlloc for MeasuredAllocator<A> {
83    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
84        // SAFETY: The caller supplies GlobalAlloc's valid layout unchanged.
85        let pointer = unsafe { self.allocator.alloc(layout) };
86        if !pointer.is_null() {
87            account(0, layout.size(), true, false);
88        }
89        pointer
90    }
91    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
92        // SAFETY: Forward the valid layout and preserve zero-initialization.
93        let pointer = unsafe { self.allocator.alloc_zeroed(layout) };
94        if !pointer.is_null() {
95            account(0, layout.size(), true, false);
96        }
97        pointer
98    }
99    unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
100        account(layout.size(), 0, false, false);
101        // SAFETY: Forward the caller's allocation and matching layout unchanged.
102        unsafe { self.allocator.dealloc(pointer, layout) };
103    }
104    unsafe fn realloc(&self, pointer: *mut u8, layout: Layout, size: usize) -> *mut u8 {
105        // SAFETY: Forward the valid pointer/layout and requested nonzero size.
106        let result = unsafe { self.allocator.realloc(pointer, layout, size) };
107        if !result.is_null() {
108            account(layout.size(), size, false, true);
109        }
110        result
111    }
112}
113
114///
115/// AlreadyMeasuring
116///
117/// A nested measurement was refused without executing its workload.
118///
119
120#[derive(Clone, Copy, Debug, Eq, PartialEq)]
121pub struct AlreadyMeasuring;
122impl std::fmt::Display for AlreadyMeasuring {
123    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
124        f.write_str("allocation measurement already active on this thread")
125    }
126}
127impl std::error::Error for AlreadyMeasuring {}
128
129struct Guard;
130impl Drop for Guard {
131    fn drop(&mut self) {
132        STATE.with(|cell| cell.set(State::default()));
133    }
134}
135
136/// Measure a workload on this thread; nested scopes refuse before executing.
137///
138/// Setup, result formatting and other threads are excluded. Panics propagate,
139/// and an unwind guard disables accounting before returning control to the caller.
140///
141/// # Errors
142/// Returns `AlreadyMeasuring` when called from an active measurement.
143pub fn measure(
144    iterations: usize,
145    mut operation: impl FnMut(),
146) -> Result<Measurement, AlreadyMeasuring> {
147    STATE.with(|cell| {
148        if cell.get().active {
149            return Err(AlreadyMeasuring);
150        }
151        cell.set(State {
152            active: true,
153            ..State::default()
154        });
155        Ok(())
156    })?;
157    let guard = Guard;
158    let started = Instant::now();
159    for _ in 0..iterations {
160        operation();
161    }
162    let elapsed = started.elapsed();
163    let mut counts = STATE.with(|cell| cell.get().counts);
164    drop(guard);
165    counts.elapsed = elapsed;
166    Ok(counts)
167}
ic_testkit_native/lib.rs

ic_testkit_native/
lib.rs
