1use 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#[must_use]
48pub fn peaks_path(paths: &AppPaths) -> PathBuf {
49 paths.state_dir().join(PEAKS_FILE)
50}
51
52pub 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
68pub 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#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
87#[serde(default)]
88pub struct Waiting {
89 pub policy_id: String,
90 pub target: String,
92 pub profile: String,
93 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 pub estimate_source: String,
101 pub margin_bytes: u64,
102}
103
104impl Waiting {
105 #[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 #[must_use]
129 pub const fn needed_bytes(&self) -> u64 {
130 self.estimate_bytes.saturating_add(self.margin_bytes)
131 }
132
133 #[must_use]
135 pub const fn available_bytes(&self) -> u64 {
136 self.free_bytes.saturating_sub(self.reserved_bytes)
137 }
138}
139
140#[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
149pub 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
168pub fn last_waiting(paths: &AppPaths) -> Result<Option<AdmissionState>, ServiceError> {
173 read_state_record(&state_path(paths))
174}
175
176#[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#[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
194pub 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}