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        ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
39        ALLOC_BYTES.fetch_add(layout.size(), Ordering::Relaxed);
40        System.alloc(layout)
41    }
42    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
43        System.dealloc(ptr, layout)
44    }
45    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
46        ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
47        ALLOC_BYTES.fetch_add(layout.size(), Ordering::Relaxed);
48        System.alloc_zeroed(layout)
49    }
50    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
51        ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
52        ALLOC_BYTES.fetch_add(new_size, Ordering::Relaxed);
53        System.realloc(ptr, layout, new_size)
54    }
55}