scirs2_io/streaming/
checkpoint.rs

1//! Checkpointing and restart support for long-running streaming jobs.
2//!
3//! Provides facilities for periodically snapshotting operator state to durable
4//! storage so that a streaming job can be resumed after a failure without
5//! reprocessing the entire input from the beginning.
6//!
7//! ## Design
8//!
9//! - [`CheckpointManager`] handles saving/loading raw byte blobs tagged with a
10//!   `state_id` string.  Files are named `{state_id}_{timestamp_ms}.ckpt` and
11//!   stored under `CheckpointConfig::storage_path`.
12//! - [`CheckpointableState`] is the serialisation contract; any type that can
13//!   write itself to `Vec<u8>` and reconstruct itself from `&[u8]` qualifies.
14//! - [`CheckpointBarrier`] is a lightweight token injected into the stream to
15//!   coordinate when all operators should snapshot their state.
16//!
17//! ## Simple byte serialisation
18//!
19//! Rather than pulling in an external dependency, primitive helpers convert
20//! `Vec<f64>` and `HashMap<String, f64>` to/from little-endian bytes.
21
22use std::collections::HashMap;
23use std::fs;
24use std::path::{Path, PathBuf};
25use std::time::{SystemTime, UNIX_EPOCH};
26
27use crate::error::{IoError, Result};
28
29// ---------------------------------------------------------------------------
30// CheckpointableState trait
31// ---------------------------------------------------------------------------
32
33/// Any streaming operator state that can be serialised and restored.
34pub trait CheckpointableState: Sized {
35    /// Serialise the state to a byte vector.
36    fn serialize(&self) -> Vec<u8>;
37
38    /// Deserialise from a byte slice.
39    ///
40    /// # Errors
41    ///
42    /// Returns an error if the bytes are malformed.
43    fn deserialize(data: &[u8]) -> Result<Self>;
44}
45
46// ---------------------------------------------------------------------------
47// Built-in CheckpointableState implementations
48// ---------------------------------------------------------------------------
49
50impl CheckpointableState for Vec<f64> {
51    fn serialize(&self) -> Vec<u8> {
52        let mut out = Vec::with_capacity(8 + self.len() * 8);
53        out.extend_from_slice(&(self.len() as u64).to_le_bytes());
54        for &v in self {
55            out.extend_from_slice(&v.to_le_bytes());
56        }
57        out
58    }
59
60    fn deserialize(data: &[u8]) -> Result<Self> {
61        if data.len() < 8 {
62            return Err(IoError::DeserializationError(
63                "Vec<f64> checkpoint too short".to_string(),
64            ));
65        }
66        let len = u64::from_le_bytes(
67            data[..8]
68                .try_into()
69                .map_err(|_| IoError::DeserializationError("bad length bytes".to_string()))?,
70        ) as usize;
71        let expected = 8 + len * 8;
72        if data.len() < expected {
73            return Err(IoError::DeserializationError(format!(
74                "Vec<f64> checkpoint: expected {expected} bytes, got {}",
75                data.len()
76            )));
77        }
78        let mut vec = Vec::with_capacity(len);
79        for i in 0..len {
80            let start = 8 + i * 8;
81            let bytes: [u8; 8] = data[start..start + 8]
82                .try_into()
83                .map_err(|_| IoError::DeserializationError("bad f64 bytes".to_string()))?;
84            vec.push(f64::from_le_bytes(bytes));
85        }
86        Ok(vec)
87    }
88}
89
90impl CheckpointableState for HashMap<String, f64> {
91    fn serialize(&self) -> Vec<u8> {
92        // Format: n_entries (8 bytes) | for each: key_len (8) | key (key_len) | value (8)
93        let mut out = Vec::new();
94        out.extend_from_slice(&(self.len() as u64).to_le_bytes());
95        // Sorted for determinism.
96        let mut entries: Vec<(&String, &f64)> = self.iter().collect();
97        entries.sort_by_key(|(k, _)| k.as_str());
98        for (k, &v) in entries {
99            let kb = k.as_bytes();
100            out.extend_from_slice(&(kb.len() as u64).to_le_bytes());
101            out.extend_from_slice(kb);
102            out.extend_from_slice(&v.to_le_bytes());
103        }
104        out
105    }
106
107    fn deserialize(data: &[u8]) -> Result<Self> {
108        if data.len() < 8 {
109            return Err(IoError::DeserializationError(
110                "HashMap checkpoint too short".to_string(),
111            ));
112        }
113        let n = u64::from_le_bytes(
114            data[..8]
115                .try_into()
116                .map_err(|_| IoError::DeserializationError("bad n bytes".to_string()))?,
117        ) as usize;
118        let mut cursor = 8usize;
119        let mut map = HashMap::with_capacity(n);
120        for _ in 0..n {
121            if cursor + 8 > data.len() {
122                return Err(IoError::DeserializationError(
123                    "HashMap checkpoint truncated (key len)".to_string(),
124                ));
125            }
126            let key_len = u64::from_le_bytes(
127                data[cursor..cursor + 8]
128                    .try_into()
129                    .map_err(|_| IoError::DeserializationError("bad key len".to_string()))?,
130            ) as usize;
131            cursor += 8;
132            if cursor + key_len + 8 > data.len() {
133                return Err(IoError::DeserializationError(
134                    "HashMap checkpoint truncated (key/value)".to_string(),
135                ));
136            }
137            let key = String::from_utf8(data[cursor..cursor + key_len].to_vec()).map_err(|e| {
138                IoError::DeserializationError(format!("non-UTF8 key in HashMap checkpoint: {e}"))
139            })?;
140            cursor += key_len;
141            let value = f64::from_le_bytes(
142                data[cursor..cursor + 8]
143                    .try_into()
144                    .map_err(|_| IoError::DeserializationError("bad f64 value".to_string()))?,
145            );
146            cursor += 8;
147            map.insert(key, value);
148        }
149        Ok(map)
150    }
151}
152
153// ---------------------------------------------------------------------------
154// CheckpointConfig
155// ---------------------------------------------------------------------------
156
157/// Configuration for the checkpoint manager.
158#[derive(Debug, Clone)]
159pub struct CheckpointConfig {
160    /// Minimum interval between automatic checkpoints in milliseconds.
161    pub interval_ms: u64,
162    /// Directory where checkpoint files are stored.
163    pub storage_path: String,
164    /// Number of most-recent checkpoints to retain (older ones are deleted).
165    pub max_checkpoints: usize,
166}
167
168impl Default for CheckpointConfig {
169    fn default() -> Self {
170        Self {
171            interval_ms: 60_000,
172            storage_path: {
173                let mut p = std::env::temp_dir();
174                p.push("scirs2_checkpoint");
175                p.to_string_lossy().into_owned()
176            },
177            max_checkpoints: 3,
178        }
179    }
180}
181
182// ---------------------------------------------------------------------------
183// CheckpointManager
184// ---------------------------------------------------------------------------
185
186/// Manages serialised state blobs on the local filesystem.
187#[derive(Debug, Clone)]
188pub struct CheckpointManager {
189    config: CheckpointConfig,
190}
191
192impl CheckpointManager {
193    /// Create a new manager with the given configuration.
194    pub fn new(config: CheckpointConfig) -> Self {
195        Self { config }
196    }
197
198    /// Create a new manager with the default configuration.
199    pub fn with_defaults() -> Self {
200        Self::new(CheckpointConfig::default())
201    }
202
203    /// Save `data` as a checkpoint for `state_id`.
204    ///
205    /// The file is named `{state_id}_{timestamp_ms}.ckpt` and written to the
206    /// configured `storage_path`.  Returns the absolute path of the file
207    /// written.
208    ///
209    /// After saving, old checkpoints beyond `max_checkpoints` are removed.
210    pub fn save(&self, state_id: &str, data: &[u8]) -> Result<String> {
211        let dir = Path::new(&self.config.storage_path);
212        fs::create_dir_all(dir).map_err(|e| {
213            IoError::FileError(format!(
214                "Cannot create checkpoint directory {}: {e}",
215                dir.display()
216            ))
217        })?;
218
219        let ts_ms = current_time_ms();
220        let filename = format!("{state_id}_{ts_ms}.ckpt");
221        let path = dir.join(&filename);
222
223        fs::write(&path, data).map_err(|e| {
224            IoError::FileError(format!(
225                "Failed to write checkpoint {}: {e}",
226                path.display()
227            ))
228        })?;
229
230        let abs = path
231            .canonicalize()
232            .map(|p| p.to_string_lossy().into_owned())
233            .unwrap_or_else(|_| path.to_string_lossy().into_owned());
234
235        self.cleanup_old(state_id);
236        Ok(abs)
237    }
238
239    /// Load the most recently saved checkpoint for `state_id`.
240    ///
241    /// # Errors
242    ///
243    /// Returns an error if no checkpoint is found or if reading fails.
244    pub fn load_latest(&self, state_id: &str) -> Result<Vec<u8>> {
245        let files = self.list_checkpoints(state_id);
246        let latest = files.last().ok_or_else(|| {
247            IoError::NotFound(format!("No checkpoint found for state_id '{state_id}'"))
248        })?;
249        fs::read(latest)
250            .map_err(|e| IoError::FileError(format!("Failed to read checkpoint {latest}: {e}")))
251    }
252
253    /// Delete all but the newest `max_checkpoints` checkpoint files for `state_id`.
254    pub fn cleanup_old(&self, state_id: &str) {
255        let files = self.list_checkpoints(state_id);
256        if files.len() <= self.config.max_checkpoints {
257            return;
258        }
259        let to_delete = files.len() - self.config.max_checkpoints;
260        for path in files.iter().take(to_delete) {
261            let _ = fs::remove_file(path);
262        }
263    }
264
265    /// Return a sorted list of checkpoint file paths for `state_id` (oldest first).
266    pub fn list_checkpoints(&self, state_id: &str) -> Vec<String> {
267        let dir = Path::new(&self.config.storage_path);
268        let prefix = format!("{state_id}_");
269        let suffix = ".ckpt";
270
271        let mut paths: Vec<PathBuf> = match fs::read_dir(dir) {
272            Ok(entries) => entries
273                .filter_map(|e| e.ok())
274                .map(|e| e.path())
275                .filter(|p| {
276                    p.file_name()
277                        .and_then(|n| n.to_str())
278                        .map(|n| n.starts_with(&prefix) && n.ends_with(suffix))
279                        .unwrap_or(false)
280                })
281                .collect(),
282            Err(_) => Vec::new(),
283        };
284
285        // Sort by the numeric timestamp embedded in the filename.
286        paths.sort_by_key(|p| extract_timestamp_ms(p));
287        paths
288            .into_iter()
289            .map(|p| p.to_string_lossy().into_owned())
290            .collect()
291    }
292}
293
294// ---------------------------------------------------------------------------
295// CheckpointBarrier
296// ---------------------------------------------------------------------------
297
298/// A barrier token that can be injected into a stream to signal all downstream
299/// operators to checkpoint their state.
300#[derive(Debug, Clone)]
301pub struct CheckpointBarrier {
302    /// Unique barrier identifier, monotonically increasing.
303    pub checkpoint_id: u64,
304    /// Wall-clock timestamp when the barrier was created (ms since epoch).
305    pub timestamp_ms: u64,
306}
307
308impl CheckpointBarrier {
309    /// Create a new barrier with the given id, stamping the current time.
310    pub fn new(checkpoint_id: u64) -> Self {
311        Self {
312            checkpoint_id,
313            timestamp_ms: current_time_ms(),
314        }
315    }
316
317    /// Create a barrier with an explicit timestamp (useful in tests).
318    pub fn with_timestamp(checkpoint_id: u64, timestamp_ms: u64) -> Self {
319        Self {
320            checkpoint_id,
321            timestamp_ms,
322        }
323    }
324}
325
326// ---------------------------------------------------------------------------
327// Helpers
328// ---------------------------------------------------------------------------
329
330fn current_time_ms() -> u64 {
331    SystemTime::now()
332        .duration_since(UNIX_EPOCH)
333        .map(|d| d.as_millis() as u64)
334        .unwrap_or(0)
335}
336
337/// Extract the numeric timestamp from a path like `{prefix}_{ts}.ckpt`.
338fn extract_timestamp_ms(path: &Path) -> u64 {
339    path.file_name()
340        .and_then(|n| n.to_str())
341        .and_then(|name| {
342            // Find the last '_' and parse everything between it and '.ckpt'.
343            let without_ext = name.strip_suffix(".ckpt")?;
344            let pos = without_ext.rfind('_')?;
345            without_ext[pos + 1..].parse::<u64>().ok()
346        })
347        .unwrap_or(0)
348}
349
350// ---------------------------------------------------------------------------
351// Tests
352// ---------------------------------------------------------------------------
353
354#[cfg(test)]
355mod tests {
356    use super::*;
357    use std::collections::HashMap;
358    use std::env;
359
360    fn temp_checkpoint_manager(suffix: &str) -> CheckpointManager {
361        let dir = env::temp_dir().join(format!("scirs2_ckpt_test_{suffix}"));
362        CheckpointManager::new(CheckpointConfig {
363            storage_path: dir.to_string_lossy().into_owned(),
364            max_checkpoints: 3,
365            ..Default::default()
366        })
367    }
368
369    #[test]
370    fn test_save_and_load_latest() {
371        let mgr = temp_checkpoint_manager("save_load");
372        let data = b"hello checkpoint world".to_vec();
373        mgr.save("job1", &data).expect("save should succeed");
374        let loaded = mgr.load_latest("job1").expect("load should succeed");
375        assert_eq!(loaded, data);
376    }
377
378    #[test]
379    fn test_cleanup_keeps_max_checkpoints() {
380        let mgr = temp_checkpoint_manager("cleanup");
381        // Save 5 checkpoints.
382        for i in 0..5u64 {
383            // Small sleep to ensure distinct ms timestamps on fast machines.
384            std::thread::sleep(std::time::Duration::from_millis(2));
385            mgr.save("job2", &i.to_le_bytes()).expect("save");
386        }
387        let files = mgr.list_checkpoints("job2");
388        assert!(
389            files.len() <= mgr.config.max_checkpoints,
390            "Expected ≤ {} checkpoints, found {}",
391            mgr.config.max_checkpoints,
392            files.len()
393        );
394    }
395
396    #[test]
397    fn test_load_latest_no_checkpoint_returns_error() {
398        let mgr = temp_checkpoint_manager("missing");
399        assert!(mgr.load_latest("nonexistent_id").is_err());
400    }
401
402    #[test]
403    fn test_vec_f64_round_trip() {
404        let v = vec![1.0_f64, 2.5, std::f64::consts::PI, -42.0];
405        let bytes = v.serialize();
406        let restored = Vec::<f64>::deserialize(&bytes).expect("deserialize");
407        assert_eq!(v, restored);
408    }
409
410    #[test]
411    fn test_hashmap_string_f64_round_trip() {
412        let mut m = HashMap::new();
413        m.insert("alpha".to_string(), 1.5_f64);
414        m.insert("beta".to_string(), -0.5_f64);
415        m.insert("gamma".to_string(), std::f64::consts::E);
416        let bytes = m.serialize();
417        let restored = HashMap::<String, f64>::deserialize(&bytes).expect("deserialize");
418        assert_eq!(m, restored);
419    }
420
421    #[test]
422    fn test_checkpoint_barrier_fields() {
423        let barrier = CheckpointBarrier::with_timestamp(42, 999_000);
424        assert_eq!(barrier.checkpoint_id, 42);
425        assert_eq!(barrier.timestamp_ms, 999_000);
426    }
427}
scirs2_io/streaming/checkpoint.rs

scirs2_io/streaming/
checkpoint.rs
