Skip to main content

oxihuman_core/
memory_tracker.rs

1//! Memory usage tracking and budget management.
2
3/// Categories of memory allocations.
4#[allow(dead_code)]
5#[derive(Clone, Debug, PartialEq, Eq, Hash)]
6pub enum MemoryCategory {
7    /// Mesh geometry buffers.
8    Meshes,
9    /// Texture / image data.
10    Textures,
11    /// Physics simulation data.
12    Physics,
13    /// Audio buffers.
14    Audio,
15    /// Uncategorised allocations.
16    Other,
17}
18
19/// A record of a single allocation or free event.
20#[allow(dead_code)]
21#[derive(Clone, Debug)]
22pub struct AllocationRecord {
23    /// Human-readable label.
24    pub label: String,
25    /// Size in bytes.
26    pub size: u64,
27    /// Category.
28    pub category: MemoryCategory,
29    /// `true` for allocation, `false` for free.
30    pub is_alloc: bool,
31}
32
33/// Tracks memory allocations, frees, budgets and peak usage.
34#[allow(dead_code)]
35#[derive(Clone, Debug)]
36pub struct MemoryTracker {
37    /// Running total of current usage in bytes.
38    current: u64,
39    /// Peak usage ever observed.
40    peak: u64,
41    /// Per-category current usage.
42    by_category: Vec<(MemoryCategory, u64)>,
43    /// Budget cap in bytes (0 = unlimited).
44    budget: u64,
45    /// Total allocation operations.
46    alloc_count: u64,
47    /// Total free operations.
48    free_count: u64,
49}
50
51// ---------------------------------------------------------------------------
52// Helpers
53// ---------------------------------------------------------------------------
54
55/// Find or insert a category entry in the tracker's category vec.
56#[allow(dead_code)]
57fn cat_index(tracker: &mut MemoryTracker, cat: &MemoryCategory) -> usize {
58    if let Some(idx) = tracker.by_category.iter().position(|(c, _)| c == cat) {
59        idx
60    } else {
61        tracker.by_category.push((cat.clone(), 0));
62        tracker.by_category.len() - 1
63    }
64}
65
66/// Find category entry (immutable).
67#[allow(dead_code)]
68fn cat_usage(tracker: &MemoryTracker, cat: &MemoryCategory) -> u64 {
69    tracker
70        .by_category
71        .iter()
72        .find(|(c, _)| c == cat)
73        .map_or(0, |(_, v)| *v)
74}
75
76// ---------------------------------------------------------------------------
77// Construction
78// ---------------------------------------------------------------------------
79
80/// Create a new `MemoryTracker` with no allocations and no budget.
81#[allow(dead_code)]
82pub fn new_memory_tracker() -> MemoryTracker {
83    MemoryTracker {
84        current: 0,
85        peak: 0,
86        by_category: Vec::new(),
87        budget: 0,
88        alloc_count: 0,
89        free_count: 0,
90    }
91}
92
93// ---------------------------------------------------------------------------
94// Track operations
95// ---------------------------------------------------------------------------
96
97/// Record an allocation of `size` bytes under `category`.
98#[allow(dead_code)]
99pub fn track_alloc(tracker: &mut MemoryTracker, size: u64, category: MemoryCategory) {
100    tracker.current += size;
101    tracker.alloc_count += 1;
102    if tracker.current > tracker.peak {
103        tracker.peak = tracker.current;
104    }
105    let idx = cat_index(tracker, &category);
106    tracker.by_category[idx].1 += size;
107}
108
109/// Record a free of `size` bytes under `category`.
110/// Saturates at zero if the free exceeds current usage.
111#[allow(dead_code)]
112pub fn track_free(tracker: &mut MemoryTracker, size: u64, category: MemoryCategory) {
113    tracker.current = tracker.current.saturating_sub(size);
114    tracker.free_count += 1;
115    let idx = cat_index(tracker, &category);
116    tracker.by_category[idx].1 = tracker.by_category[idx].1.saturating_sub(size);
117}
118
119// ---------------------------------------------------------------------------
120// Queries
121// ---------------------------------------------------------------------------
122
123/// Return the current total memory usage in bytes.
124#[allow(dead_code)]
125pub fn current_usage(tracker: &MemoryTracker) -> u64 {
126    tracker.current
127}
128
129/// Return the current usage for a specific category.
130#[allow(dead_code)]
131pub fn usage_by_category(tracker: &MemoryTracker, category: &MemoryCategory) -> u64 {
132    cat_usage(tracker, category)
133}
134
135/// Return the peak memory usage ever recorded.
136#[allow(dead_code)]
137pub fn peak_usage(tracker: &MemoryTracker) -> u64 {
138    tracker.peak
139}
140
141/// Return the remaining bytes under the budget. Returns 0 if no budget is set
142/// or if usage exceeds the budget.
143#[allow(dead_code)]
144pub fn budget_remaining(tracker: &MemoryTracker) -> u64 {
145    if tracker.budget == 0 {
146        return 0;
147    }
148    tracker.budget.saturating_sub(tracker.current)
149}
150
151/// Set the memory budget in bytes.  0 means unlimited.
152#[allow(dead_code)]
153pub fn set_budget(tracker: &mut MemoryTracker, budget: u64) {
154    tracker.budget = budget;
155}
156
157/// Return `true` if current usage exceeds the budget (and budget > 0).
158#[allow(dead_code)]
159pub fn over_budget(tracker: &MemoryTracker) -> bool {
160    tracker.budget > 0 && tracker.current > tracker.budget
161}
162
163/// Return the total number of allocation operations.
164#[allow(dead_code)]
165pub fn allocation_count(tracker: &MemoryTracker) -> u64 {
166    tracker.alloc_count
167}
168
169/// Return the total number of free operations.
170#[allow(dead_code)]
171pub fn free_count(tracker: &MemoryTracker) -> u64 {
172    tracker.free_count
173}
174
175/// Return the category with the highest current usage, or `None` if empty.
176#[allow(dead_code)]
177pub fn largest_category(tracker: &MemoryTracker) -> Option<MemoryCategory> {
178    tracker
179        .by_category
180        .iter()
181        .max_by_key(|(_, v)| *v)
182        .map(|(c, _)| c.clone())
183}
184
185// ---------------------------------------------------------------------------
186// Reset
187// ---------------------------------------------------------------------------
188
189/// Reset all counters and category data to zero.
190#[allow(dead_code)]
191pub fn reset_tracker(tracker: &mut MemoryTracker) {
192    tracker.current = 0;
193    tracker.peak = 0;
194    tracker.by_category.clear();
195    tracker.budget = 0;
196    tracker.alloc_count = 0;
197    tracker.free_count = 0;
198}
199
200// ---------------------------------------------------------------------------
201// Serialization
202// ---------------------------------------------------------------------------
203
204/// Produce a minimal JSON representation of the tracker state.
205#[allow(dead_code)]
206pub fn memory_tracker_to_json(tracker: &MemoryTracker) -> String {
207    let mut s = String::from("{");
208    s.push_str(&format!("\"current\":{}", tracker.current));
209    s.push_str(&format!(",\"peak\":{}", tracker.peak));
210    s.push_str(&format!(",\"budget\":{}", tracker.budget));
211    s.push_str(&format!(",\"alloc_count\":{}", tracker.alloc_count));
212    s.push_str(&format!(",\"free_count\":{}", tracker.free_count));
213    s.push_str(",\"categories\":{");
214    for (i, (cat, val)) in tracker.by_category.iter().enumerate() {
215        if i > 0 {
216            s.push(',');
217        }
218        s.push_str(&format!("\"{:?}\":{}", cat, val));
219    }
220    s.push_str("}}");
221    s
222}
223
224// ---------------------------------------------------------------------------
225// Tests
226// ---------------------------------------------------------------------------
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231
232    #[test]
233    fn test_new_memory_tracker() {
234        let t = new_memory_tracker();
235        assert_eq!(current_usage(&t), 0);
236        assert_eq!(peak_usage(&t), 0);
237    }
238
239    #[test]
240    fn test_track_alloc() {
241        let mut t = new_memory_tracker();
242        track_alloc(&mut t, 1024, MemoryCategory::Meshes);
243        assert_eq!(current_usage(&t), 1024);
244        assert_eq!(allocation_count(&t), 1);
245    }
246
247    #[test]
248    fn test_track_free() {
249        let mut t = new_memory_tracker();
250        track_alloc(&mut t, 1000, MemoryCategory::Textures);
251        track_free(&mut t, 400, MemoryCategory::Textures);
252        assert_eq!(current_usage(&t), 600);
253        assert_eq!(free_count(&t), 1);
254    }
255
256    #[test]
257    fn test_track_free_saturates() {
258        let mut t = new_memory_tracker();
259        track_alloc(&mut t, 100, MemoryCategory::Audio);
260        track_free(&mut t, 999, MemoryCategory::Audio);
261        assert_eq!(current_usage(&t), 0);
262    }
263
264    #[test]
265    fn test_usage_by_category() {
266        let mut t = new_memory_tracker();
267        track_alloc(&mut t, 500, MemoryCategory::Meshes);
268        track_alloc(&mut t, 300, MemoryCategory::Textures);
269        assert_eq!(usage_by_category(&t, &MemoryCategory::Meshes), 500);
270        assert_eq!(usage_by_category(&t, &MemoryCategory::Textures), 300);
271        assert_eq!(usage_by_category(&t, &MemoryCategory::Physics), 0);
272    }
273
274    #[test]
275    fn test_peak_usage() {
276        let mut t = new_memory_tracker();
277        track_alloc(&mut t, 1000, MemoryCategory::Meshes);
278        track_alloc(&mut t, 500, MemoryCategory::Textures);
279        track_free(&mut t, 1000, MemoryCategory::Meshes);
280        assert_eq!(peak_usage(&t), 1500);
281        assert_eq!(current_usage(&t), 500);
282    }
283
284    #[test]
285    fn test_set_budget_and_remaining() {
286        let mut t = new_memory_tracker();
287        set_budget(&mut t, 2000);
288        track_alloc(&mut t, 800, MemoryCategory::Audio);
289        assert_eq!(budget_remaining(&t), 1200);
290    }
291
292    #[test]
293    fn test_budget_remaining_no_budget() {
294        let t = new_memory_tracker();
295        assert_eq!(budget_remaining(&t), 0);
296    }
297
298    #[test]
299    fn test_over_budget() {
300        let mut t = new_memory_tracker();
301        set_budget(&mut t, 100);
302        track_alloc(&mut t, 200, MemoryCategory::Physics);
303        assert!(over_budget(&t));
304    }
305
306    #[test]
307    fn test_not_over_budget() {
308        let mut t = new_memory_tracker();
309        set_budget(&mut t, 1000);
310        track_alloc(&mut t, 500, MemoryCategory::Meshes);
311        assert!(!over_budget(&t));
312    }
313
314    #[test]
315    fn test_allocation_and_free_counts() {
316        let mut t = new_memory_tracker();
317        track_alloc(&mut t, 100, MemoryCategory::Meshes);
318        track_alloc(&mut t, 200, MemoryCategory::Textures);
319        track_free(&mut t, 50, MemoryCategory::Meshes);
320        assert_eq!(allocation_count(&t), 2);
321        assert_eq!(free_count(&t), 1);
322    }
323
324    #[test]
325    fn test_largest_category() {
326        let mut t = new_memory_tracker();
327        track_alloc(&mut t, 100, MemoryCategory::Meshes);
328        track_alloc(&mut t, 500, MemoryCategory::Textures);
329        track_alloc(&mut t, 200, MemoryCategory::Audio);
330        assert_eq!(largest_category(&t), Some(MemoryCategory::Textures));
331    }
332
333    #[test]
334    fn test_largest_category_empty() {
335        let t = new_memory_tracker();
336        assert!(largest_category(&t).is_none());
337    }
338
339    #[test]
340    fn test_reset_tracker() {
341        let mut t = new_memory_tracker();
342        track_alloc(&mut t, 1000, MemoryCategory::Meshes);
343        set_budget(&mut t, 5000);
344        reset_tracker(&mut t);
345        assert_eq!(current_usage(&t), 0);
346        assert_eq!(peak_usage(&t), 0);
347        assert_eq!(allocation_count(&t), 0);
348        assert_eq!(free_count(&t), 0);
349    }
350
351    #[test]
352    fn test_memory_tracker_to_json() {
353        let mut t = new_memory_tracker();
354        track_alloc(&mut t, 256, MemoryCategory::Meshes);
355        let json = memory_tracker_to_json(&t);
356        assert!(json.contains("\"current\":256"));
357        assert!(json.contains("\"peak\":256"));
358        assert!(json.contains("\"alloc_count\":1"));
359    }
360
361    #[test]
362    fn test_multiple_categories_independent() {
363        let mut t = new_memory_tracker();
364        track_alloc(&mut t, 100, MemoryCategory::Meshes);
365        track_alloc(&mut t, 200, MemoryCategory::Audio);
366        track_free(&mut t, 100, MemoryCategory::Meshes);
367        assert_eq!(usage_by_category(&t, &MemoryCategory::Meshes), 0);
368        assert_eq!(usage_by_category(&t, &MemoryCategory::Audio), 200);
369        assert_eq!(current_usage(&t), 200);
370    }
371}