tx2_pack/
checkpoint.rs

1use crate::error::Result;
2use crate::format::PackedSnapshot;
3use crate::metadata::SnapshotMetadata;
4use crate::storage::{SnapshotWriter, SnapshotReader, SnapshotStore};
5use std::path::Path;
6use ahash::AHashMap;
7
8#[derive(Debug, Clone)]
9pub struct Checkpoint {
10    pub id: String,
11    pub snapshot: PackedSnapshot,
12    pub metadata: SnapshotMetadata,
13    pub parent_id: Option<String>,
14}
15
16impl Checkpoint {
17    pub fn new(id: String, snapshot: PackedSnapshot) -> Self {
18        let metadata = SnapshotMetadata::new(id.clone());
19
20        Self {
21            id: id.clone(),
22            snapshot,
23            metadata,
24            parent_id: None,
25        }
26    }
27
28    pub fn with_parent(mut self, parent_id: String) -> Self {
29        self.parent_id = Some(parent_id);
30        self
31    }
32
33    pub fn with_metadata(mut self, metadata: SnapshotMetadata) -> Self {
34        self.metadata = metadata;
35        self
36    }
37}
38
39pub struct CheckpointManager {
40    store: SnapshotStore,
41    writer: SnapshotWriter,
42    reader: SnapshotReader,
43    checkpoints: AHashMap<String, Checkpoint>,
44    checkpoint_chain: Vec<String>,
45}
46
47impl CheckpointManager {
48    pub fn new<P: AsRef<Path>>(root_dir: P) -> Result<Self> {
49        let store = SnapshotStore::new(root_dir)?;
50        let writer = SnapshotWriter::new();
51        let reader = SnapshotReader::new();
52
53        Ok(Self {
54            store,
55            writer,
56            reader,
57            checkpoints: AHashMap::new(),
58            checkpoint_chain: Vec::new(),
59        })
60    }
61
62    pub fn with_writer(mut self, writer: SnapshotWriter) -> Self {
63        self.writer = writer;
64        self
65    }
66
67    pub fn with_reader(mut self, reader: SnapshotReader) -> Self {
68        self.reader = reader;
69        self
70    }
71
72    pub fn create_checkpoint(
73        &mut self,
74        id: String,
75        snapshot: PackedSnapshot,
76    ) -> Result<()> {
77        let parent_id = self.checkpoint_chain.last().cloned();
78
79        let mut checkpoint = Checkpoint::new(id.clone(), snapshot);
80        if let Some(parent) = parent_id {
81            checkpoint = checkpoint.with_parent(parent);
82        }
83
84        self.store.save(&checkpoint.snapshot, &checkpoint.metadata, &self.writer)?;
85
86        self.checkpoint_chain.push(id.clone());
87        self.checkpoints.insert(id, checkpoint);
88
89        Ok(())
90    }
91
92    pub fn load_checkpoint(&mut self, id: &str) -> Result<Checkpoint> {
93        if let Some(checkpoint) = self.checkpoints.get(id) {
94            return Ok(checkpoint.clone());
95        }
96
97        let (snapshot, metadata) = self.store.load(id, &self.reader)?;
98
99        let checkpoint = Checkpoint {
100            id: id.to_string(),
101            snapshot,
102            metadata,
103            parent_id: None,
104        };
105
106        self.checkpoints.insert(id.to_string(), checkpoint.clone());
107
108        Ok(checkpoint)
109    }
110
111    pub fn delete_checkpoint(&mut self, id: &str) -> Result<()> {
112        self.store.delete(id)?;
113        self.checkpoints.remove(id);
114        self.checkpoint_chain.retain(|cid| cid != id);
115        Ok(())
116    }
117
118    pub fn list_checkpoints(&self) -> Result<Vec<String>> {
119        self.store.list()
120    }
121
122    pub fn get_checkpoint_chain(&self) -> &[String] {
123        &self.checkpoint_chain
124    }
125
126    pub fn get_latest_checkpoint(&self) -> Option<&str> {
127        self.checkpoint_chain.last().map(|s| s.as_str())
128    }
129
130    pub fn prune_old_checkpoints(&mut self, keep_count: usize) -> Result<()> {
131        let chain_len = self.checkpoint_chain.len();
132
133        if chain_len <= keep_count {
134            return Ok(());
135        }
136
137        let to_remove = chain_len - keep_count;
138
139        for _ in 0..to_remove {
140            if let Some(id) = self.checkpoint_chain.first().cloned() {
141                self.delete_checkpoint(&id)?;
142            }
143        }
144
145        Ok(())
146    }
147
148    pub fn clear_all_checkpoints(&mut self) -> Result<()> {
149        for id in self.checkpoint_chain.clone() {
150            self.delete_checkpoint(&id)?;
151        }
152
153        Ok(())
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160    use tempfile::TempDir;
161
162    #[test]
163    fn test_checkpoint_manager() {
164        let temp_dir = TempDir::new().unwrap();
165        let mut manager = CheckpointManager::new(temp_dir.path()).unwrap();
166
167        let snapshot1 = PackedSnapshot::new();
168        manager.create_checkpoint("cp1".to_string(), snapshot1).unwrap();
169
170        let snapshot2 = PackedSnapshot::new();
171        manager.create_checkpoint("cp2".to_string(), snapshot2).unwrap();
172
173        assert_eq!(manager.get_checkpoint_chain().len(), 2);
174        assert_eq!(manager.get_latest_checkpoint(), Some("cp2"));
175
176        let loaded = manager.load_checkpoint("cp1").unwrap();
177        assert_eq!(loaded.id, "cp1");
178
179        manager.prune_old_checkpoints(1).unwrap();
180        assert_eq!(manager.get_checkpoint_chain().len(), 1);
181        assert_eq!(manager.get_latest_checkpoint(), Some("cp2"));
182    }
183
184    #[test]
185    fn test_checkpoint_clear() {
186        let temp_dir = TempDir::new().unwrap();
187        let mut manager = CheckpointManager::new(temp_dir.path()).unwrap();
188
189        for i in 0..5 {
190            let snapshot = PackedSnapshot::new();
191            manager.create_checkpoint(format!("cp{}", i), snapshot).unwrap();
192        }
193
194        assert_eq!(manager.get_checkpoint_chain().len(), 5);
195
196        manager.clear_all_checkpoints().unwrap();
197        assert_eq!(manager.get_checkpoint_chain().len(), 0);
198    }
199}
tx2_pack/checkpoint.rs

tx2_pack/
checkpoint.rs
