Skip to main content

trellis_runner/engine/checkpoint/
in_memory.rs

1use std::sync::{Arc, Mutex};
2
3use crate::engine::checkpoint::{Checkpoint, CheckpointBackend, CheckpointView};
4
5#[derive(Clone, Default)]
6pub struct InMemoryCheckpointStore<SN, F> {
7    pub saved: Arc<Mutex<Vec<Checkpoint<SN, F>>>>,
8}
9
10impl<SN, F> InMemoryCheckpointStore<SN, F> {
11    pub fn new() -> Self {
12        Self {
13            saved: Arc::new(Mutex::new(Vec::new())),
14        }
15    }
16
17    pub fn saved_count(&self) -> usize {
18        self.saved.lock().unwrap().len()
19    }
20}
21
22impl<SN, F> CheckpointBackend<SN, F> for InMemoryCheckpointStore<SN, F>
23where
24    SN: Clone + Send + Sync,
25    F: Clone + Send + Sync,
26{
27    fn save(
28        &self,
29        checkpoint: CheckpointView<'_, SN, F>,
30    ) -> Result<(), crate::engine::checkpoint::CheckpointError> {
31        dbg!("saving");
32        self.saved.lock().unwrap().push(Checkpoint::new(checkpoint));
33        dbg!("saved");
34        Ok(())
35    }
36
37    fn load(
38        &self,
39    ) -> Result<Option<Checkpoint<SN, F>>, crate::engine::checkpoint::CheckpointError> {
40        Ok(self.saved.lock().unwrap().last().cloned())
41    }
42}