Skip to main content

oximedia_workflow/
workflow_snapshot.rs

1#![allow(dead_code)]
2//! Workflow state snapshotting for rollback, replay, and audit.
3//!
4//! Captures the full state of a workflow at a given point in time so that
5//! execution can be replayed, rolled back to a previous state, or audited.
6//! Snapshots are immutable once created and carry a monotonic sequence number.
7
8use std::collections::HashMap;
9
10/// Monotonically increasing snapshot sequence number.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
12pub struct SnapshotSeq(u64);
13
14impl SnapshotSeq {
15    /// Create a snapshot sequence number from a raw value.
16    #[must_use]
17    pub fn new(value: u64) -> Self {
18        Self(value)
19    }
20
21    /// Return the underlying value.
22    #[must_use]
23    pub fn value(self) -> u64 {
24        self.0
25    }
26}
27
28impl std::fmt::Display for SnapshotSeq {
29    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30        write!(f, "seq:{}", self.0)
31    }
32}
33
34/// The state of an individual task captured in the snapshot.
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub enum CapturedTaskState {
37    /// Task has not started.
38    Pending,
39    /// Task is currently running.
40    Running,
41    /// Task completed successfully.
42    Completed,
43    /// Task failed with a reason.
44    Failed(String),
45    /// Task was skipped.
46    Skipped,
47}
48
49impl std::fmt::Display for CapturedTaskState {
50    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
51        match self {
52            Self::Pending => write!(f, "pending"),
53            Self::Running => write!(f, "running"),
54            Self::Completed => write!(f, "completed"),
55            Self::Failed(reason) => write!(f, "failed: {reason}"),
56            Self::Skipped => write!(f, "skipped"),
57        }
58    }
59}
60
61/// A snapshot of a single task.
62#[derive(Debug, Clone)]
63pub struct TaskSnapshot {
64    /// Task identifier.
65    pub task_id: String,
66    /// State at snapshot time.
67    pub state: CapturedTaskState,
68    /// Number of retries that have occurred.
69    pub retry_count: u32,
70    /// Optional output data produced by the task.
71    pub output: Option<String>,
72}
73
74impl TaskSnapshot {
75    /// Create a new task snapshot.
76    pub fn new(task_id: impl Into<String>, state: CapturedTaskState) -> Self {
77        Self {
78            task_id: task_id.into(),
79            state,
80            retry_count: 0,
81            output: None,
82        }
83    }
84
85    /// Set retry count.
86    #[must_use]
87    pub fn with_retries(mut self, count: u32) -> Self {
88        self.retry_count = count;
89        self
90    }
91
92    /// Set output data.
93    pub fn with_output(mut self, output: impl Into<String>) -> Self {
94        self.output = Some(output.into());
95        self
96    }
97
98    /// Return true if the task has finished (completed, failed, or skipped).
99    #[must_use]
100    pub fn is_terminal(&self) -> bool {
101        matches!(
102            self.state,
103            CapturedTaskState::Completed
104                | CapturedTaskState::Failed(_)
105                | CapturedTaskState::Skipped
106        )
107    }
108}
109
110/// An immutable snapshot of the entire workflow at a point in time.
111#[derive(Debug, Clone)]
112pub struct WorkflowSnapshot {
113    /// Sequence number of this snapshot.
114    pub seq: SnapshotSeq,
115    /// Workflow identifier.
116    pub workflow_id: String,
117    /// Wall-clock time the snapshot was taken (seconds since epoch).
118    pub timestamp_secs: u64,
119    /// Label describing why this snapshot was taken.
120    pub label: String,
121    /// Per-task state snapshots.
122    pub tasks: Vec<TaskSnapshot>,
123    /// Arbitrary metadata attached to the snapshot.
124    pub metadata: HashMap<String, String>,
125}
126
127impl WorkflowSnapshot {
128    /// Create a new workflow snapshot.
129    pub fn new(
130        seq: SnapshotSeq,
131        workflow_id: impl Into<String>,
132        timestamp_secs: u64,
133        label: impl Into<String>,
134    ) -> Self {
135        Self {
136            seq,
137            workflow_id: workflow_id.into(),
138            timestamp_secs,
139            label: label.into(),
140            tasks: Vec::new(),
141            metadata: HashMap::new(),
142        }
143    }
144
145    /// Add a task snapshot.
146    pub fn add_task(&mut self, task: TaskSnapshot) {
147        self.tasks.push(task);
148    }
149
150    /// Attach metadata.
151    pub fn set_metadata(&mut self, key: impl Into<String>, value: impl Into<String>) {
152        self.metadata.insert(key.into(), value.into());
153    }
154
155    /// Return the number of tasks in this snapshot.
156    #[must_use]
157    pub fn task_count(&self) -> usize {
158        self.tasks.len()
159    }
160
161    /// Count tasks in a given state.
162    #[must_use]
163    pub fn count_in_state(&self, state: &CapturedTaskState) -> usize {
164        self.tasks.iter().filter(|t| &t.state == state).count()
165    }
166
167    /// Return the completion ratio (0.0..=1.0).
168    #[allow(clippy::cast_precision_loss)]
169    #[must_use]
170    pub fn completion_ratio(&self) -> f64 {
171        if self.tasks.is_empty() {
172            return 0.0;
173        }
174        let terminal = self.tasks.iter().filter(|t| t.is_terminal()).count();
175        terminal as f64 / self.tasks.len() as f64
176    }
177
178    /// Find a task snapshot by id.
179    #[must_use]
180    pub fn find_task(&self, task_id: &str) -> Option<&TaskSnapshot> {
181        self.tasks.iter().find(|t| t.task_id == task_id)
182    }
183}
184
185/// Manages a series of snapshots for a workflow.
186#[derive(Debug)]
187pub struct SnapshotStore {
188    /// Workflow identifier this store belongs to.
189    workflow_id: String,
190    /// Ordered list of snapshots.
191    snapshots: Vec<WorkflowSnapshot>,
192    /// Counter for assigning sequence numbers.
193    next_seq: u64,
194    /// Maximum number of snapshots to keep (0 = unlimited).
195    max_snapshots: usize,
196}
197
198impl SnapshotStore {
199    /// Create a new snapshot store for the given workflow.
200    pub fn new(workflow_id: impl Into<String>) -> Self {
201        Self {
202            workflow_id: workflow_id.into(),
203            snapshots: Vec::new(),
204            next_seq: 1,
205            max_snapshots: 0,
206        }
207    }
208
209    /// Set a maximum retention limit for snapshots.
210    #[must_use]
211    pub fn with_max_snapshots(mut self, max: usize) -> Self {
212        self.max_snapshots = max;
213        self
214    }
215
216    /// Take a new snapshot and add it to the store.
217    pub fn take_snapshot(
218        &mut self,
219        timestamp_secs: u64,
220        label: impl Into<String>,
221        tasks: Vec<TaskSnapshot>,
222    ) -> SnapshotSeq {
223        let seq = SnapshotSeq::new(self.next_seq);
224        self.next_seq += 1;
225
226        let mut snapshot =
227            WorkflowSnapshot::new(seq, self.workflow_id.clone(), timestamp_secs, label);
228        for t in tasks {
229            snapshot.add_task(t);
230        }
231
232        self.snapshots.push(snapshot);
233
234        // Enforce retention
235        if self.max_snapshots > 0 && self.snapshots.len() > self.max_snapshots {
236            let to_remove = self.snapshots.len() - self.max_snapshots;
237            self.snapshots.drain(..to_remove);
238        }
239
240        seq
241    }
242
243    /// Return the number of stored snapshots.
244    #[must_use]
245    pub fn len(&self) -> usize {
246        self.snapshots.len()
247    }
248
249    /// Check if the store has no snapshots.
250    #[must_use]
251    pub fn is_empty(&self) -> bool {
252        self.snapshots.is_empty()
253    }
254
255    /// Get the latest snapshot.
256    #[must_use]
257    pub fn latest(&self) -> Option<&WorkflowSnapshot> {
258        self.snapshots.last()
259    }
260
261    /// Get a snapshot by sequence number.
262    #[must_use]
263    pub fn get(&self, seq: SnapshotSeq) -> Option<&WorkflowSnapshot> {
264        self.snapshots.iter().find(|s| s.seq == seq)
265    }
266
267    /// List all snapshot sequence numbers and labels.
268    #[must_use]
269    pub fn list(&self) -> Vec<(SnapshotSeq, &str)> {
270        self.snapshots
271            .iter()
272            .map(|s| (s.seq, s.label.as_str()))
273            .collect()
274    }
275
276    /// Compare two snapshots and return task ids whose state changed.
277    #[must_use]
278    pub fn diff(&self, a: SnapshotSeq, b: SnapshotSeq) -> Option<Vec<String>> {
279        let snap_a = self.get(a)?;
280        let snap_b = self.get(b)?;
281
282        let map_a: HashMap<&str, &CapturedTaskState> = snap_a
283            .tasks
284            .iter()
285            .map(|t| (t.task_id.as_str(), &t.state))
286            .collect();
287
288        let mut changed = Vec::new();
289        for task in &snap_b.tasks {
290            match map_a.get(task.task_id.as_str()) {
291                Some(prev_state) if **prev_state != task.state => {
292                    changed.push(task.task_id.clone());
293                }
294                None => {
295                    changed.push(task.task_id.clone());
296                }
297                _ => {}
298            }
299        }
300
301        Some(changed)
302    }
303
304    /// Clear all snapshots.
305    pub fn clear(&mut self) {
306        self.snapshots.clear();
307    }
308}
309
310#[cfg(test)]
311mod tests {
312    use super::*;
313
314    fn sample_tasks() -> Vec<TaskSnapshot> {
315        vec![
316            TaskSnapshot::new("t1", CapturedTaskState::Completed),
317            TaskSnapshot::new("t2", CapturedTaskState::Running),
318            TaskSnapshot::new("t3", CapturedTaskState::Pending),
319        ]
320    }
321
322    #[test]
323    fn test_snapshot_seq_ordering() {
324        let a = SnapshotSeq::new(1);
325        let b = SnapshotSeq::new(2);
326        assert!(a < b);
327        assert_eq!(a.value(), 1);
328    }
329
330    #[test]
331    fn test_snapshot_seq_display() {
332        let s = SnapshotSeq::new(42);
333        assert_eq!(format!("{s}"), "seq:42");
334    }
335
336    #[test]
337    fn test_captured_task_state_display() {
338        assert_eq!(format!("{}", CapturedTaskState::Pending), "pending");
339        assert_eq!(format!("{}", CapturedTaskState::Running), "running");
340        assert_eq!(format!("{}", CapturedTaskState::Completed), "completed");
341        assert_eq!(
342            format!("{}", CapturedTaskState::Failed("boom".into())),
343            "failed: boom"
344        );
345        assert_eq!(format!("{}", CapturedTaskState::Skipped), "skipped");
346    }
347
348    #[test]
349    fn test_task_snapshot_terminal() {
350        assert!(TaskSnapshot::new("t", CapturedTaskState::Completed).is_terminal());
351        assert!(TaskSnapshot::new("t", CapturedTaskState::Failed("x".into())).is_terminal());
352        assert!(TaskSnapshot::new("t", CapturedTaskState::Skipped).is_terminal());
353        assert!(!TaskSnapshot::new("t", CapturedTaskState::Pending).is_terminal());
354        assert!(!TaskSnapshot::new("t", CapturedTaskState::Running).is_terminal());
355    }
356
357    #[test]
358    fn test_task_snapshot_builder() {
359        let ts = TaskSnapshot::new("t1", CapturedTaskState::Completed)
360            .with_retries(3)
361            .with_output("done");
362        assert_eq!(ts.retry_count, 3);
363        assert_eq!(ts.output.as_deref(), Some("done"));
364    }
365
366    #[test]
367    fn test_workflow_snapshot_basic() {
368        let mut ws = WorkflowSnapshot::new(SnapshotSeq::new(1), "wf-1", 1000, "initial");
369        ws.add_task(TaskSnapshot::new("t1", CapturedTaskState::Pending));
370        ws.set_metadata("user", "admin");
371        assert_eq!(ws.task_count(), 1);
372        assert_eq!(
373            ws.metadata.get("user").expect("should succeed in test"),
374            "admin"
375        );
376    }
377
378    #[test]
379    fn test_workflow_snapshot_completion_ratio() {
380        let mut ws = WorkflowSnapshot::new(SnapshotSeq::new(1), "wf-1", 1000, "mid");
381        for t in sample_tasks() {
382            ws.add_task(t);
383        }
384        // 1 completed out of 3
385        let ratio = ws.completion_ratio();
386        assert!((ratio - 1.0 / 3.0).abs() < 0.01);
387    }
388
389    #[test]
390    fn test_workflow_snapshot_empty_ratio() {
391        let ws = WorkflowSnapshot::new(SnapshotSeq::new(1), "wf-1", 1000, "empty");
392        assert!((ws.completion_ratio() - 0.0).abs() < f64::EPSILON);
393    }
394
395    #[test]
396    fn test_workflow_snapshot_find_task() {
397        let mut ws = WorkflowSnapshot::new(SnapshotSeq::new(1), "wf-1", 1000, "mid");
398        for t in sample_tasks() {
399            ws.add_task(t);
400        }
401        assert!(ws.find_task("t1").is_some());
402        assert!(ws.find_task("nonexistent").is_none());
403    }
404
405    #[test]
406    fn test_snapshot_store_take_and_get() {
407        let mut store = SnapshotStore::new("wf-1");
408        assert!(store.is_empty());
409        let seq = store.take_snapshot(1000, "first", sample_tasks());
410        assert_eq!(store.len(), 1);
411        let snap = store.get(seq).expect("should succeed in test");
412        assert_eq!(snap.label, "first");
413        assert_eq!(snap.task_count(), 3);
414    }
415
416    #[test]
417    fn test_snapshot_store_latest() {
418        let mut store = SnapshotStore::new("wf-1");
419        store.take_snapshot(1000, "first", vec![]);
420        store.take_snapshot(2000, "second", vec![]);
421        let latest = store.latest().expect("should succeed in test");
422        assert_eq!(latest.label, "second");
423    }
424
425    #[test]
426    fn test_snapshot_store_max_retention() {
427        let mut store = SnapshotStore::new("wf-1").with_max_snapshots(2);
428        store.take_snapshot(1000, "one", vec![]);
429        store.take_snapshot(2000, "two", vec![]);
430        store.take_snapshot(3000, "three", vec![]);
431        assert_eq!(store.len(), 2);
432        // "one" should have been evicted
433        let labels: Vec<&str> = store.list().iter().map(|(_, l)| *l).collect();
434        assert_eq!(labels, vec!["two", "three"]);
435    }
436
437    #[test]
438    fn test_snapshot_store_diff() {
439        let mut store = SnapshotStore::new("wf-1");
440        let tasks_a = vec![
441            TaskSnapshot::new("t1", CapturedTaskState::Pending),
442            TaskSnapshot::new("t2", CapturedTaskState::Pending),
443        ];
444        let seq_a = store.take_snapshot(1000, "before", tasks_a);
445
446        let tasks_b = vec![
447            TaskSnapshot::new("t1", CapturedTaskState::Completed),
448            TaskSnapshot::new("t2", CapturedTaskState::Pending),
449        ];
450        let seq_b = store.take_snapshot(2000, "after", tasks_b);
451
452        let changed = store.diff(seq_a, seq_b).expect("should succeed in test");
453        assert_eq!(changed, vec!["t1"]);
454    }
455
456    #[test]
457    fn test_snapshot_store_diff_new_task() {
458        let mut store = SnapshotStore::new("wf-1");
459        let seq_a = store.take_snapshot(1000, "before", vec![]);
460        let tasks_b = vec![TaskSnapshot::new("t1", CapturedTaskState::Pending)];
461        let seq_b = store.take_snapshot(2000, "after", tasks_b);
462
463        let changed = store.diff(seq_a, seq_b).expect("should succeed in test");
464        assert_eq!(changed, vec!["t1"]);
465    }
466
467    #[test]
468    fn test_snapshot_store_clear() {
469        let mut store = SnapshotStore::new("wf-1");
470        store.take_snapshot(1000, "one", vec![]);
471        store.take_snapshot(2000, "two", vec![]);
472        assert_eq!(store.len(), 2);
473        store.clear();
474        assert!(store.is_empty());
475    }
476}