polars_utils/
sys.rs

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

polars_utils/
sys.rs
