Skip to main content

tale_ndjson/metrics/
memory.rs

1//! Pressure pushing down on me pressing down on you
2//! no man ask for under pressure
3//! that burns a building down
4//! splits a family in two
5//! puts people on streets
6
7pub use memimpl::*;
8
9use crate::defaults::memory::MEMORY_LIMIT_BYTES;
10
11#[derive(Debug, Clone, Copy)]
12pub enum MemoryPressure {
13    Unknown,
14    None,     // < 50% usage
15    Low,      // 50-70%
16    Medium,   // 70-85%
17    High,     // 85-95%
18    Critical, // > 95%
19}
20
21pub fn detect_memory_pressure(limit_bytes: Option<usize>) -> MemoryPressure {
22    memimpl::detect_memory_pressure(limit_bytes)
23}
24
25impl From<f64> for MemoryPressure {
26    fn from(p: f64) -> Self {
27        match p {
28            p if p < 50.0 => MemoryPressure::None,
29            p if p < 70.0 => MemoryPressure::Low,
30            p if p < 85.0 => MemoryPressure::Medium,
31            p if p < 95.0 => MemoryPressure::High,
32            _ => MemoryPressure::Critical,
33        }
34    }
35}
36
37pub mod memimpl {
38    use std::sync::{LazyLock, Mutex, OnceLock};
39
40    use sysinfo::{CpuRefreshKind, MemoryRefreshKind, Pid, ProcessRefreshKind, ProcessesToUpdate, RefreshKind, System};
41
42    use super::*;
43
44    static SYSTEM: LazyLock<Mutex<System>> = LazyLock::new(|| {
45        let processes = ProcessRefreshKind::nothing().with_cpu().with_memory();
46        let memory = MemoryRefreshKind::nothing().with_ram();
47        let cpu = CpuRefreshKind::nothing(); // .with_cpu_usage();
48        let refreshes = RefreshKind::nothing()
49            .with_cpu(cpu)
50            .with_memory(memory)
51            .with_processes(processes);
52        Mutex::new(sysinfo::System::new_with_specifics(refreshes))
53    });
54    pub fn system() -> std::sync::MutexGuard<'static, System> {
55        SYSTEM
56            .lock()
57            .expect("Unrecoverable runtime problem: cannot acquire system lock. Exiting.")
58    }
59
60    static PID: OnceLock<Pid> = OnceLock::new();
61    pub fn pid() -> Option<Pid> {
62        if let Some(pid) = PID.get() {
63            Some(*pid)
64        } else {
65            match sysinfo::get_current_pid() {
66                Ok(p) => {
67                    PID.set(p).expect("unexpected failure to store PID in oncelock");
68                    Some(p)
69                }
70                Err(_) => None,
71            }
72        }
73    }
74
75    /// Are we under memory pressure?
76    pub fn detect_memory_pressure(max_allowed_bytes: Option<usize>) -> MemoryPressure {
77        // Return appropriate level
78        let Some(pid) = pid() else {
79            return MemoryPressure::Unknown;
80        };
81        let proc_list = vec![pid];
82        let to_update = ProcessesToUpdate::Some(proc_list.as_slice());
83        let mut system = system();
84        system.refresh_processes(to_update, true);
85        let Some(process) = system.process(pid) else {
86            return MemoryPressure::Unknown;
87        };
88        let rss = process.memory() as usize;
89        let max_allowed =
90            max_allowed_bytes.unwrap_or_else(|| crate::config::config().max_memory.unwrap_or(MEMORY_LIMIT_BYTES));
91        let free = system.free_memory() as usize;
92        // let total = system.total_memory();
93        // eprintln!("memory: rss={rss}; max_allowed={max_allowed}; free={free};");
94
95        if rss > max_allowed || rss > free {
96            MemoryPressure::Critical
97        } else {
98            let percent = (rss * 100) as f64 / max_allowed as f64;
99            MemoryPressure::from(percent)
100        }
101    }
102
103    pub fn get_system_ram_bytes() -> usize {
104        let system = system();
105        system.total_memory() as usize
106    }
107
108    pub fn available_memory_bytes() -> usize {
109        let Some(pid) = pid() else {
110            return 0;
111        };
112        let proc_list = vec![pid];
113        let to_update = ProcessesToUpdate::Some(proc_list.as_slice());
114        let mut system = system();
115        system.refresh_processes(to_update, true);
116        let free = system.free_memory() as usize;
117        // let total = system.total_memory() as usize;
118
119        let Some(process) = system.process(pid) else {
120            return free as usize;
121        };
122        let rss = process.memory() as usize;
123        let max_allowed = crate::config::config().max_memory.unwrap_or(MEMORY_LIMIT_BYTES);
124
125        // eprintln!("system: total={total} free={free}; process: rss={rss}; max:
126        // {max_allowed}");
127
128        let remaining_budget = max_allowed.saturating_sub(rss);
129        std::cmp::min(free, remaining_budget)
130    }
131
132    pub fn process_memory_bytes() -> usize {
133        let Some(pid) = pid() else {
134            return 0;
135        };
136        let system = system();
137        let Some(process) = system.process(pid) else {
138            return 0;
139        };
140        process.memory() as usize
141    }
142}