rivet_logger/processors/
memory_peak_usage.rs1use 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}