Skip to main content

awsim_core/
persistence.rs

1use std::path::PathBuf;
2
3use tracing::{error, info};
4
5/// Manages JSON snapshot-based persistence for service state.
6///
7/// When a `data_dir` is configured, each service can serialize its state to a JSON
8/// file under `{data_dir}/snapshots/{service}.json`.  On startup those snapshots are
9/// loaded and passed back to each service so it can rebuild its in-memory state.
10pub struct PersistenceManager {
11    data_dir: PathBuf,
12}
13
14impl PersistenceManager {
15    /// Create a new `PersistenceManager` rooted at `data_dir`.
16    pub fn new(data_dir: impl Into<PathBuf>) -> Self {
17        Self {
18            data_dir: data_dir.into(),
19        }
20    }
21
22    /// Save a service's state snapshot to `{data_dir}/snapshots/{service_name}.json`.
23    ///
24    /// Uses atomic write (write to temp file, then rename) to prevent corruption
25    /// if the process is killed mid-write.
26    pub fn save_snapshot(&self, service_name: &str, data: &[u8]) -> std::io::Result<()> {
27        let dir = self.data_dir.join("snapshots");
28        std::fs::create_dir_all(&dir)?;
29        let path = dir.join(format!("{service_name}.json"));
30        let tmp_path = dir.join(format!("{service_name}.json.tmp"));
31        // Write to temp file first
32        std::fs::write(&tmp_path, data)?;
33        // Atomic rename — either the old file remains or the new one replaces it
34        std::fs::rename(&tmp_path, &path)?;
35        info!(service = service_name, path = %path.display(), "Saved snapshot");
36        Ok(())
37    }
38
39    /// Load a service's state snapshot from disk.  Returns `None` if no snapshot exists.
40    pub fn load_snapshot(&self, service_name: &str) -> Option<Vec<u8>> {
41        let path = self
42            .data_dir
43            .join("snapshots")
44            .join(format!("{service_name}.json"));
45        match std::fs::read(&path) {
46            Ok(data) => {
47                info!(service = service_name, path = %path.display(), "Loaded snapshot");
48                Some(data)
49            }
50            Err(_) => None,
51        }
52    }
53
54    /// List the names of all saved snapshots (without the `.json` suffix).
55    pub fn list_snapshots(&self) -> Vec<String> {
56        let dir = self.data_dir.join("snapshots");
57        std::fs::read_dir(&dir)
58            .ok()
59            .map(|entries| {
60                entries
61                    .filter_map(|e| e.ok())
62                    .filter_map(|e| {
63                        let name = e.file_name().to_string_lossy().to_string();
64                        name.strip_suffix(".json").map(|s| s.to_string())
65                    })
66                    .collect()
67            })
68            .unwrap_or_default()
69    }
70
71    /// Save snapshots for all services that support it.
72    pub fn save_all(
73        &self,
74        services: &std::collections::HashMap<String, std::sync::Arc<dyn crate::ServiceHandler>>,
75    ) {
76        for (name, handler) in services {
77            if let Some(data) = handler.snapshot() {
78                if let Err(e) = self.save_snapshot(name, &data) {
79                    error!(service = %name, error = %e, "Failed to save snapshot");
80                }
81            }
82        }
83    }
84
85    /// Restore snapshots for all services that support it.
86    pub fn restore_all(
87        &self,
88        services: &std::collections::HashMap<String, std::sync::Arc<dyn crate::ServiceHandler>>,
89    ) {
90        for (name, handler) in services {
91            if let Some(data) = self.load_snapshot(name) {
92                if let Err(e) = handler.restore(&data) {
93                    tracing::warn!(service = %name, error = %e, "Failed to restore snapshot");
94                }
95            }
96        }
97    }
98}