velesdb_core/agent/
snapshot.rs1#![allow(clippy::cast_possible_truncation)]
30
31use std::fs::File;
32use std::io::{self, Read, Write};
33use std::path::Path;
34
35use crate::storage::snapshot::crc32_hash;
36
37pub const SNAPSHOT_MAGIC: &[u8; 4] = b"VAMM";
39
40pub const SNAPSHOT_VERSION: u8 = 2;
47
48#[derive(Debug, Clone, Default)]
50pub struct MemoryState {
51 pub semantic: Vec<u8>,
53 pub episodic: Vec<u8>,
55 pub procedural: Vec<u8>,
57 pub ttl: Vec<u8>,
59}
60
61#[derive(Debug, Clone)]
63pub struct SnapshotMetadata {
64 pub version: u8,
66 pub total_size: usize,
68 pub checksum: u32,
70}
71
72#[derive(Debug)]
74#[non_exhaustive]
75pub enum SnapshotError {
76 Io(io::Error),
78 InvalidMagic,
80 UnsupportedVersion(u8),
82 ChecksumMismatch {
84 expected: u32,
86 actual: u32,
88 },
89 CorruptedData(String),
91}
92
93impl std::fmt::Display for SnapshotError {
94 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95 match self {
96 Self::Io(e) => write!(f, "IO error: {e}"),
97 Self::InvalidMagic => write!(f, "Invalid snapshot magic bytes"),
98 Self::UnsupportedVersion(v) => write!(f, "Unsupported snapshot version: {v}"),
99 Self::ChecksumMismatch { expected, actual } => {
100 write!(
101 f,
102 "Checksum mismatch: expected {expected:08x}, got {actual:08x}"
103 )
104 }
105 Self::CorruptedData(msg) => write!(f, "Corrupted data: {msg}"),
106 }
107 }
108}
109
110impl std::error::Error for SnapshotError {
111 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
112 match self {
113 Self::Io(e) => Some(e),
114 _ => None,
115 }
116 }
117}
118
119impl From<io::Error> for SnapshotError {
120 fn from(e: io::Error) -> Self {
121 Self::Io(e)
122 }
123}
124
125#[must_use]
135pub fn create_snapshot(state: &MemoryState) -> Vec<u8> {
136 let total_size = 4
137 + 1
138 + 8
139 + state.semantic.len()
140 + 8
141 + state.episodic.len()
142 + 8
143 + state.procedural.len()
144 + 8
145 + state.ttl.len()
146 + 4;
147 let mut buf = Vec::with_capacity(total_size);
148
149 buf.extend_from_slice(SNAPSHOT_MAGIC);
150 buf.push(SNAPSHOT_VERSION);
151
152 buf.extend_from_slice(&(state.semantic.len() as u64).to_le_bytes());
153 buf.extend_from_slice(&state.semantic);
154
155 buf.extend_from_slice(&(state.episodic.len() as u64).to_le_bytes());
156 buf.extend_from_slice(&state.episodic);
157
158 buf.extend_from_slice(&(state.procedural.len() as u64).to_le_bytes());
159 buf.extend_from_slice(&state.procedural);
160
161 buf.extend_from_slice(&(state.ttl.len() as u64).to_le_bytes());
162 buf.extend_from_slice(&state.ttl);
163
164 let crc = crc32_hash(&buf);
165 buf.extend_from_slice(&crc.to_le_bytes());
166
167 buf
168}
169
170pub fn load_snapshot(data: &[u8]) -> Result<MemoryState, SnapshotError> {
180 validate_snapshot_header(data)?;
181
182 let mut offset = 5; let payload_end = data.len() - 4; let semantic = read_section(data, &mut offset, payload_end, "Semantic")?;
186 let episodic = read_section(data, &mut offset, payload_end, "Episodic")?;
187 let procedural = read_section(data, &mut offset, payload_end, "Procedural")?;
188 let ttl = read_section(data, &mut offset, payload_end, "TTL")?;
189
190 Ok(MemoryState {
191 semantic,
192 episodic,
193 procedural,
194 ttl,
195 })
196}
197
198fn validate_snapshot_header(data: &[u8]) -> Result<(), SnapshotError> {
200 const MIN_SIZE: usize = 4 + 1 + 8 + 8 + 8 + 8 + 4;
201
202 if data.len() < MIN_SIZE {
203 return Err(SnapshotError::CorruptedData(
204 "Snapshot too small".to_string(),
205 ));
206 }
207 if &data[0..4] != SNAPSHOT_MAGIC {
208 return Err(SnapshotError::InvalidMagic);
209 }
210 let version = data[4];
211 if version != SNAPSHOT_VERSION {
212 return Err(SnapshotError::UnsupportedVersion(version));
213 }
214
215 let stored_crc = u32::from_le_bytes(
216 data[data.len() - 4..]
217 .try_into()
218 .map_err(|_| SnapshotError::CorruptedData("Invalid CRC bytes".to_string()))?,
219 );
220 let computed_crc = crc32_hash(&data[..data.len() - 4]);
221 if stored_crc != computed_crc {
222 return Err(SnapshotError::ChecksumMismatch {
223 expected: stored_crc,
224 actual: computed_crc,
225 });
226 }
227 Ok(())
228}
229
230fn read_section(
232 data: &[u8],
233 offset: &mut usize,
234 payload_end: usize,
235 label: &str,
236) -> Result<Vec<u8>, SnapshotError> {
237 let section_len = read_u64(&data[*offset..])? as usize;
238 *offset += 8;
239 let end = offset
243 .checked_add(section_len)
244 .filter(|end| *end <= payload_end)
245 .ok_or_else(|| SnapshotError::CorruptedData(format!("{label} data truncated")))?;
246 let section = data[*offset..end].to_vec();
247 *offset = end;
248 Ok(section)
249}
250
251pub fn save_snapshot_to_file<P: AsRef<Path>>(
259 path: P,
260 state: &MemoryState,
261) -> Result<(), SnapshotError> {
262 let path = path.as_ref();
263 let snapshot_data = create_snapshot(state);
264
265 let temp_path = path.with_extension("tmp");
266 let mut file = File::create(&temp_path)?;
267 file.write_all(&snapshot_data)?;
268 file.sync_all()?;
269 drop(file);
270
271 std::fs::rename(&temp_path, path)?;
272
273 Ok(())
274}
275
276pub fn load_snapshot_from_file<P: AsRef<Path>>(path: P) -> Result<MemoryState, SnapshotError> {
282 let mut file = File::open(path)?;
283 let mut data = Vec::new();
284 file.read_to_end(&mut data)?;
285 load_snapshot(&data)
286}
287
288fn read_u64(data: &[u8]) -> Result<u64, SnapshotError> {
290 if data.len() < 8 {
291 return Err(SnapshotError::CorruptedData(
292 "Not enough bytes for u64".to_string(),
293 ));
294 }
295 Ok(u64::from_le_bytes(data[0..8].try_into().map_err(|_| {
296 SnapshotError::CorruptedData("Invalid u64 bytes".to_string())
297 })?))
298}
299
300pub struct SnapshotManager {
302 base_path: std::path::PathBuf,
304 max_snapshots: usize,
306}
307
308impl SnapshotManager {
309 pub fn new<P: AsRef<Path>>(base_path: P, max_snapshots: usize) -> Self {
316 Self {
317 base_path: base_path.as_ref().to_path_buf(),
318 max_snapshots,
319 }
320 }
321
322 pub fn create_versioned_snapshot(&self, state: &MemoryState) -> Result<u64, SnapshotError> {
332 std::fs::create_dir_all(&self.base_path)?;
333
334 let version = self.next_version()?;
335 let filename = format!("snapshot_{version:08}.vamm");
336 let path = self.base_path.join(filename);
337
338 save_snapshot_to_file(&path, state)?;
339 self.cleanup_old_snapshots()?;
340
341 Ok(version)
342 }
343
344 pub fn load_latest(&self) -> Result<(u64, MemoryState), SnapshotError> {
350 let version = self
351 .latest_version()?
352 .ok_or_else(|| SnapshotError::CorruptedData("No snapshots found".to_string()))?;
353 let state = self.load_version(version)?;
354 Ok((version, state))
355 }
356
357 pub fn load_version(&self, version: u64) -> Result<MemoryState, SnapshotError> {
363 let filename = format!("snapshot_{version:08}.vamm");
364 let path = self.base_path.join(filename);
365 load_snapshot_from_file(&path)
366 }
367
368 pub fn list_versions(&self) -> Result<Vec<u64>, SnapshotError> {
374 if !self.base_path.exists() {
375 return Ok(Vec::new());
376 }
377
378 let mut versions: Vec<u64> = std::fs::read_dir(&self.base_path)?
379 .filter_map(Result::ok)
380 .filter_map(|e| parse_snapshot_version(&e.file_name().to_string_lossy()))
381 .collect();
382
383 versions.sort_unstable();
384 Ok(versions)
385 }
386
387 fn latest_version(&self) -> Result<Option<u64>, SnapshotError> {
389 Ok(self.list_versions()?.into_iter().max())
390 }
391
392 fn next_version(&self) -> Result<u64, SnapshotError> {
394 Ok(self.latest_version()?.map_or(1, |v| v + 1))
395 }
396
397 fn cleanup_old_snapshots(&self) -> Result<(), SnapshotError> {
399 let versions = self.list_versions()?;
400 if versions.len() <= self.max_snapshots {
401 return Ok(());
402 }
403
404 let to_remove = versions.len() - self.max_snapshots;
405 for version in versions.into_iter().take(to_remove) {
406 let filename = format!("snapshot_{version:08}.vamm");
407 let path = self.base_path.join(filename);
408 let _ = std::fs::remove_file(path);
409 }
410
411 Ok(())
412 }
413}
414
415fn parse_snapshot_version(filename: &str) -> Option<u64> {
417 filename
418 .strip_prefix("snapshot_")
419 .and_then(|s| s.strip_suffix(".vamm"))
420 .and_then(|s| s.parse::<u64>().ok())
421}