1use std::fs;
10use std::io::Write as _;
11use std::path::{Path, PathBuf};
12
13use chrono::{DateTime, Utc};
14use serde::{Deserialize, Serialize};
15
16use super::discovery::{escaped_name_with_digest, validate_distribution_name};
17use crate::paths::AppPaths;
18use runner_manager_domain::model::ScaleTarget;
19
20pub const GUEST_CONFIG_FILE: &str = "wsl-recovery.toml";
21pub const REQUEST_FILE: &str = "drain-request.json";
22pub const HEARTBEAT_FILE: &str = "guest-heartbeat.json";
23pub const FENCE_DIRECTORY: &str = "launch-fence";
24pub const OWNER_FILE: &str = "owner.json";
25pub const RECOVERY_STATUS_FILE: &str = "recovery-status.json";
26pub const SCHEMA_VERSION: u32 = 1;
27
28pub fn recovery_root(paths: &AppPaths, distribution: &str) -> Result<PathBuf, super::WslError> {
30 validate_distribution_name(distribution)?;
31 Ok(paths
32 .config_dir()
33 .join("wsl-recovery")
34 .join(escaped_name_with_digest(distribution)))
35}
36
37#[derive(Debug, thiserror::Error)]
38pub enum FenceError {
39 #[error("cannot {operation} WSL recovery state at {}: {source}", path.display())]
40 Io {
41 operation: &'static str,
42 path: PathBuf,
43 #[source]
44 source: std::io::Error,
45 },
46 #[error("cannot decode WSL recovery state at {}: {source}", path.display())]
47 Decode {
48 path: PathBuf,
49 #[source]
50 source: serde_json::Error,
51 },
52 #[error("WSL recovery state at {} has schema {found}, but this build supports {SCHEMA_VERSION}", path.display())]
53 Schema { path: PathBuf, found: u32 },
54}
55
56fn io(operation: &'static str, path: &Path, source: std::io::Error) -> FenceError {
57 FenceError::Io {
58 operation,
59 path: path.to_path_buf(),
60 source,
61 }
62}
63
64#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
65#[serde(deny_unknown_fields)]
66pub struct GuestRecoveryConfig {
67 pub schema_version: u32,
68 pub shared_root: PathBuf,
69}
70
71impl GuestRecoveryConfig {
72 #[must_use]
73 pub fn new(shared_root: PathBuf) -> Self {
74 Self {
75 schema_version: SCHEMA_VERSION,
76 shared_root,
77 }
78 }
79
80 #[must_use]
81 pub fn path(paths: &AppPaths) -> PathBuf {
82 paths.config_dir().join(GUEST_CONFIG_FILE)
83 }
84
85 pub fn read(paths: &AppPaths) -> Result<Option<Self>, FenceError> {
86 let path = Self::path(paths);
87 let text = match fs::read_to_string(&path) {
88 Ok(text) => text,
89 Err(source) if source.kind() == std::io::ErrorKind::NotFound => return Ok(None),
90 Err(source) => return Err(io("read", &path, source)),
91 };
92 let value: Self = toml::from_str(&text).map_err(|source| FenceError::Io {
93 operation: "decode",
94 path: path.clone(),
95 source: std::io::Error::new(std::io::ErrorKind::InvalidData, source),
96 })?;
97 if value.schema_version != SCHEMA_VERSION {
98 return Err(FenceError::Schema {
99 path,
100 found: value.schema_version,
101 });
102 }
103 Ok(Some(value))
104 }
105
106 pub fn write(&self, paths: &AppPaths) -> Result<(), FenceError> {
107 let path = Self::path(paths);
108 let text = toml::to_string_pretty(self).map_err(|source| FenceError::Io {
109 operation: "encode",
110 path: path.clone(),
111 source: std::io::Error::new(std::io::ErrorKind::InvalidData, source),
112 })?;
113 atomic_write(&path, text.as_bytes())
114 }
115}
116
117#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
118#[serde(deny_unknown_fields)]
119pub struct DrainRequest {
120 pub schema_version: u32,
121 pub generation: u64,
122 pub requested_at: DateTime<Utc>,
123}
124
125impl DrainRequest {
126 #[must_use]
127 pub fn new(generation: u64, requested_at: DateTime<Utc>) -> Self {
128 Self {
129 schema_version: SCHEMA_VERSION,
130 generation,
131 requested_at,
132 }
133 }
134
135 pub fn write(&self, root: &Path) -> Result<(), FenceError> {
136 write_json(&root.join(REQUEST_FILE), self)
137 }
138
139 pub fn read(root: &Path) -> Result<Option<Self>, FenceError> {
140 read_json(&root.join(REQUEST_FILE))
141 }
142}
143
144#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
145#[serde(deny_unknown_fields)]
146pub struct GuestHeartbeat {
147 pub schema_version: u32,
148 pub observed_at: DateTime<Utc>,
149 pub acknowledged_generation: Option<u64>,
150 pub local_active_attempts: Option<u32>,
151 pub managed_targets: Vec<ScaleTarget>,
152 pub unmanaged_runner_services: Option<u32>,
153}
154
155#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
156#[serde(rename_all = "snake_case")]
157pub enum RecoveryPhase {
158 Healthy,
159 Degraded,
160 Draining,
161 Recovering,
162 Backoff,
163 RecoveryBlocked,
164}
165
166#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
167#[serde(deny_unknown_fields)]
168pub struct RecoveryStatus {
169 pub schema_version: u32,
170 pub observed_at: DateTime<Utc>,
171 pub phase: RecoveryPhase,
172 pub consecutive_probe_failures: u8,
173 pub reason: Option<String>,
174 pub last_recovered_at: Option<DateTime<Utc>>,
175}
176
177impl RecoveryStatus {
178 pub fn write(&self, root: &Path) -> Result<(), FenceError> {
179 write_json(&root.join(RECOVERY_STATUS_FILE), self)
180 }
181
182 pub fn read(root: &Path) -> Result<Option<Self>, FenceError> {
183 read_json(&root.join(RECOVERY_STATUS_FILE))
184 }
185}
186
187impl GuestHeartbeat {
188 pub fn write(&self, root: &Path) -> Result<(), FenceError> {
189 write_json(&root.join(HEARTBEAT_FILE), self)
190 }
191
192 pub fn read(root: &Path) -> Result<Option<Self>, FenceError> {
193 read_json(&root.join(HEARTBEAT_FILE))
194 }
195}
196
197#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
198#[serde(rename_all = "snake_case")]
199pub enum FenceOwnerKind {
200 GuestLaunch,
201 WindowsRecovery,
202}
203
204#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
205#[serde(deny_unknown_fields)]
206pub struct FenceOwner {
207 pub schema_version: u32,
208 pub kind: FenceOwnerKind,
209 pub generation: Option<u64>,
210 pub process_id: u32,
211 pub acquired_at: DateTime<Utc>,
212}
213
214#[derive(Debug)]
216pub struct FenceClaim {
217 directory: PathBuf,
218 release_on_drop: bool,
219}
220
221impl FenceClaim {
222 pub fn try_claim(
225 root: &Path,
226 kind: FenceOwnerKind,
227 generation: Option<u64>,
228 ) -> Result<Option<Self>, FenceError> {
229 fs::create_dir_all(root).map_err(|source| io("create", root, source))?;
230 let directory = root.join(FENCE_DIRECTORY);
231 match fs::create_dir(&directory) {
232 Ok(()) => {}
233 Err(source) if source.kind() == std::io::ErrorKind::AlreadyExists => return Ok(None),
234 Err(source) => return Err(io("claim", &directory, source)),
235 }
236 let owner = FenceOwner {
237 schema_version: SCHEMA_VERSION,
238 kind,
239 generation,
240 process_id: std::process::id(),
241 acquired_at: Utc::now(),
242 };
243 if let Err(error) = write_json(&directory.join(OWNER_FILE), &owner) {
244 let _ = fs::remove_dir_all(&directory);
245 return Err(error);
246 }
247 Ok(Some(Self {
248 directory,
249 release_on_drop: true,
250 }))
251 }
252
253 pub fn make_durable(mut self) {
256 self.release_on_drop = false;
257 }
258
259 pub fn owner(root: &Path) -> Result<Option<FenceOwner>, FenceError> {
260 read_json(&root.join(FENCE_DIRECTORY).join(OWNER_FILE))
261 }
262
263 pub fn release(mut self) -> Result<(), FenceError> {
264 self.release_on_drop = false;
265 remove_claim(&self.directory)
266 }
267}
268
269impl Drop for FenceClaim {
270 fn drop(&mut self) {
271 if self.release_on_drop {
272 let _ = remove_claim(&self.directory);
273 }
274 }
275}
276
277pub fn clear_recovery(root: &Path, generation: u64) -> Result<(), FenceError> {
278 let owner = FenceClaim::owner(root)?;
279 if owner.as_ref().is_some_and(|owner| {
280 owner.kind == FenceOwnerKind::WindowsRecovery && owner.generation == Some(generation)
281 }) {
282 remove_claim(&root.join(FENCE_DIRECTORY))?;
283 }
284 let request_path = root.join(REQUEST_FILE);
285 if DrainRequest::read(root)?
286 .as_ref()
287 .is_some_and(|request| request.generation != generation)
288 {
289 return Ok(());
290 }
291 match fs::remove_file(&request_path) {
292 Ok(()) => Ok(()),
293 Err(source) if source.kind() == std::io::ErrorKind::NotFound => Ok(()),
294 Err(source) => Err(io("remove", &request_path, source)),
295 }
296}
297
298pub fn retire_windows_recovery(root: &Path) -> Result<(), FenceError> {
303 let recovery_generation = FenceClaim::owner(root)?.and_then(|owner| {
304 (owner.kind == FenceOwnerKind::WindowsRecovery)
305 .then_some(owner.generation)
306 .flatten()
307 });
308 if let Some(generation) = recovery_generation {
309 clear_recovery(root, generation)?;
310 }
311 if let Some(request) = DrainRequest::read(root)? {
312 clear_recovery(root, request.generation)?;
313 }
314 let status_path = root.join(RECOVERY_STATUS_FILE);
315 match fs::remove_file(&status_path) {
316 Ok(()) => Ok(()),
317 Err(source) if source.kind() == std::io::ErrorKind::NotFound => Ok(()),
318 Err(source) => Err(io("remove", &status_path, source)),
319 }
320}
321
322fn remove_claim(directory: &Path) -> Result<(), FenceError> {
323 match fs::remove_dir_all(directory) {
324 Ok(()) => Ok(()),
325 Err(source) if source.kind() == std::io::ErrorKind::NotFound => Ok(()),
326 Err(source) => Err(io("release", directory, source)),
327 }
328}
329
330fn atomic_write(path: &Path, bytes: &[u8]) -> Result<(), FenceError> {
331 let parent = path.parent().unwrap_or_else(|| Path::new("."));
332 fs::create_dir_all(parent).map_err(|source| io("create", parent, source))?;
333 let mut temporary =
334 tempfile::NamedTempFile::new_in(parent).map_err(|source| io("write", path, source))?;
335 temporary
336 .write_all(bytes)
337 .and_then(|()| temporary.as_file().sync_all())
338 .map_err(|source| io("write", path, source))?;
339 temporary
340 .persist(path)
341 .map(|_| ())
342 .map_err(|error| io("replace", path, error.error))
343}
344
345fn write_json<T: Serialize>(path: &Path, value: &T) -> Result<(), FenceError> {
346 let bytes = serde_json::to_vec_pretty(value).map_err(|source| FenceError::Decode {
347 path: path.to_path_buf(),
348 source,
349 })?;
350 atomic_write(path, &bytes)
351}
352
353fn read_json<T: for<'de> Deserialize<'de>>(path: &Path) -> Result<Option<T>, FenceError> {
354 let bytes = match fs::read(path) {
355 Ok(bytes) => bytes,
356 Err(source) if source.kind() == std::io::ErrorKind::NotFound => return Ok(None),
357 Err(source) => return Err(io("read", path, source)),
358 };
359 let value: T = serde_json::from_slice(&bytes).map_err(|source| FenceError::Decode {
360 path: path.to_path_buf(),
361 source,
362 })?;
363 Ok(Some(value))
364}
365
366#[must_use]
369pub fn unmanaged_runner_service_count() -> Option<u32> {
370 if !cfg!(target_os = "linux") {
371 return Some(0);
372 }
373 let mut names = std::collections::BTreeSet::new();
374 for directory in [
375 "/etc/systemd/system",
376 "/usr/lib/systemd/system",
377 "/lib/systemd/system",
378 ] {
379 let entries = match fs::read_dir(directory) {
380 Ok(entries) => entries,
381 Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue,
382 Err(_) => return None,
383 };
384 for entry in entries.flatten() {
385 let name = entry.file_name().to_string_lossy().into_owned();
386 if name.starts_with("actions.runner.") && name.ends_with(".service") {
387 names.insert(name);
388 }
389 }
390 }
391 u32::try_from(names.len()).ok()
392}
393
394#[cfg(test)]
395mod tests {
396 use super::*;
397
398 #[test]
399 fn one_directory_has_exactly_one_owner_and_drop_releases_guest_claim() {
400 let root = tempfile::tempdir().unwrap();
401 let first = FenceClaim::try_claim(root.path(), FenceOwnerKind::GuestLaunch, None)
402 .unwrap()
403 .unwrap();
404 assert!(
405 FenceClaim::try_claim(root.path(), FenceOwnerKind::WindowsRecovery, Some(1))
406 .unwrap()
407 .is_none()
408 );
409 drop(first);
410 assert!(
411 FenceClaim::try_claim(root.path(), FenceOwnerKind::WindowsRecovery, Some(1))
412 .unwrap()
413 .is_some()
414 );
415 }
416
417 #[test]
418 fn only_matching_recovery_generation_is_cleared() {
419 let root = tempfile::tempdir().unwrap();
420 let claim = FenceClaim::try_claim(root.path(), FenceOwnerKind::WindowsRecovery, Some(9))
421 .unwrap()
422 .unwrap();
423 claim.make_durable();
424 DrainRequest::new(9, Utc::now()).write(root.path()).unwrap();
425 clear_recovery(root.path(), 8).unwrap();
426 assert!(root.path().join(FENCE_DIRECTORY).exists());
427 assert_eq!(
428 DrainRequest::read(root.path()).unwrap().unwrap().generation,
429 9
430 );
431 clear_recovery(root.path(), 9).unwrap();
432 assert!(!root.path().join(FENCE_DIRECTORY).exists());
433 assert!(!root.path().join(REQUEST_FILE).exists());
434 }
435
436 #[test]
437 fn retiring_windows_recovery_handles_restarted_generations_but_keeps_guest_claims() {
438 let root = tempfile::tempdir().unwrap();
439 let claim = FenceClaim::try_claim(root.path(), FenceOwnerKind::WindowsRecovery, Some(7))
440 .unwrap()
441 .unwrap();
442 claim.make_durable();
443 DrainRequest::new(8, Utc::now()).write(root.path()).unwrap();
444 RecoveryStatus {
445 schema_version: SCHEMA_VERSION,
446 observed_at: Utc::now(),
447 phase: RecoveryPhase::RecoveryBlocked,
448 consecutive_probe_failures: 9,
449 reason: Some("stale".into()),
450 last_recovered_at: None,
451 }
452 .write(root.path())
453 .unwrap();
454
455 retire_windows_recovery(root.path()).unwrap();
456
457 assert!(!root.path().join(FENCE_DIRECTORY).exists());
458 assert!(!root.path().join(REQUEST_FILE).exists());
459 assert!(!root.path().join(RECOVERY_STATUS_FILE).exists());
460
461 let guest = FenceClaim::try_claim(root.path(), FenceOwnerKind::GuestLaunch, None)
462 .unwrap()
463 .unwrap();
464 guest.make_durable();
465 DrainRequest::new(10, Utc::now())
466 .write(root.path())
467 .unwrap();
468 retire_windows_recovery(root.path()).unwrap();
469 assert!(root.path().join(FENCE_DIRECTORY).exists());
470 assert!(!root.path().join(REQUEST_FILE).exists());
471 }
472
473 #[test]
474 fn heartbeat_and_status_documents_replace_atomically() {
475 let root = tempfile::tempdir().unwrap();
476 let first = GuestHeartbeat {
477 schema_version: SCHEMA_VERSION,
478 observed_at: Utc::now(),
479 acknowledged_generation: None,
480 local_active_attempts: Some(1),
481 managed_targets: Vec::new(),
482 unmanaged_runner_services: Some(0),
483 };
484 let mut second = first.clone();
485 second.local_active_attempts = Some(0);
486 first.write(root.path()).unwrap();
487 second.write(root.path()).unwrap();
488 assert_eq!(GuestHeartbeat::read(root.path()).unwrap(), Some(second));
489
490 let first = RecoveryStatus {
491 schema_version: SCHEMA_VERSION,
492 observed_at: Utc::now(),
493 phase: RecoveryPhase::Degraded,
494 consecutive_probe_failures: 1,
495 reason: None,
496 last_recovered_at: None,
497 };
498 let mut second = first.clone();
499 second.phase = RecoveryPhase::Healthy;
500 first.write(root.path()).unwrap();
501 second.write(root.path()).unwrap();
502 assert_eq!(RecoveryStatus::read(root.path()).unwrap(), Some(second));
503 }
504}