fakecloud_dsql/
persistence.rs

1//! Snapshot save/load helpers for the DSQL control-plane state.
2
3use std::sync::Arc;
4
5use tokio::sync::Mutex as AsyncMutex;
6
7use fakecloud_persistence::SnapshotStore;
8
9use crate::state::{DsqlSnapshot, SharedDsqlState, DSQL_SNAPSHOT_SCHEMA_VERSION};
10
11#[derive(Debug, PartialEq, Eq)]
12pub enum LoadOutcome {
13    /// No snapshot file on disk; start with fresh state.
14    Empty,
15    /// Snapshot loaded successfully; returns the restored account count.
16    Loaded(usize),
17}
18
19#[derive(Debug, thiserror::Error)]
20pub enum LoadError {
21    #[error("failed to read dsql persistence snapshot: {0}")]
22    Io(String),
23    #[error("failed to parse dsql persistence snapshot: {0}")]
24    Parse(String),
25    #[error("dsql persistence schema too new: on-disk={on_disk}, max supported={supported}")]
26    SchemaTooNew { on_disk: u32, supported: u32 },
27}
28
29/// Load a snapshot into `state`. `Empty` when nothing is saved, `Loaded(n)`
30/// after a successful restore, or a descriptive error for a fatal startup log.
31pub fn load_into(
32    store: &dyn SnapshotStore,
33    state: &SharedDsqlState,
34) -> Result<LoadOutcome, LoadError> {
35    let Some(bytes) = store.load().map_err(|e| LoadError::Io(e.to_string()))? else {
36        return Ok(LoadOutcome::Empty);
37    };
38    let snapshot: DsqlSnapshot =
39        serde_json::from_slice(&bytes).map_err(|e| LoadError::Parse(e.to_string()))?;
40    if snapshot.schema_version > DSQL_SNAPSHOT_SCHEMA_VERSION {
41        return Err(LoadError::SchemaTooNew {
42            on_disk: snapshot.schema_version,
43            supported: DSQL_SNAPSHOT_SCHEMA_VERSION,
44        });
45    }
46    let accounts = snapshot.accounts.account_count();
47    *state.write() = snapshot.accounts;
48    Ok(LoadOutcome::Loaded(accounts))
49}
50
51/// Serialize and persist the current DSQL state. No-op when there is no store
52/// (memory mode). Serialized off the async runtime to avoid blocking.
53pub async fn save_snapshot(
54    state: &SharedDsqlState,
55    store: Option<Arc<dyn SnapshotStore>>,
56    lock: &AsyncMutex<()>,
57) {
58    let Some(store) = store else {
59        return;
60    };
61    let _guard = lock.lock().await;
62    let snapshot = DsqlSnapshot {
63        schema_version: DSQL_SNAPSHOT_SCHEMA_VERSION,
64        accounts: state.read().clone(),
65    };
66    let join = tokio::task::spawn_blocking(move || -> std::io::Result<()> {
67        let bytes = serde_json::to_vec(&snapshot)
68            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?;
69        store.save(&bytes)
70    })
71    .await;
72    match join {
73        Ok(Ok(())) => {}
74        Ok(Err(err)) => tracing::error!(%err, "failed to write dsql snapshot"),
75        Err(err) => tracing::error!(%err, "dsql snapshot task panicked"),
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82    use crate::state::{DsqlSnapshot, DsqlState};
83    use fakecloud_core::multi_account::MultiAccountState;
84    use parking_lot::RwLock;
85    use std::sync::Mutex;
86
87    fn make_state() -> SharedDsqlState {
88        Arc::new(RwLock::new(MultiAccountState::new(
89            "000000000000",
90            "us-east-1",
91            "",
92        )))
93    }
94
95    struct MemStore {
96        data: Mutex<Option<Vec<u8>>>,
97    }
98    impl MemStore {
99        fn new(data: Option<Vec<u8>>) -> Self {
100            Self {
101                data: Mutex::new(data),
102            }
103        }
104    }
105    impl SnapshotStore for MemStore {
106        fn load(&self) -> std::io::Result<Option<Vec<u8>>> {
107            Ok(self.data.lock().unwrap().clone())
108        }
109        fn save(&self, bytes: &[u8]) -> std::io::Result<()> {
110            *self.data.lock().unwrap() = Some(bytes.to_vec());
111            Ok(())
112        }
113    }
114
115    #[test]
116    fn load_into_empty_returns_empty() {
117        let state = make_state();
118        let store = MemStore::new(None);
119        assert_eq!(load_into(&store, &state).unwrap(), LoadOutcome::Empty);
120    }
121
122    #[test]
123    fn load_into_valid_snapshot_restores_accounts() {
124        let mut accounts: MultiAccountState<DsqlState> =
125            MultiAccountState::new("000000000000", "us-east-1", "");
126        accounts.get_or_create("111122223333");
127        let snap = DsqlSnapshot {
128            schema_version: DSQL_SNAPSHOT_SCHEMA_VERSION,
129            accounts,
130        };
131        let bytes = serde_json::to_vec(&snap).unwrap();
132        let state = make_state();
133        let store = MemStore::new(Some(bytes));
134        assert_eq!(load_into(&store, &state).unwrap(), LoadOutcome::Loaded(2));
135    }
136
137    #[test]
138    fn load_into_rejects_future_schema() {
139        let mut accounts: MultiAccountState<DsqlState> =
140            MultiAccountState::new("000000000000", "us-east-1", "");
141        accounts.get_or_create("111122223333");
142        let bytes = serde_json::to_vec(&serde_json::json!({
143            "schema_version": DSQL_SNAPSHOT_SCHEMA_VERSION + 1,
144            "accounts": accounts,
145        }))
146        .unwrap();
147        let state = make_state();
148        let store = MemStore::new(Some(bytes));
149        assert!(matches!(
150            load_into(&store, &state),
151            Err(LoadError::SchemaTooNew { .. })
152        ));
153    }
154}
fakecloud_dsql/persistence.rs

fakecloud_dsql/
persistence.rs
