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_or_else(|e| e.into_inner());
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
64            .peak_usage_bytes
65            .lock()
66            .unwrap_or_else(|e| e.into_inner());
67        *peak = current_usage;
68    }
69
70    /// Updates the monitor's internal state by querying the global allocator stats.
71    fn update_internal_stats(&self) {
72        let current_usage = get_currently_allocated_bytes();
73        let extended_stats = get_extended_memory_stats();
74
75        // Update peak tracking
76        let mut peak = self
77            .peak_usage_bytes
78            .lock()
79            .unwrap_or_else(|e| e.into_inner());
80        if current_usage > *peak {
81            *peak = current_usage;
82        }
83
84        // Calculate allocation delta
85        let mut last_alloc = self
86            .last_allocation_bytes
87            .lock()
88            .unwrap_or_else(|e| e.into_inner());
89        let allocation_delta = current_usage.saturating_sub(*last_alloc);
90        *last_alloc = current_usage;
91
92        // Update sample count
93        let mut count = self.sample_count.lock().unwrap_or_else(|e| e.into_inner());
94        *count += 1;
95
96        // Create comprehensive report with extended statistics
97        let report = MemoryReport {
98            current_usage_bytes: current_usage,
99            peak_usage_bytes: *peak,
100            allocation_delta_bytes: allocation_delta,
101            sample_count: *count,
102
103            // Extended statistics from allocator
104            total_allocations: extended_stats.total_allocations,
105            total_deallocations: extended_stats.total_deallocations,
106            total_reallocations: extended_stats.total_reallocations,
107            bytes_allocated_lifetime: extended_stats.bytes_allocated_lifetime,
108            bytes_deallocated_lifetime: extended_stats.bytes_deallocated_lifetime,
109            large_allocations: extended_stats.large_allocations,
110            large_allocation_bytes: extended_stats.large_allocation_bytes,
111            small_allocations: extended_stats.small_allocations,
112            small_allocation_bytes: extended_stats.small_allocation_bytes,
113            fragmentation_ratio: extended_stats.fragmentation_ratio,
114            allocation_efficiency: extended_stats.allocation_efficiency,
115            average_allocation_size: extended_stats.average_allocation_size,
116        };
117
118        let mut last_report = self.last_report.lock().unwrap_or_else(|e| e.into_inner());
119        *last_report = Some(report);
120    }
121}
122
123impl ResourceMonitor for MemoryMonitor {
124    fn monitor_id(&self) -> Cow<'static, str> {
125        Cow::Owned(self.id.clone())
126    }
127
128    fn resource_type(&self) -> MonitoredResourceType {
129        MonitoredResourceType::SystemRam
130    }
131
132    fn get_usage_report(&self) -> ResourceUsageReport {
133        let current_usage = get_currently_allocated_bytes();
134        let peak_usage = *self
135            .peak_usage_bytes
136            .lock()
137            .unwrap_or_else(|e| e.into_inner());
138
139        ResourceUsageReport {
140            current_bytes: current_usage as u64,
141            peak_bytes: Some(peak_usage as u64),
142            total_capacity_bytes: None, // System memory limit not easily available
143        }
144    }
145
146    fn get_metrics(
147        &self,
148    ) -> Vec<(
149        khora_core::telemetry::metrics::MetricId,
150        khora_core::telemetry::metrics::MetricValue,
151    )> {
152        use khora_core::telemetry::metrics::{MetricId, MetricValue};
153        let stats = get_extended_memory_stats();
154        // Named gauges flow through the standard push pipeline into the DCC's
155        // metric store, where the heuristics read `memory.current_bytes` for
156        // pressure + allocation-churn signals (the allocator finally has teeth).
157        vec![
158            (
159                MetricId::new("memory", "current_bytes"),
160                MetricValue::Gauge(stats.current_allocated_bytes as f64),
161            ),
162            (
163                MetricId::new("memory", "peak_bytes"),
164                MetricValue::Gauge(stats.peak_allocated_bytes as f64),
165            ),
166            (
167                MetricId::new("memory", "bytes_allocated_lifetime"),
168                MetricValue::Gauge(stats.bytes_allocated_lifetime as f64),
169            ),
170            (
171                MetricId::new("memory", "net_allocations"),
172                MetricValue::Gauge(stats.net_allocations as f64),
173            ),
174        ]
175    }
176
177    fn as_any(&self) -> &dyn std::any::Any {
178        self
179    }
180
181    fn update(&self) {
182        self.update_internal_stats();
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189
190    #[test]
191    fn memory_monitor_creation() {
192        let monitor = MemoryMonitor::new("TestMemory".to_string());
193        assert_eq!(monitor.monitor_id(), "TestMemory");
194        assert_eq!(monitor.resource_type(), MonitoredResourceType::SystemRam);
195    }
196
197    #[test]
198    fn memory_monitor_update_stats() {
199        let monitor = MemoryMonitor::new("TestMemory".to_string());
200
201        // Initially no report available until update is called
202        assert!(monitor.get_memory_report().is_none());
203
204        // Update stats
205        monitor.update_internal_stats();
206
207        // After update, report should be available
208        let report = monitor.get_memory_report().unwrap();
209        // In test environment, memory usage might be 0, so we just check report exists
210        assert_eq!(report.sample_count, 1);
211    }
212
213    #[test]
214    fn memory_monitor_peak_tracking() {
215        let monitor = MemoryMonitor::new("TestMemory".to_string());
216
217        monitor.update_internal_stats();
218        let _initial_report = monitor.get_memory_report().unwrap();
219        let _initial_peak = _initial_report.peak_usage_bytes;
220
221        // Reset peak and update again
222        monitor.reset_peak_usage();
223        monitor.update_internal_stats();
224
225        let after_reset_report = monitor.get_memory_report().unwrap();
226        // Peak should be reset to current usage
227        assert_eq!(after_reset_report.sample_count, 2);
228    }
229
230    #[test]
231    fn memory_monitor_reset_peak() {
232        let monitor = MemoryMonitor::new("TestMemory".to_string());
233
234        monitor.update_internal_stats();
235        let before_reset = monitor.get_memory_report().unwrap();
236
237        monitor.reset_peak_usage();
238        monitor.update_internal_stats();
239        let after_reset = monitor.get_memory_report().unwrap();
240
241        // Sample count should increment
242        assert_eq!(after_reset.sample_count, before_reset.sample_count + 1);
243    }
244
245    #[test]
246    fn memory_monitor_integration_test() {
247        let monitor = MemoryMonitor::new("TestMemory".to_string());
248
249        // Test monitor identification
250        assert_eq!(monitor.monitor_id(), "TestMemory");
251        assert_eq!(monitor.resource_type(), MonitoredResourceType::SystemRam);
252
253        // Test memory tracking over time
254        monitor.update_internal_stats();
255        let updated_report = monitor.get_usage_report();
256        assert!(updated_report.peak_bytes.is_some());
257
258        // Test specific report methods
259        assert!(monitor.get_memory_report().is_some());
260    }
261}
khora_infra/telemetry/memory_monitor.rs

khora_infra/telemetry/
memory_monitor.rs
