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});