khora_infra/telemetry/
memory_monitor.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//! System Memory Resource Monitor
16//!
17//! Provides memory monitoring capabilities through integration with
18//! the SaaTrackingAllocator for heap allocation tracking.
19
20use std::borrow::Cow;
21use std::sync::Mutex;
22
23use khora_core::memory::{get_currently_allocated_bytes, get_extended_memory_stats};
24use khora_core::telemetry::monitoring::{
25    MemoryReport, MonitoredResourceType, ResourceMonitor, ResourceUsageReport,
26};
27
28/// System memory resource monitor.
29///
30/// Tracks heap allocation statistics through the SaaTrackingAllocator
31/// and provides real-time memory usage information.
32#[derive(Debug)]
33pub struct MemoryMonitor {
34    id: String,
35    last_report: Mutex<Option<MemoryReport>>,
36    peak_usage_bytes: Mutex<usize>,
37    last_allocation_bytes: Mutex<usize>,
38    sample_count: Mutex<u64>,
39}
40
41impl MemoryMonitor {
42    /// Creates a new memory monitor with the given identifier.
43    pub fn new(id: String) -> Self {
44        let current_usage = get_currently_allocated_bytes();
45        Self {
46            id,
47            last_report: Mutex::new(None),
48            peak_usage_bytes: Mutex::new(current_usage),
49            last_allocation_bytes: Mutex::new(current_usage),
50            sample_count: Mutex::new(0),
51        }
52    }
53
54    /// Returns the latest detailed memory report.
55    pub fn get_memory_report(&self) -> Option<MemoryReport> {
56        let last_report = self.last_report.lock().unwrap();
57        *last_report
58    }
59
60    /// Resets the peak usage counter to the current memory usage.
61    pub fn reset_peak_usage(&self) {
62        let current_usage = get_currently_allocated_bytes();
63        let mut peak = self.peak_usage_bytes.lock().unwrap();
64        *peak = current_usage;
65    }
66
67    /// Updates the monitor's internal state by querying the global allocator stats.
68    fn update_internal_stats(&self) {
69        let current_usage = get_currently_allocated_bytes();
70        let extended_stats = get_extended_memory_stats();
71
72        // Update peak tracking
73        let mut peak = self.peak_usage_bytes.lock().unwrap();
74        if current_usage > *peak {
75            *peak = current_usage;
76        }
77
78        // Calculate allocation delta
79        let mut last_alloc = self.last_allocation_bytes.lock().unwrap();
80        let allocation_delta = current_usage.saturating_sub(*last_alloc);
81        *last_alloc = current_usage;
82
83        // Update sample count
84        let mut count = self.sample_count.lock().unwrap();
85        *count += 1;
86
87        // Create comprehensive report with extended statistics
88        let report = MemoryReport {
89            current_usage_bytes: current_usage,
90            peak_usage_bytes: *peak,
91            allocation_delta_bytes: allocation_delta,
92            sample_count: *count,
93
94            // Extended statistics from allocator
95            total_allocations: extended_stats.total_allocations,
96            total_deallocations: extended_stats.total_deallocations,
97            total_reallocations: extended_stats.total_reallocations,
98            bytes_allocated_lifetime: extended_stats.bytes_allocated_lifetime,
99            bytes_deallocated_lifetime: extended_stats.bytes_deallocated_lifetime,
100            large_allocations: extended_stats.large_allocations,
101            large_allocation_bytes: extended_stats.large_allocation_bytes,
102            small_allocations: extended_stats.small_allocations,
103            small_allocation_bytes: extended_stats.small_allocation_bytes,
104            fragmentation_ratio: extended_stats.fragmentation_ratio,
105            allocation_efficiency: extended_stats.allocation_efficiency,
106            average_allocation_size: extended_stats.average_allocation_size,
107        };
108
109        let mut last_report = self.last_report.lock().unwrap();
110        *last_report = Some(report);
111    }
112}
113
114impl ResourceMonitor for MemoryMonitor {
115    fn monitor_id(&self) -> Cow<'static, str> {
116        Cow::Owned(self.id.clone())
117    }
118
119    fn resource_type(&self) -> MonitoredResourceType {
120        MonitoredResourceType::SystemRam
121    }
122
123    fn get_usage_report(&self) -> ResourceUsageReport {
124        let current_usage = get_currently_allocated_bytes();
125        let peak_usage = *self.peak_usage_bytes.lock().unwrap();
126
127        ResourceUsageReport {
128            current_bytes: current_usage as u64,
129            peak_bytes: Some(peak_usage as u64),
130            total_capacity_bytes: None, // System memory limit not easily available
131        }
132    }
133
134    fn as_any(&self) -> &dyn std::any::Any {
135        self
136    }
137
138    fn update(&self) {
139        self.update_internal_stats();
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn memory_monitor_creation() {
149        let monitor = MemoryMonitor::new("TestMemory".to_string());
150        assert_eq!(monitor.monitor_id(), "TestMemory");
151        assert_eq!(monitor.resource_type(), MonitoredResourceType::SystemRam);
152    }
153
154    #[test]
155    fn memory_monitor_update_stats() {
156        let monitor = MemoryMonitor::new("TestMemory".to_string());
157
158        // Initially no report available until update is called
159        assert!(monitor.get_memory_report().is_none());
160
161        // Update stats
162        monitor.update_internal_stats();
163
164        // After update, report should be available
165        let report = monitor.get_memory_report().unwrap();
166        // In test environment, memory usage might be 0, so we just check report exists
167        assert_eq!(report.sample_count, 1);
168    }
169
170    #[test]
171    fn memory_monitor_peak_tracking() {
172        let monitor = MemoryMonitor::new("TestMemory".to_string());
173
174        monitor.update_internal_stats();
175        let _initial_report = monitor.get_memory_report().unwrap();
176        let _initial_peak = _initial_report.peak_usage_bytes;
177
178        // Reset peak and update again
179        monitor.reset_peak_usage();
180        monitor.update_internal_stats();
181
182        let after_reset_report = monitor.get_memory_report().unwrap();
183        // Peak should be reset to current usage
184        assert_eq!(after_reset_report.sample_count, 2);
185    }
186
187    #[test]
188    fn memory_monitor_reset_peak() {
189        let monitor = MemoryMonitor::new("TestMemory".to_string());
190
191        monitor.update_internal_stats();
192        let before_reset = monitor.get_memory_report().unwrap();
193
194        monitor.reset_peak_usage();
195        monitor.update_internal_stats();
196        let after_reset = monitor.get_memory_report().unwrap();
197
198        // Sample count should increment
199        assert_eq!(after_reset.sample_count, before_reset.sample_count + 1);
200    }
201
202    #[test]
203    fn memory_monitor_integration_test() {
204        let monitor = MemoryMonitor::new("TestMemory".to_string());
205
206        // Test monitor identification
207        assert_eq!(monitor.monitor_id(), "TestMemory");
208        assert_eq!(monitor.resource_type(), MonitoredResourceType::SystemRam);
209
210        // Test memory tracking over time
211        monitor.update_internal_stats();
212        let updated_report = monitor.get_usage_report();
213        assert!(updated_report.peak_bytes.is_some());
214
215        // Test specific report methods
216        assert!(monitor.get_memory_report().is_some());
217    }
218}
khora_infra/telemetry/memory_monitor.rs

khora_infra/telemetry/
memory_monitor.rs
