kglite_c/
alloc.rs

1//! Tracking global allocator + memory stats for the C ABI.
2//!
3//! kglite-c installs a tracking allocator (wrapping the System
4//! allocator) so a binding can observe the Rust-side heap via
5//! [`kglite_memory_stats`] — current live bytes, peak since process
6//! start, and total allocation count. Counters are process-wide.
7//!
8//! Only allocations made through the Rust global allocator are counted;
9//! the host runtime's own heap (Go, the JVM, Node, …) is separate and
10//! invisible here. The counters are maintained with `Relaxed` atomics —
11//! cheap, and exact accounting across threads isn't required for a
12//! monitoring stat.
13
14use std::alloc::{GlobalAlloc, Layout, System};
15use std::sync::atomic::{AtomicU64, Ordering};
16
17static CURRENT: AtomicU64 = AtomicU64::new(0);
18static PEAK: AtomicU64 = AtomicU64::new(0);
19static TOTAL_ALLOCS: AtomicU64 = AtomicU64::new(0);
20
21/// System allocator wrapper that tallies bytes + allocation count.
22struct TrackingAllocator;
23
24// SAFETY: every method forwards to `System` (a sound `GlobalAlloc`) and
25// only adds bookkeeping; we never hand back a pointer System didn't
26// produce. realloc is left to the default `GlobalAlloc` impl, which
27// routes through our `alloc`/`dealloc` so byte accounting stays correct.
28unsafe impl GlobalAlloc for TrackingAllocator {
29    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
30        let ptr = unsafe { System.alloc(layout) };
31        if !ptr.is_null() {
32            let size = layout.size() as u64;
33            TOTAL_ALLOCS.fetch_add(1, Ordering::Relaxed);
34            let now = CURRENT.fetch_add(size, Ordering::Relaxed) + size;
35            PEAK.fetch_max(now, Ordering::Relaxed);
36        }
37        ptr
38    }
39
40    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
41        unsafe { System.dealloc(ptr, layout) };
42        CURRENT.fetch_sub(layout.size() as u64, Ordering::Relaxed);
43    }
44}
45
46#[global_allocator]
47static GLOBAL: TrackingAllocator = TrackingAllocator;
48
49/// Rust-heap statistics from kglite's tracking allocator.
50#[repr(C)]
51#[derive(Debug, Clone, Copy)]
52pub struct KgMemStats {
53    /// Current live Rust-heap bytes (allocated minus freed).
54    pub current_bytes: u64,
55    /// Peak live Rust-heap bytes since process start.
56    pub peak_bytes: u64,
57    /// Total number of allocations since process start (monotonic).
58    pub total_allocs: u64,
59}
60
61/// Return current Rust-heap statistics from kglite's tracking allocator.
62/// Counts only allocations through the Rust global allocator — the host
63/// runtime's own heap is separate. Useful for a binding to surface
64/// kglite's memory footprint in its own metrics.
65#[no_mangle]
66pub extern "C" fn kglite_memory_stats() -> KgMemStats {
67    crate::ffi::value_boundary(
68        KgMemStats {
69            current_bytes: 0,
70            peak_bytes: 0,
71            total_allocs: 0,
72        },
73        || {
74            let current_bytes = CURRENT.load(Ordering::Relaxed);
75            // Another thread can be between CURRENT.fetch_add and PEAK.fetch_max.
76            // Fold this observation into PEAK so every returned snapshot preserves
77            // the public peak >= current invariant without serializing allocations.
78            let peak_bytes = PEAK.fetch_max(current_bytes, Ordering::Relaxed);
79            KgMemStats {
80                current_bytes,
81                peak_bytes: peak_bytes.max(current_bytes),
82                total_allocs: TOTAL_ALLOCS.load(Ordering::Relaxed),
83            }
84        },
85    )
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91
92    #[test]
93    fn memory_stats_track_allocations() {
94        let before = kglite_memory_stats();
95        // Force some heap traffic the optimizer can't elide.
96        let v: Vec<u64> = (0..10_000).collect();
97        let after = kglite_memory_stats();
98        assert!(after.total_allocs >= before.total_allocs);
99        assert!(after.peak_bytes >= after.current_bytes);
100        // Keep `v` alive across the second reading.
101        assert_eq!(v.len(), 10_000);
102    }
103}
kglite_c/alloc.rs

kglite_c/
alloc.rs
