ic_testkit_native/
lib.rs

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

ic_testkit_native/
lib.rs
