Skip to main content

rivet_logger/processors/
memory_peak_usage.rs

1use std::sync::atomic::{AtomicU64, Ordering};
2
3use crate::logger::{BoxError, LogRecord, LogValue, Processor};
4
5use super::memory_common::{current_process_memory, format_memory, MemoryOptions};
6
7pub struct MemoryPeakUsage {
8    options: MemoryOptions,
9    peak: AtomicU64,
10}
11
12impl MemoryPeakUsage {
13    pub fn new(real_usage: bool, use_formatting: bool) -> Self {
14        Self {
15            options: MemoryOptions {
16                real_usage,
17                use_formatting,
18            },
19            peak: AtomicU64::new(0),
20        }
21    }
22}
23
24impl Default for MemoryPeakUsage {
25    fn default() -> Self {
26        Self::new(true, true)
27    }
28}
29
30impl Processor for MemoryPeakUsage {
31    fn process(&self, mut record: LogRecord) -> Result<LogRecord, BoxError> {
32        if let Some(bytes) = current_process_memory(self.options.real_usage) {
33            let mut current_peak = self.peak.load(Ordering::Relaxed);
34            while bytes > current_peak {
35                match self.peak.compare_exchange(
36                    current_peak,
37                    bytes,
38                    Ordering::SeqCst,
39                    Ordering::SeqCst,
40                ) {
41                    Ok(_) => {
42                        break;
43                    }
44                    Err(next) => current_peak = next,
45                }
46            }
47
48            let peak = self.peak.load(Ordering::Relaxed);
49            let value = if self.options.use_formatting {
50                format_memory(peak)
51            } else {
52                LogValue::U64(peak)
53            };
54            record.extra.insert("memory_peak_usage".to_string(), value);
55        }
56
57        Ok(record)
58    }
59}