Skip to main content

scx_mlfq/
alloc.rs

1// SPDX-License-Identifier: GPL-2.0
2// Copyright (c) 2026 Galih Tama <galpt@v.recipes>
3
4#![allow(dead_code, unused_imports)]
5
6use std::sync::atomic::{AtomicUsize, Ordering};
7
8#[cfg(feature = "count_alloc")]
9use std::alloc::{GlobalAlloc, Layout, System};
10
11/// Counters for the optional allocation tracker. Only active when the
12/// `count_alloc` feature is enabled. The global allocator below bumps
13/// these on every heap allocation, so a hot-path iteration that stays at
14/// zero proves the zero-allocation guarantee.
15pub static ALLOC_COUNT: AtomicUsize = AtomicUsize::new(0);
16pub static ALLOC_BYTES: AtomicUsize = AtomicUsize::new(0);
17
18/// Reset the counters to zero. Call before a measurement window.
19pub fn reset_counters() {
20    ALLOC_COUNT.store(0, Ordering::Relaxed);
21    ALLOC_BYTES.store(0, Ordering::Relaxed);
22}
23
24/// Read the current counters.
25pub fn counters() -> (usize, usize) {
26    (
27        ALLOC_COUNT.load(Ordering::Relaxed),
28        ALLOC_BYTES.load(Ordering::Relaxed),
29    )
30}
31
32#[cfg(feature = "count_alloc")]
33pub struct TrackingAllocator;
34
35#[cfg(feature = "count_alloc")]
36unsafe impl GlobalAlloc for TrackingAllocator {
37    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
38        unsafe {
39            ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
40            ALLOC_BYTES.fetch_add(layout.size(), Ordering::Relaxed);
41            System.alloc(layout)
42        }
43    }
44    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
45        unsafe { System.dealloc(ptr, layout) }
46    }
47    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
48        unsafe {
49            ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
50            ALLOC_BYTES.fetch_add(layout.size(), Ordering::Relaxed);
51            System.alloc_zeroed(layout)
52        }
53    }
54    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
55        unsafe {
56            ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
57            ALLOC_BYTES.fetch_add(new_size, Ordering::Relaxed);
58            System.realloc(ptr, layout, new_size)
59        }
60    }
61}