Skip to main content

khora_core/memory/
mod.rs

1// Copyright 2025 eraflo
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Provides a public interface for querying engine-wide memory allocation statistics.
16//!
17//! This module defines a set of global atomic counters for detailed memory tracking.
18//! It forms a "contract" where a registered global allocator is responsible for
19//! incrementing these counters, and any part of the engine can read them in a
20//! thread-safe manner to monitor memory usage.
21//!
22//! The primary use case is for the `khora-telemetry` crate to collect these stats
23//! and feed them into the Dynamic Context Core (DCC) for adaptive decision-making.
24
25use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
26
27mod tracking_allocator;
28pub use tracking_allocator::SaaTrackingAllocator;
29
30// --- Global Memory Counters ---
31
32/// Tracks the total number of bytes currently allocated by the registered global allocator.
33pub static CURRENTLY_ALLOCATED_BYTES: AtomicUsize = AtomicUsize::new(0);
34
35/// Tracks the peak number of bytes ever allocated simultaneously during the application's lifetime.
36pub static PEAK_ALLOCATED_BYTES: AtomicU64 = AtomicU64::new(0);
37
38/// Tracks the total number of allocation calls made.
39pub static TOTAL_ALLOCATIONS: AtomicU64 = AtomicU64::new(0);
40
41/// Tracks the total number of deallocation calls made.
42pub static TOTAL_DEALLOCATIONS: AtomicU64 = AtomicU64::new(0);
43
44/// Tracks the total number of reallocation calls made.
45pub static TOTAL_REALLOCATIONS: AtomicU64 = AtomicU64::new(0);
46
47/// Tracks the cumulative total of bytes ever allocated over the application's lifetime.
48pub static BYTES_ALLOCATED_LIFETIME: AtomicU64 = AtomicU64::new(0);
49
50/// Tracks the cumulative total of bytes ever deallocated over the application's lifetime.
51pub static BYTES_DEALLOCATED_LIFETIME: AtomicU64 = AtomicU64::new(0);
52
53/// Tracks the number of "large" allocations (e.g., >= 1MB).
54pub static LARGE_ALLOCATIONS: AtomicU64 = AtomicU64::new(0);
55
56/// Tracks the cumulative total of bytes from "large" allocations.
57pub static LARGE_ALLOCATION_BYTES: AtomicU64 = AtomicU64::new(0);
58
59/// Tracks the number of "small" allocations (e.g., < 1KB).
60pub static SMALL_ALLOCATIONS: AtomicU64 = AtomicU64::new(0);
61
62/// Tracks the cumulative total of bytes from "small" allocations.
63pub static SMALL_ALLOCATION_BYTES: AtomicU64 = AtomicU64::new(0);
64
65// --- Data Structures for Reporting ---
66
67/// A snapshot of comprehensive memory allocation statistics, including derived metrics.
68#[derive(Debug, Clone, Copy, Default)]
69pub struct ExtendedMemoryStats {
70    // --- Current State ---
71    /// The total number of bytes currently in use.
72    pub current_allocated_bytes: usize,
73    /// The maximum number of bytes that were ever in use simultaneously.
74    pub peak_allocated_bytes: u64,
75
76    // --- Allocation Counters ---
77    /// The total number of times an allocation was requested.
78    pub total_allocations: u64,
79    /// The total number of times a deallocation was requested.
80    pub total_deallocations: u64,
81    /// The total number of times a reallocation was requested.
82    pub total_reallocations: u64,
83    /// The net number of active allocations (`total_allocations` - `total_deallocations`).
84    pub net_allocations: i64,
85
86    // --- Lifetime Totals ---
87    /// The cumulative sum of all bytes ever allocated.
88    pub bytes_allocated_lifetime: u64,
89    /// The cumulative sum of all bytes ever deallocated.
90    pub bytes_deallocated_lifetime: u64,
91    /// The net number of bytes allocated over the lifetime. Should be equal to `current_allocated_bytes`.
92    pub bytes_net_lifetime: i64,
93
94    // --- Size Category Tracking ---
95    /// The number of allocations classified as "large".
96    pub large_allocations: u64,
97    /// The total byte size of all "large" allocations.
98    pub large_allocation_bytes: u64,
99    /// The number of allocations classified as "small".
100    pub small_allocations: u64,
101    /// The total byte size of all "small" allocations.
102    pub small_allocation_bytes: u64,
103    /// The number of allocations not classified as small or large.
104    pub medium_allocations: u64,
105    /// The total byte size of all "medium" allocations.
106    pub medium_allocation_bytes: u64,
107
108    // --- Calculated Metrics ---
109    /// The average size of a single allocation (`bytes_allocated_lifetime` / `total_allocations`).
110    pub average_allocation_size: f64,
111    /// A rough measure of memory fragmentation (`1.0 - current / peak`).
112    pub fragmentation_ratio: f64,
113    /// The ratio of memory still in use compared to all memory ever allocated.
114    pub allocation_efficiency: f64,
115}
116
117impl ExtendedMemoryStats {
118    /// Populates the derived metrics based on the raw counter values.
119    pub fn calculate_derived_metrics(&mut self) {
120        if self.total_allocations > 0 {
121            self.average_allocation_size =
122                self.bytes_allocated_lifetime as f64 / self.total_allocations as f64;
123        }
124
125        if self.peak_allocated_bytes > 0 {
126            self.fragmentation_ratio =
127                1.0 - (self.current_allocated_bytes as f64 / self.peak_allocated_bytes as f64);
128        }
129
130        if self.bytes_allocated_lifetime > 0 {
131            self.allocation_efficiency =
132                self.current_allocated_bytes as f64 / self.bytes_allocated_lifetime as f64;
133        }
134
135        self.medium_allocations =
136            self.total_allocations - self.small_allocations - self.large_allocations;
137        self.medium_allocation_bytes = self.bytes_allocated_lifetime
138            - self.small_allocation_bytes
139            - self.large_allocation_bytes;
140    }
141}
142
143// --- Public API for Reading Stats ---
144
145/// Takes a snapshot of all global memory counters and returns them in a structured format.
146///
147/// This function is the primary entry point for querying memory statistics. It reads
148/// all counters atomically (using `Ordering::Relaxed`) and calculates several
149/// derived metrics.
150pub fn get_extended_memory_stats() -> ExtendedMemoryStats {
151    let current_allocated = CURRENTLY_ALLOCATED_BYTES.load(Ordering::Relaxed);
152    let peak_allocated = PEAK_ALLOCATED_BYTES.load(Ordering::Relaxed);
153    let total_allocs = TOTAL_ALLOCATIONS.load(Ordering::Relaxed);
154    let total_deallocs = TOTAL_DEALLOCATIONS.load(Ordering::Relaxed);
155    let total_reallocs = TOTAL_REALLOCATIONS.load(Ordering::Relaxed);
156    let bytes_alloc_lifetime = BYTES_ALLOCATED_LIFETIME.load(Ordering::Relaxed);
157    let bytes_dealloc_lifetime = BYTES_DEALLOCATED_LIFETIME.load(Ordering::Relaxed);
158    let large_allocs = LARGE_ALLOCATIONS.load(Ordering::Relaxed);
159    let large_alloc_bytes = LARGE_ALLOCATION_BYTES.load(Ordering::Relaxed);
160    let small_allocs = SMALL_ALLOCATIONS.load(Ordering::Relaxed);
161    let small_alloc_bytes = SMALL_ALLOCATION_BYTES.load(Ordering::Relaxed);
162
163    let mut stats = ExtendedMemoryStats {
164        current_allocated_bytes: current_allocated,
165        peak_allocated_bytes: peak_allocated,
166        total_allocations: total_allocs,
167        total_deallocations: total_deallocs,
168        total_reallocations: total_reallocs,
169        net_allocations: total_allocs as i64 - total_deallocs as i64,
170        bytes_allocated_lifetime: bytes_alloc_lifetime,
171        bytes_deallocated_lifetime: bytes_dealloc_lifetime,
172        bytes_net_lifetime: bytes_alloc_lifetime as i64 - bytes_dealloc_lifetime as i64,
173        large_allocations: large_allocs,
174        large_allocation_bytes: large_alloc_bytes,
175        small_allocations: small_allocs,
176        small_allocation_bytes: small_alloc_bytes,
177        ..Default::default()
178    };
179
180    stats.calculate_derived_metrics();
181    stats
182}
183
184/// Gets the total number of bytes currently allocated by the global allocator.
185///
186/// This is a lightweight alternative to `get_extended_memory_stats` for when only
187/// the current usage is needed.
188pub fn get_currently_allocated_bytes() -> usize {
189    CURRENTLY_ALLOCATED_BYTES.load(Ordering::Relaxed)
190}