Skip to main content

runner_manager_platform/
memory_admission.rs

1//! Host memory admission on disk: its setting, what the daemon is waiting for,
2//! and the peaks this host measured.
3//!
4//! * **`config/memory.toml`** holds `host set-memory-margin`: whether admission
5//!   is on and its margin. An absent file is the default (on, 1 GiB). A file
6//!   rather than a column for the reason `polling.toml` gives: a schema step
7//!   makes the database refuse to open under an older build.
8//! * **`state/memory-admission.toml`** is written by the daemon whenever a
9//!   profile starts or stops waiting for memory, and read by `status`, which
10//!   runs in the operator's terminal while the daemon runs as a service.
11//! * **`state/memory-peaks.toml`** keeps each profile's newest measured peaks,
12//!   the host's own estimate when no coordinator gives one. It survives a
13//!   daemon restart so the estimate does not start over.
14
15use std::collections::BTreeMap;
16use std::path::PathBuf;
17
18use chrono::{DateTime, Utc};
19use runner_manager_domain::memory::{MemoryFigures, MemorySettings};
20use serde::{Deserialize, Serialize};
21
22use crate::paths::AppPaths;
23use crate::service::{ServiceError, read_state_record, write_state_record};
24
25pub const SETTINGS_FILE: &str = "memory.toml";
26pub const STATE_FILE: &str = "memory-admission.toml";
27pub const PEAKS_FILE: &str = "memory-peaks.toml";
28
29const SCHEMA_VERSION: u32 = 1;
30
31#[derive(Debug, Clone, Serialize, Deserialize)]
32struct Settings {
33    schema_version: u32,
34    enabled: bool,
35    margin_bytes: u64,
36}
37
38fn settings_path(paths: &AppPaths) -> PathBuf {
39    paths.config_dir().join(SETTINGS_FILE)
40}
41
42fn state_path(paths: &AppPaths) -> PathBuf {
43    paths.state_dir().join(STATE_FILE)
44}
45
46/// Where the measured peaks are kept.
47#[must_use]
48pub fn peaks_path(paths: &AppPaths) -> PathBuf {
49    paths.state_dir().join(PEAKS_FILE)
50}
51
52/// The configured setting, or the default when none was set.
53///
54/// # Errors
55/// [`ServiceError::Record`] when the file exists and cannot be read or parsed.
56pub fn settings(paths: &AppPaths) -> Result<MemorySettings, ServiceError> {
57    Ok(
58        read_state_record::<Settings>(&settings_path(paths))?.map_or_else(
59            MemorySettings::default,
60            |settings| MemorySettings {
61                enabled: settings.enabled,
62                margin_bytes: settings.margin_bytes,
63            },
64        ),
65    )
66}
67
68/// Stores the setting.
69///
70/// # Errors
71/// [`ServiceError::Record`] when `config/` cannot be written.
72pub fn set_settings(paths: &AppPaths, settings: MemorySettings) -> Result<(), ServiceError> {
73    write_state_record(
74        &settings_path(paths),
75        &Settings {
76            schema_version: SCHEMA_VERSION,
77            enabled: settings.enabled,
78            margin_bytes: settings.margin_bytes,
79        },
80    )
81}
82
83/// One profile waiting for memory, as `status` shows it.
84// `default` on every field: a CLI newer or older than the service's copy of
85// the binary must still read what the service wrote.
86#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
87#[serde(default)]
88pub struct Waiting {
89    pub policy_id: String,
90    /// `repo:owner/name` or `org:login`, as `status` names targets.
91    pub target: String,
92    pub profile: String,
93    /// The first refusal of this wait, and the latest.
94    pub since: Option<DateTime<Utc>>,
95    pub at: Option<DateTime<Utc>>,
96    pub free_bytes: u64,
97    pub reserved_bytes: u64,
98    pub estimate_bytes: u64,
99    /// `coordinator`, `local` or `default`.
100    pub estimate_source: String,
101    pub margin_bytes: u64,
102}
103
104impl Waiting {
105    /// A wait recorded at `at` from `figures`.
106    #[must_use]
107    pub fn new(
108        (policy_id, target, profile): (String, String, String),
109        figures: &MemoryFigures,
110        since: DateTime<Utc>,
111        at: DateTime<Utc>,
112    ) -> Self {
113        Self {
114            policy_id,
115            target,
116            profile,
117            since: Some(since),
118            at: Some(at),
119            free_bytes: figures.free,
120            reserved_bytes: figures.reserved,
121            estimate_bytes: figures.estimate.bytes,
122            estimate_source: figures.estimate.source.as_str().to_string(),
123            margin_bytes: figures.margin,
124        }
125    }
126
127    /// What a start needs available: the estimate and the margin.
128    #[must_use]
129    pub const fn needed_bytes(&self) -> u64 {
130        self.estimate_bytes.saturating_add(self.margin_bytes)
131    }
132
133    /// What a start can count on: free memory less the reservations.
134    #[must_use]
135    pub const fn available_bytes(&self) -> u64 {
136        self.free_bytes.saturating_sub(self.reserved_bytes)
137    }
138}
139
140/// What the daemon last wrote.
141#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
142#[serde(default)]
143pub struct AdmissionState {
144    pub schema_version: u32,
145    pub written_at: Option<DateTime<Utc>>,
146    pub waiting: Vec<Waiting>,
147}
148
149/// Records the profiles waiting for memory.
150///
151/// # Errors
152/// [`ServiceError::Record`] when `state/` cannot be written.
153pub fn record_waiting(
154    paths: &AppPaths,
155    waiting: Vec<Waiting>,
156    now: DateTime<Utc>,
157) -> Result<(), ServiceError> {
158    write_state_record(
159        &state_path(paths),
160        &AdmissionState {
161            schema_version: SCHEMA_VERSION,
162            written_at: Some(now),
163            waiting,
164        },
165    )
166}
167
168/// What the daemon last recorded; `None` before it ever wrote.
169///
170/// # Errors
171/// [`ServiceError::Record`] when the file exists and cannot be read or parsed.
172pub fn last_waiting(paths: &AppPaths) -> Result<Option<AdmissionState>, ServiceError> {
173    read_state_record(&state_path(paths))
174}
175
176/// Each profile's newest measured peaks, oldest first, keyed by policy id.
177#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
178#[serde(default)]
179pub struct Peaks {
180    pub schema_version: u32,
181    pub policies: BTreeMap<String, Vec<u64>>,
182}
183
184/// The measured peaks, empty when none were kept or the file is unreadable:
185/// an estimate is a guess either way, and the default stands in.
186#[must_use]
187pub fn load_peaks(path: &std::path::Path) -> Peaks {
188    read_state_record::<Peaks>(path)
189        .ok()
190        .flatten()
191        .unwrap_or_default()
192}
193
194/// Stores the measured peaks.
195///
196/// # Errors
197/// [`ServiceError::Record`] when `state/` cannot be written.
198pub fn save_peaks(path: &std::path::Path, peaks: &Peaks) -> Result<(), ServiceError> {
199    write_state_record(
200        path,
201        &Peaks {
202            schema_version: SCHEMA_VERSION,
203            policies: peaks.policies.clone(),
204        },
205    )
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211    use runner_manager_domain::memory::{Estimate, EstimateSource};
212
213    fn paths(dir: &tempfile::TempDir) -> AppPaths {
214        AppPaths::rooted_at(dir.path())
215    }
216
217    #[test]
218    fn the_setting_defaults_to_on_with_a_one_gib_margin_and_round_trips() {
219        let dir = tempfile::tempdir().expect("dir");
220        let paths = paths(&dir);
221        assert_eq!(
222            settings(&paths).expect("default"),
223            MemorySettings::default()
224        );
225        assert!(MemorySettings::default().enabled);
226        assert_eq!(MemorySettings::default().margin_bytes, 1 << 30);
227        let off = MemorySettings {
228            enabled: false,
229            margin_bytes: 2 << 30,
230        };
231        set_settings(&paths, off).expect("written");
232        assert_eq!(settings(&paths).expect("read"), off);
233    }
234
235    #[test]
236    fn the_waiting_record_and_the_peaks_round_trip() {
237        let dir = tempfile::tempdir().expect("dir");
238        let paths = paths(&dir);
239        assert_eq!(last_waiting(&paths).expect("absent"), None);
240        let figures = MemoryFigures {
241            free: 19 << 30,
242            reserved: 16 << 30,
243            estimate: Estimate {
244                bytes: 6 << 30,
245                source: EstimateSource::Coordinator,
246            },
247            margin: 1 << 30,
248        };
249        let now = Utc::now();
250        let waiting = Waiting::new(
251            ("p".into(), "repo:o/r".into(), "default".into()),
252            &figures,
253            now,
254            now,
255        );
256        assert_eq!(waiting.needed_bytes(), 7 << 30);
257        assert_eq!(waiting.available_bytes(), 3 << 30);
258        record_waiting(&paths, vec![waiting.clone()], now).expect("written");
259        let state = last_waiting(&paths).expect("read").expect("present");
260        assert_eq!(state.waiting, vec![waiting]);
261
262        let file = peaks_path(&paths);
263        assert_eq!(load_peaks(&file), Peaks::default());
264        let mut peaks = Peaks::default();
265        peaks.policies.insert("p".into(), vec![1, 2, 3]);
266        save_peaks(&file, &peaks).expect("written");
267        assert_eq!(load_peaks(&file).policies, peaks.policies);
268    }
269}