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_kspin::SpinNoIrq;
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
16/// Metadata for each allocation made by the global allocator.
17#[derive(Debug)]
18pub struct AllocationInfo {
19    /// Layout of the allocation.
20    pub layout: Layout,
21    /// Backtrace at the time of allocation.
22    pub backtrace: Backtrace,
23    /// Generation at which the allocation was made.
24    pub generation: u64,
25}
26
27pub(crate) struct GlobalState {
28    // FIXME: don't know why using HashMap causes crash
29    pub map: BTreeMap<usize, AllocationInfo>,
30    pub generation: u64,
31}
32
33static STATE: SpinNoIrq<GlobalState> = SpinNoIrq::new(GlobalState {
34    map: BTreeMap::new(),
35    generation: 0,
36});
37
38/// Enables allocation tracking.
39pub fn enable_tracking() {
40    TRACKING_ENABLED.store(true, Ordering::SeqCst);
41}
42
43/// Disables allocation tracking.
44pub fn disable_tracking() {
45    TRACKING_ENABLED.store(false, Ordering::SeqCst);
46}
47
48/// Returns whether allocation tracking is enabled.
49pub fn tracking_enabled() -> bool {
50    TRACKING_ENABLED.load(Ordering::SeqCst)
51}
52
53pub(crate) fn with_state<R>(f: impl FnOnce(Option<&mut GlobalState>) -> R) -> R {
54    IN_GLOBAL_ALLOCATOR.with_current(|in_global| {
55        if *in_global || !tracking_enabled() {
56            f(None)
57        } else {
58            *in_global = true;
59            let mut state = STATE.lock();
60            let result = f(Some(&mut state));
61            *in_global = false;
62            result
63        }
64    })
65}
66
67/// Returns the current generation of the global allocator.
68///
69/// The generation is incremented every time a new allocation is made. It
70/// can be utilized to track the changes in the allocation state over time.
71///
72/// See [`allocations_in`].
73pub fn current_generation() -> u64 {
74    STATE.lock().generation
75}
76
77/// Visits all allocations made by the global allocator within the given
78/// generation range.
79pub fn allocations_in(range: Range<u64>, visitor: impl FnMut(&AllocationInfo)) {
80    with_state(|state| {
81        state
82            .unwrap()
83            .map
84            .values()
85            .filter(move |info| range.contains(&info.generation))
86            .for_each(visitor)
87    });
88}
ax_alloc/tracking.rs

ax_alloc/
tracking.rs
