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

ic_testkit_native/
lib.rs
