polars_utils/
sys.rs

1use std::sync::{LazyLock, Mutex};
2
3use sysinfo::{MemoryRefreshKind, System};
4
5use crate::config::verbose;
6
7/// Return the total system memory in bytes.
8pub fn total_memory() -> u64 {
9    return *TOTAL_MEMORY;
10
11    static TOTAL_MEMORY: LazyLock<u64> = LazyLock::new(|| {
12        let mut sys = System::new();
13
14        sys.refresh_memory_specifics(MemoryRefreshKind::nothing().with_ram());
15
16        let mut v: u64 = match sys.cgroup_limits() {
17            Some(limits) => limits.total_memory,
18            None => sys.total_memory(),
19        };
20
21        if let Ok(s) = std::env::var("POLARS_OVERRIDE_TOTAL_MEMORY") {
22            v = s
23                .parse::<u64>()
24                .unwrap_or_else(|_| panic!("invalid value for POLARS_OVERRIDE_TOTAL_MEMORY: {s}"))
25        }
26
27        if verbose() {
28            let gib = (v as f64) / (1024.0 * 1024.0 * 1024.0);
29            eprintln!("total memory: {gib:.3} GiB ({v} bytes)")
30        }
31
32        v
33    });
34}
35
36/// Startup system is expensive, so we do it once
37pub struct MemInfo {
38    sys: Mutex<System>,
39}
40
41impl MemInfo {
42    /// This call is quite expensive, cache the results.
43    pub fn free(&self) -> u64 {
44        let mut sys = self.sys.lock().unwrap();
45        sys.refresh_memory();
46        match sys.cgroup_limits() {
47            Some(limits) => limits.free_memory,
48            None => sys.available_memory(),
49        }
50    }
51}
52
53pub static MEMINFO: LazyLock<MemInfo> = LazyLock::new(|| MemInfo {
54    sys: Mutex::new(System::new()),
55});
56
57/// Check whether a process with the given PID is currently alive.
58///
59/// Used by `polars_ooc::cleaner::cleanup_stale_dirs` to remove spill
60/// directories left behind by dead processes on startup.
61pub fn is_process_alive(pid: u32) -> bool {
62    use sysinfo::{Pid, ProcessRefreshKind, System, UpdateKind};
63    let pid = Pid::from_u32(pid);
64    let mut sys = System::new();
65    sys.refresh_processes_specifics(
66        sysinfo::ProcessesToUpdate::Some(&[pid]),
67        true,
68        ProcessRefreshKind::nothing().with_cmd(UpdateKind::Never),
69    );
70    sys.process(pid).is_some()
71}
polars_utils/sys.rs

polars_utils/
sys.rs
