starship/modules/
memory_usage.rs

1use systemstat::{
2    Platform, System,
3    data::{ByteSize, saturating_sub_bytes},
4};
5
6use super::{Context, Module, ModuleConfig};
7
8use crate::configs::memory_usage::MemoryConfig;
9use crate::formatter::StringFormatter;
10
11// Display a `ByteSize` in a human readable format.
12fn display_bs(bs: ByteSize) -> String {
13    let mut display_bytes = bs.to_string_as(true);
14    let mut keep = true;
15    // Skip decimals and the space before the byte unit.
16    display_bytes.retain(|c| match c {
17        ' ' => {
18            keep = true;
19            false
20        }
21        '.' => {
22            keep = false;
23            false
24        }
25        _ => keep,
26    });
27    display_bytes
28}
29
30// Calculate the memory usage from total and free memory
31fn pct(total: ByteSize, free: ByteSize) -> f64 {
32    100.0 * saturating_sub_bytes(total, free).0 as f64 / total.0 as f64
33}
34
35// Print usage string used/total
36fn format_usage_total(total: ByteSize, free: ByteSize) -> String {
37    format!(
38        "{}/{}",
39        display_bs(saturating_sub_bytes(total, free)),
40        display_bs(total)
41    )
42}
43
44/// Creates a module with system memory usage information
45pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
46    let mut module = context.new_module("memory_usage");
47    let config = MemoryConfig::try_load(module.config);
48
49    // As we default to disabled=true, we have to check here after loading our config module,
50    // before it was only checking against whatever is in the config starship.toml
51    if config.disabled {
52        return None;
53    }
54
55    let system = System::new();
56
57    // `memory_and_swap` only works on platforms that have an implementation for swap memory
58    // But getting both together is faster on some platforms (Windows/Linux)
59    let (memory, swap) = match system.memory_and_swap() {
60        // Ignore swap if total is 0
61        Ok((mem, swap)) if swap.total.0 > 0 => (mem, Some(swap)),
62        Ok((mem, _)) => (mem, None),
63        Err(e) => {
64            log::debug!(
65                "Failed to retrieve both memory and swap, falling back to memory only: {e}"
66            );
67            let mem = match system.memory() {
68                Ok(mem) => mem,
69                Err(e) => {
70                    log::warn!("Failed to retrieve memory: {e}");
71                    return None;
72                }
73            };
74
75            (mem, None)
76        }
77    };
78
79    let used_pct = pct(memory.total, memory.free);
80
81    if (used_pct.round() as i64) < config.threshold {
82        return None;
83    }
84
85    let parsed = StringFormatter::new(config.format).and_then(|formatter| {
86        formatter
87            .map_meta(|var, _| match var {
88                "symbol" => Some(config.symbol),
89                _ => None,
90            })
91            .map_style(|variable| match variable {
92                "style" => Some(Ok(config.style)),
93                _ => None,
94            })
95            .map(|variable| match variable {
96                "ram" => Some(Ok(format_usage_total(memory.total, memory.free))),
97                "ram_pct" => Some(Ok(format!("{used_pct:.0}%"))),
98                "swap" => Some(Ok(format_usage_total(
99                    swap.as_ref()?.total,
100                    swap.as_ref()?.free,
101                ))),
102                "swap_pct" => Some(Ok(format!(
103                    "{:.0}%",
104                    pct(swap.as_ref()?.total, swap.as_ref()?.free)
105                ))),
106                _ => None,
107            })
108            .parse(None, Some(context))
109    });
110
111    module.set_segments(match parsed {
112        Ok(segments) => segments,
113        Err(error) => {
114            log::warn!("Error in module `memory_usage`:\n{error}");
115            return None;
116        }
117    });
118
119    Some(module)
120}
121
122#[cfg(test)]
123mod test {
124    use super::*;
125
126    use crate::test::ModuleRenderer;
127
128    #[test]
129    fn test_format_usage_total() {
130        assert_eq!(
131            format_usage_total(ByteSize(1024 * 1024 * 1024), ByteSize(1024 * 1024 * 1024)),
132            "0B/1GiB"
133        );
134        assert_eq!(
135            format_usage_total(
136                ByteSize(1024 * 1024 * 1024),
137                ByteSize(1024 * 1024 * 1024 / 2)
138            ),
139            "512MiB/1GiB"
140        );
141        assert_eq!(
142            format_usage_total(ByteSize(1024 * 1024 * 1024), ByteSize(0)),
143            "1GiB/1GiB"
144        );
145    }
146
147    #[test]
148    fn test_pct() {
149        assert_eq!(
150            pct(ByteSize(1024 * 1024 * 1024), ByteSize(1024 * 1024 * 1024)),
151            0.0
152        );
153        assert_eq!(
154            pct(
155                ByteSize(1024 * 1024 * 1024),
156                ByteSize(1024 * 1024 * 1024 / 2)
157            ),
158            50.0
159        );
160        assert_eq!(pct(ByteSize(1024 * 1024 * 1024), ByteSize(0)), 100.0);
161    }
162
163    #[test]
164    fn zero_threshold() {
165        let output = ModuleRenderer::new("memory_usage")
166            .config(toml::toml! {
167                [memory_usage]
168                disabled = false
169                threshold = 0
170            })
171            .collect();
172
173        assert!(output.is_some());
174    }
175
176    #[test]
177    fn impossible_threshold() {
178        let output = ModuleRenderer::new("memory_usage")
179            .config(toml::toml! {
180                [memory_usage]
181                disabled = false
182                threshold = 9999
183            })
184            .collect();
185
186        assert!(output.is_none());
187    }
188}
starship/modules/memory_usage.rs

starship/modules/
memory_usage.rs
