Skip to main content

fakecloud_translate/
persistence.rs

1//! Snapshot save/load for Amazon Translate 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::{SharedTranslateState, TranslateSnapshot, TRANSLATE_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 translate persistence snapshot: {0}")]
21    Io(String),
22    #[error("failed to parse translate persistence snapshot: {0}")]
23    Parse(String),
24    #[error("translate 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: &SharedTranslateState,
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: TranslateSnapshot =
36        serde_json::from_slice(&bytes).map_err(|e| LoadError::Parse(e.to_string()))?;
37    if snapshot.schema_version > TRANSLATE_SNAPSHOT_SCHEMA_VERSION {
38        return Err(LoadError::SchemaTooNew {
39            on_disk: snapshot.schema_version,
40            supported: TRANSLATE_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: &SharedTranslateState,
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 = TranslateSnapshot {
64        schema_version: TRANSLATE_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 translate snapshot"),
76        Err(err) => tracing::error!(%err, "translate snapshot task panicked"),
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83    use crate::state::TranslateData;
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() -> SharedTranslateState {
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<TranslateData> =
119            MultiAccountState::new("000000000000", "us-east-1", "");
120        let data = accounts.get_or_create("111122223333");
121        data.text_translation_jobs.insert(
122            "job-1".into(),
123            json!({ "JobId": "job-1", "JobStatus": "SUBMITTED" }),
124        );
125        let snap = TranslateSnapshot {
126            schema_version: TRANSLATE_SNAPSHOT_SCHEMA_VERSION,
127            accounts,
128        };
129        let store = MemStore(Mutex::new(Some(serde_json::to_vec(&snap).unwrap())));
130        let restored = state();
131        assert_eq!(
132            load_into(&store, &restored).unwrap(),
133            LoadOutcome::Loaded(2)
134        );
135        let guard = restored.read();
136        let job = &guard.get("111122223333").unwrap().text_translation_jobs["job-1"];
137        assert_eq!(job["JobStatus"], "COMPLETED");
138    }
139
140    #[test]
141    fn rejects_future_schema() {
142        let accounts: MultiAccountState<TranslateData> =
143            MultiAccountState::new("000000000000", "us-east-1", "");
144        let bytes = serde_json::to_vec(&serde_json::json!({
145            "schema_version": TRANSLATE_SNAPSHOT_SCHEMA_VERSION + 1,
146            "accounts": accounts,
147        }))
148        .unwrap();
149        let store = MemStore(Mutex::new(Some(bytes)));
150        assert!(matches!(
151            load_into(&store, &state()),
152            Err(LoadError::SchemaTooNew { .. })
153        ));
154    }
155}