ax_alloc/
tracking.rs

1use alloc::collections::btree_map::BTreeMap;
2use core::{
3    alloc::Layout,
4    ops::Range,
5    sync::atomic::{AtomicBool, Ordering},
6};
7
8use ax_sync::SpinLock;
9use axbacktrace::Backtrace;
10
11pub(crate) static TRACKING_ENABLED: AtomicBool = AtomicBool::new(false);
12
13#[ax_percpu::def_percpu]
14pub(crate) static IN_GLOBAL_ALLOCATOR: bool = false;
15// Re-entrancy note: `Backtrace::capture()` now uses `InlineFrames` (stack-allocated
16// array, no heap) so it does NOT re-enter the allocator. The `IN_GLOBAL_ALLOCATOR`
17// guard remains as a safety net for any future code paths that might allocate
18// during backtrace capture.
19
20/// Metadata for each allocation made by the global allocator.
21#[derive(Debug)]
22pub struct AllocationInfo {
23    /// Layout of the allocation.
24    pub layout: Layout,
25    /// Backtrace at the time of allocation.
26    pub backtrace: Backtrace,
27    /// Generation at which the allocation was made.
28    pub generation: u64,
29}
30
31pub(crate) struct GlobalState {
32    // FIXME: don't know why using HashMap causes crash
33    pub map: BTreeMap<usize, AllocationInfo>,
34    pub generation: u64,
35}
36
37static STATE: SpinLock<GlobalState> = SpinLock::new(GlobalState {
38    map: BTreeMap::new(),
39    generation: 0,
40});
41
42/// Enables allocation tracking.
43pub fn enable_tracking() {
44    TRACKING_ENABLED.store(true, Ordering::SeqCst);
45}
46
47/// Disables allocation tracking.
48pub fn disable_tracking() {
49    TRACKING_ENABLED.store(false, Ordering::SeqCst);
50}
51
52/// Returns whether allocation tracking is enabled.
53pub fn tracking_enabled() -> bool {
54    TRACKING_ENABLED.load(Ordering::SeqCst)
55}
56
57pub(crate) fn with_state<R>(f: impl FnOnce(Option<&mut GlobalState>) -> R) -> R {
58    let _guard = ax_sync::PreemptGuard::new();
59    // SAFETY: the guard prevents migration throughout all accesses below.
60    unsafe {
61        ax_percpu::with_cpu_pin(|pin| {
62            if IN_GLOBAL_ALLOCATOR.read_current(pin) || !tracking_enabled() {
63                return f(None);
64            }
65
66            IN_GLOBAL_ALLOCATOR.write_current(pin, true);
67            let mut state = STATE.lock_irqsave();
68            let result = f(Some(&mut state));
69            drop(state);
70            IN_GLOBAL_ALLOCATOR.write_current(pin, false);
71            result
72        })
73    }
74    .expect("allocator tracking requires an installed CPU area")
75}
76
77/// Returns the current generation of the global allocator.
78///
79/// The generation is incremented every time a new allocation is made. It
80/// can be utilized to track the changes in the allocation state over time.
81///
82/// See [`allocations_in`].
83pub fn current_generation() -> u64 {
84    STATE.lock_irqsave().generation
85}
86
87/// Visits all allocations made by the global allocator within the given
88/// generation range.
89pub fn allocations_in(range: Range<u64>, visitor: impl FnMut(&AllocationInfo)) {
90    with_state(|state| {
91        state
92            .unwrap()
93            .map
94            .values()
95            .filter(move |info| range.contains(&info.generation))
96            .for_each(visitor)
97    });
98}
ax_alloc/tracking.rs

ax_alloc/
tracking.rs
