fakecloud_comprehend/
persistence.rs

1//! Snapshot save/load for Amazon Comprehend state, with lifecycle
2//! reconciliation on restart.
3
4use std::sync::Arc;
5
6use tokio::sync::Mutex as AsyncMutex;
7
8use fakecloud_persistence::SnapshotStore;
9
10use crate::state::{ComprehendSnapshot, SharedComprehendState, COMPREHEND_SNAPSHOT_SCHEMA_VERSION};
11
12#[derive(Debug, PartialEq, Eq)]
13pub enum LoadOutcome {
14    Empty,
15    Loaded(usize),
16}
17
18#[derive(Debug, thiserror::Error)]
19pub enum LoadError {
20    #[error("failed to read comprehend persistence snapshot: {0}")]
21    Io(String),
22    #[error("failed to parse comprehend persistence snapshot: {0}")]
23    Parse(String),
24    #[error("comprehend persistence schema too new: on-disk={on_disk}, max supported={supported}")]
25    SchemaTooNew { on_disk: u32, supported: u32 },
26}
27
28pub fn load_into(
29    store: &dyn SnapshotStore,
30    state: &SharedComprehendState,
31) -> Result<LoadOutcome, LoadError> {
32    let Some(bytes) = store.load().map_err(|e| LoadError::Io(e.to_string()))? else {
33        return Ok(LoadOutcome::Empty);
34    };
35    let mut snapshot: ComprehendSnapshot =
36        serde_json::from_slice(&bytes).map_err(|e| LoadError::Parse(e.to_string()))?;
37    if snapshot.schema_version > COMPREHEND_SNAPSHOT_SCHEMA_VERSION {
38        return Err(LoadError::SchemaTooNew {
39            on_disk: snapshot.schema_version,
40            supported: COMPREHEND_SNAPSHOT_SCHEMA_VERSION,
41        });
42    }
43    // Reconcile any lifecycle transition that was in flight when the process
44    // stopped: it advances toward its settled state (there is no timer to
45    // resume).
46    for (_account_id, account) in snapshot.accounts.iter_mut() {
47        account.reconcile();
48    }
49    let accounts = snapshot.accounts.account_count();
50    *state.write() = snapshot.accounts;
51    Ok(LoadOutcome::Loaded(accounts))
52}
53
54pub async fn save_snapshot(
55    state: &SharedComprehendState,
56    store: Option<Arc<dyn SnapshotStore>>,
57    lock: &AsyncMutex<()>,
58) {
59    let Some(store) = store else {
60        return;
61    };
62    let _guard = lock.lock().await;
63    let snapshot = ComprehendSnapshot {
64        schema_version: COMPREHEND_SNAPSHOT_SCHEMA_VERSION,
65        accounts: state.read().clone(),
66    };
67    let join = tokio::task::spawn_blocking(move || -> std::io::Result<()> {
68        let bytes = serde_json::to_vec(&snapshot)
69            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()))?;
70        store.save(&bytes)
71    })
72    .await;
73    match join {
74        Ok(Ok(())) => {}
75        Ok(Err(err)) => tracing::error!(%err, "failed to write comprehend snapshot"),
76        Err(err) => tracing::error!(%err, "comprehend snapshot task panicked"),
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use crate::state::ComprehendData;
84    use fakecloud_core::multi_account::MultiAccountState;
85    use parking_lot::RwLock;
86    use serde_json::json;
87    use std::sync::Mutex;
88
89    struct MemStore(Mutex<Option<Vec<u8>>>);
90    impl SnapshotStore for MemStore {
91        fn load(&self) -> std::io::Result<Option<Vec<u8>>> {
92            Ok(self.0.lock().unwrap().clone())
93        }
94        fn save(&self, bytes: &[u8]) -> std::io::Result<()> {
95            *self.0.lock().unwrap() = Some(bytes.to_vec());
96            Ok(())
97        }
98    }
99
100    fn state() -> SharedComprehendState {
101        Arc::new(RwLock::new(MultiAccountState::new(
102            "000000000000",
103            "us-east-1",
104            "",
105        )))
106    }
107
108    #[test]
109    fn empty_store_is_empty() {
110        assert_eq!(
111            load_into(&MemStore(Mutex::new(None)), &state()).unwrap(),
112            LoadOutcome::Empty
113        );
114    }
115
116    #[test]
117    fn round_trip_restores_and_reconciles_lifecycle() {
118        let mut accounts: MultiAccountState<ComprehendData> =
119            MultiAccountState::new("000000000000", "us-east-1", "");
120        let data = accounts.get_or_create("111122223333");
121        data.jobs.insert(
122            "j1".into(),
123            json!({ "JobId": "j1", "JobStatus": "SUBMITTED" }),
124        );
125        data.job_families.insert("j1".into(), "sentiment".into());
126        let snap = ComprehendSnapshot {
127            schema_version: COMPREHEND_SNAPSHOT_SCHEMA_VERSION,
128            accounts,
129        };
130        let store = MemStore(Mutex::new(Some(serde_json::to_vec(&snap).unwrap())));
131        let restored = state();
132        assert_eq!(
133            load_into(&store, &restored).unwrap(),
134            LoadOutcome::Loaded(2)
135        );
136        let guard = restored.read();
137        let job = &guard.get("111122223333").unwrap().jobs["j1"];
138        assert_eq!(job["JobStatus"], "COMPLETED");
139    }
140
141    #[test]
142    fn rejects_future_schema() {
143        let accounts: MultiAccountState<ComprehendData> =
144            MultiAccountState::new("000000000000", "us-east-1", "");
145        let bytes = serde_json::to_vec(&serde_json::json!({
146            "schema_version": COMPREHEND_SNAPSHOT_SCHEMA_VERSION + 1,
147            "accounts": accounts,
148        }))
149        .unwrap();
150        let store = MemStore(Mutex::new(Some(bytes)));
151        assert!(matches!(
152            load_into(&store, &state()),
153            Err(LoadError::SchemaTooNew { .. })
154        ));
155    }
156}
fakecloud_comprehend/persistence.rs

fakecloud_comprehend/
persistence.rs
