1//! How much memory a new program could use right now, for the headroom gate.
2//!
3//! The built-in measure is the kernel's own (`MemAvailable` on Linux). That
4//! is the whole machine's truth on bare metal, but not everywhere: a virtual
5//! machine sees its own memory, not its host's, and other systems have no
6//! such file. So the measure can be replaced by a command in the config
7//! file, and where there is neither, the load of a command is simply not
8//! judged. Nothing here assumes a particular host.
9
10use std::path::PathBuf;
11use std::sync::{Mutex, PoisonError};
12use std::time::{Duration, Instant};
13
14use serde::Deserialize;
15use serde_json::{Value, json};
16
17/// A measurement is reused for this long: commands arrive in bursts, and a
18/// configured command may cost more than reading a file.
19const FRESH_FOR: Duration = Duration::from_secs(5);
20
21/// How long a configured command may take before it counts as failed.
22const COMMAND_WITHIN: Duration = Duration::from_secs(2);
23
24/// `$XDG_CONFIG_HOME/jevhooks/config.toml`, or `~/.config/jevhooks/config.toml`.
25pub fn config_path() -> PathBuf {
26    let base = std::env::var_os("XDG_CONFIG_HOME")
27        .filter(|v| !v.is_empty())
28        .map(PathBuf::from)
29        .or_else(|| std::env::var_os("HOME").map(|home| PathBuf::from(home).join(".config")))
30        .unwrap_or_else(|| PathBuf::from("."));
31    base.join("jevhooks").join("config.toml")
32}
33
34/// The daemon's config file. Every key is optional; an unknown key is an
35/// error, so a misspelt setting is reported instead of ignored.
36#[derive(Debug, Default, Deserialize)]
37#[serde(deny_unknown_fields)]
38pub struct Config {
39    /// A shell command that prints, as its first number, the megabytes of
40    /// memory a new program could use. Replaces the built-in measure.
41    pub available_memory_command: Option<String>,
42}
43
44impl Config {
45    /// The config file's settings; the defaults when there is no file; an
46    /// error when there is one that cannot be read or understood.
47    pub fn load() -> Result<Self, String> {
48        let path = config_path();
49        match std::fs::read_to_string(&path) {
50            Ok(text) => toml::from_str(&text).map_err(|e| format!("{}: {e}", path.display())),
51            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(Self::default()),
52            Err(e) => Err(format!("{}: {e}", path.display())),
53        }
54    }
55}
56
57/// Where the available-memory figure comes from. Exactly one source: a
58/// configured command is never silently replaced by the built-in measure,
59/// because the user configured it to say the built-in one is wrong here.
60enum Source {
61    Command(String),
62    /// `MemAvailable` in `/proc/meminfo`.
63    Meminfo,
64    /// No measure on this system and none configured.
65    None,
66}
67
68/// One measurement: megabytes, or why there is no figure.
69type Measured = Result<f64, String>;
70
71pub struct Headroom {
72    source: Source,
73    last: Mutex<Option<(Instant, Measured)>>,
74}
75
76/// The first number in `text`, as megabytes.
77fn first_number(text: &str) -> Option<f64> {
78    text.split(|c: char| !(c.is_ascii_digit() || c == '.')).find(|part| !part.is_empty())?.parse().ok()
79}
80
81fn meminfo_mb() -> Measured {
82    let meminfo = std::fs::read_to_string("/proc/meminfo").map_err(|e| format!("/proc/meminfo: {e}"))?;
83    let line = meminfo.lines().find(|line| line.starts_with("MemAvailable:")).ok_or("no MemAvailable in /proc/meminfo")?;
84    let kilobytes: f64 =
85        line.split_whitespace().nth(1).and_then(|n| n.parse().ok()).ok_or("MemAvailable is not a number")?;
86    Ok(kilobytes / 1024.0)
87}
88
89async fn command_mb(command: &str) -> Measured {
90    let run = tokio::process::Command::new("sh").arg("-c").arg(command).kill_on_drop(true).output();
91    let output = tokio::time::timeout(COMMAND_WITHIN, run)
92        .await
93        .map_err(|_| format!("available_memory_command did not finish within {} s", COMMAND_WITHIN.as_secs()))?
94        .map_err(|e| format!("available_memory_command could not run: {e}"))?;
95    if !output.status.success() {
96        return Err(format!("available_memory_command exited with {}", output.status));
97    }
98    first_number(&String::from_utf8_lossy(&output.stdout))
99        .ok_or_else(|| "available_memory_command printed no number".to_owned())
100}
101
102impl Headroom {
103    pub fn new(config: &Config) -> Self {
104        let source = match &config.available_memory_command {
105            Some(command) => Source::Command(command.clone()),
106            None if std::path::Path::new("/proc/meminfo").exists() => Source::Meminfo,
107            None => Source::None,
108        };
109        Self { source, last: Mutex::new(None) }
110    }
111
112    async fn measure(&self) -> Measured {
113        let fresh = self.last.lock().unwrap_or_else(PoisonError::into_inner).clone();
114        if let Some((at, measured)) = fresh
115            && at.elapsed() < FRESH_FOR
116        {
117            return measured;
118        }
119        let measured = match &self.source {
120            Source::Command(command) => command_mb(command).await,
121            Source::Meminfo => meminfo_mb(),
122            Source::None => Err("no built-in measure on this system; set available_memory_command".to_owned()),
123        };
124        *self.last.lock().unwrap_or_else(PoisonError::into_inner) = Some((Instant::now(), measured.clone()));
125        measured
126    }
127
128    /// Megabytes available now, or nothing when it cannot be measured, in
129    /// which case a command's load is not judged.
130    pub async fn available_mb(&self) -> Option<f64> {
131        self.measure().await.ok()
132    }
133
134    /// For `status`: where the figure comes from and what it is now.
135    pub async fn status(&self) -> Value {
136        let source = match &self.source {
137            Source::Command(command) => format!("command: {command}"),
138            Source::Meminfo => "MemAvailable in /proc/meminfo".to_owned(),
139            Source::None => "none".to_owned(),
140        };
141        match self.measure().await {
142            Ok(available) => json!({ "source": source, "available_mb": available.round() }),
143            Err(why) => json!({ "source": source, "error": why }),
144        }
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn the_first_number_is_read_whatever_surrounds_it() {
154        assert_eq!(first_number("8123\n"), Some(8123.0));
155        assert_eq!(first_number("available: 5120.5 MB of 16384"), Some(5120.5));
156        assert_eq!(first_number("none"), None);
157        assert_eq!(first_number(""), None);
158    }
159
160    #[test]
161    fn a_misspelt_setting_is_an_error() {
162        assert!(toml::from_str::<Config>("available_memory_comand = \"x\"").is_err());
163        assert!(toml::from_str::<Config>("").unwrap().available_memory_command.is_none());
164    }
165
166    #[tokio::test]
167    async fn a_configured_command_is_the_only_source() {
168        let config = Config { available_memory_command: Some("echo 1234".into()) };
169        assert_eq!(Headroom::new(&config).available_mb().await, Some(1234.0));
170        // A failing command gives no figure; it does not fall back to the built-in measure.
171        let config = Config { available_memory_command: Some("exit 3".into()) };
172        assert_eq!(Headroom::new(&config).available_mb().await, None);
173    }
174}