Skip to main content

oximedia_workflow/
workflow_checkpoint.rs

1//! Workflow checkpoint management for `oximedia-workflow`.
2//!
3//! [`CheckpointManager`] snapshots workflow execution state at configurable
4//! intervals so that a failed workflow can be resumed from the last checkpoint
5//! instead of being re-run from scratch.
6
7#![allow(dead_code)]
8
9use serde::{Deserialize, Serialize};
10use std::collections::HashMap;
11use std::time::{Duration, SystemTime, UNIX_EPOCH};
12
13// ---------------------------------------------------------------------------
14// Checkpoint policy
15// ---------------------------------------------------------------------------
16
17/// Describes when checkpoints should be captured.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
19pub enum CheckpointPolicy {
20    /// Never checkpoint (fastest, but no recovery).
21    Never,
22    /// Checkpoint after every completed task.
23    AfterEveryTask,
24    /// Checkpoint at a fixed wall-clock interval.
25    Interval,
26    /// Checkpoint only at explicitly marked tasks.
27    Explicit,
28}
29
30impl CheckpointPolicy {
31    /// Returns a human-readable label.
32    #[must_use]
33    pub const fn label(self) -> &'static str {
34        match self {
35            Self::Never => "Never",
36            Self::AfterEveryTask => "After Every Task",
37            Self::Interval => "Interval",
38            Self::Explicit => "Explicit",
39        }
40    }
41
42    /// Returns all variants.
43    #[must_use]
44    pub const fn all() -> &'static [CheckpointPolicy] {
45        &[
46            Self::Never,
47            Self::AfterEveryTask,
48            Self::Interval,
49            Self::Explicit,
50        ]
51    }
52}
53
54impl std::fmt::Display for CheckpointPolicy {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        f.write_str(self.label())
57    }
58}
59
60// ---------------------------------------------------------------------------
61// Checkpoint
62// ---------------------------------------------------------------------------
63
64/// A snapshot of workflow execution state at a point in time.
65#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
66pub struct Checkpoint {
67    /// Unique checkpoint identifier.
68    pub id: u64,
69    /// ID of the workflow this checkpoint belongs to.
70    pub workflow_id: String,
71    /// Unix timestamp (seconds) when the checkpoint was taken.
72    pub timestamp_secs: u64,
73    /// IDs of tasks that had completed at checkpoint time.
74    pub completed_tasks: Vec<String>,
75    /// IDs of tasks that were running at checkpoint time.
76    pub running_tasks: Vec<String>,
77    /// Arbitrary key-value state to persist (e.g. intermediate outputs).
78    pub state_data: HashMap<String, String>,
79    /// Human-readable label (optional).
80    pub label: Option<String>,
81}
82
83impl Checkpoint {
84    /// Creates a new checkpoint with the current timestamp.
85    #[must_use]
86    pub fn new(id: u64, workflow_id: impl Into<String>) -> Self {
87        let ts = SystemTime::now()
88            .duration_since(UNIX_EPOCH)
89            .unwrap_or(Duration::ZERO)
90            .as_secs();
91        Self {
92            id,
93            workflow_id: workflow_id.into(),
94            timestamp_secs: ts,
95            completed_tasks: Vec::new(),
96            running_tasks: Vec::new(),
97            state_data: HashMap::new(),
98            label: None,
99        }
100    }
101
102    /// Sets a state key-value pair.
103    pub fn set_state(&mut self, key: impl Into<String>, value: impl Into<String>) {
104        self.state_data.insert(key.into(), value.into());
105    }
106
107    /// Gets a state value by key.
108    #[must_use]
109    pub fn get_state(&self, key: &str) -> Option<&String> {
110        self.state_data.get(key)
111    }
112
113    /// Returns `true` if a given task was completed at this checkpoint.
114    #[must_use]
115    pub fn is_task_completed(&self, task_id: &str) -> bool {
116        self.completed_tasks.iter().any(|t| t == task_id)
117    }
118
119    /// Returns the total number of tasks (completed + running).
120    #[must_use]
121    pub fn task_count(&self) -> usize {
122        self.completed_tasks.len() + self.running_tasks.len()
123    }
124}
125
126// ---------------------------------------------------------------------------
127// Checkpoint manager
128// ---------------------------------------------------------------------------
129
130/// Manages a series of checkpoints for one workflow.
131#[derive(Debug, Clone)]
132pub struct CheckpointManager {
133    /// Active policy.
134    policy: CheckpointPolicy,
135    /// Interval in seconds (used only with `Interval` policy).
136    interval_secs: u64,
137    /// Stored checkpoints, newest last.
138    checkpoints: Vec<Checkpoint>,
139    /// Auto-incrementing ID counter.
140    next_id: u64,
141    /// Maximum number of checkpoints to retain (0 = unlimited).
142    max_retained: usize,
143}
144
145impl Default for CheckpointManager {
146    fn default() -> Self {
147        Self {
148            policy: CheckpointPolicy::AfterEveryTask,
149            interval_secs: 60,
150            checkpoints: Vec::new(),
151            next_id: 1,
152            max_retained: 0,
153        }
154    }
155}
156
157impl CheckpointManager {
158    /// Creates a new manager with the given policy.
159    #[must_use]
160    pub fn new(policy: CheckpointPolicy) -> Self {
161        Self {
162            policy,
163            ..Default::default()
164        }
165    }
166
167    /// Creates a manager with a fixed interval policy.
168    #[must_use]
169    pub fn with_interval(interval_secs: u64) -> Self {
170        Self {
171            policy: CheckpointPolicy::Interval,
172            interval_secs,
173            ..Default::default()
174        }
175    }
176
177    /// Sets the maximum number of checkpoints to retain.
178    pub fn set_max_retained(&mut self, max: usize) {
179        self.max_retained = max;
180    }
181
182    /// Returns the current policy.
183    #[must_use]
184    pub fn policy(&self) -> CheckpointPolicy {
185        self.policy
186    }
187
188    /// Returns the number of stored checkpoints.
189    #[must_use]
190    pub fn count(&self) -> usize {
191        self.checkpoints.len()
192    }
193
194    /// Creates and stores a new checkpoint for the given workflow.
195    pub fn create_checkpoint(&mut self, workflow_id: &str) -> &Checkpoint {
196        let id = self.next_id;
197        self.next_id += 1;
198        let cp = Checkpoint::new(id, workflow_id);
199        self.checkpoints.push(cp);
200        self.enforce_retention();
201        self.checkpoints
202            .last()
203            .expect("invariant: checkpoint just pushed above")
204    }
205
206    /// Returns the latest checkpoint, if any.
207    #[must_use]
208    pub fn latest(&self) -> Option<&Checkpoint> {
209        self.checkpoints.last()
210    }
211
212    /// Returns all checkpoints for a given workflow ID.
213    #[must_use]
214    pub fn checkpoints_for(&self, workflow_id: &str) -> Vec<&Checkpoint> {
215        self.checkpoints
216            .iter()
217            .filter(|c| c.workflow_id == workflow_id)
218            .collect()
219    }
220
221    /// Removes all checkpoints.
222    pub fn clear(&mut self) {
223        self.checkpoints.clear();
224    }
225
226    /// Returns the configured interval in seconds.
227    #[must_use]
228    pub fn interval_secs(&self) -> u64 {
229        self.interval_secs
230    }
231
232    /// Enforces the retention limit by dropping the oldest checkpoints.
233    fn enforce_retention(&mut self) {
234        if self.max_retained > 0 && self.checkpoints.len() > self.max_retained {
235            let excess = self.checkpoints.len() - self.max_retained;
236            self.checkpoints.drain(0..excess);
237        }
238    }
239}
240
241// ===========================================================================
242// Tests
243// ===========================================================================
244
245#[cfg(test)]
246mod tests {
247    use super::*;
248
249    // -- CheckpointPolicy ---------------------------------------------------
250
251    #[test]
252    fn test_policy_label() {
253        assert_eq!(CheckpointPolicy::Never.label(), "Never");
254        assert_eq!(CheckpointPolicy::AfterEveryTask.label(), "After Every Task");
255    }
256
257    #[test]
258    fn test_policy_display() {
259        assert_eq!(format!("{}", CheckpointPolicy::Interval), "Interval");
260    }
261
262    #[test]
263    fn test_policy_all() {
264        assert_eq!(CheckpointPolicy::all().len(), 4);
265    }
266
267    // -- Checkpoint ---------------------------------------------------------
268
269    #[test]
270    fn test_checkpoint_new() {
271        let cp = Checkpoint::new(1, "wf-001");
272        assert_eq!(cp.id, 1);
273        assert_eq!(cp.workflow_id, "wf-001");
274        assert!(cp.completed_tasks.is_empty());
275    }
276
277    #[test]
278    fn test_checkpoint_state() {
279        let mut cp = Checkpoint::new(1, "wf-001");
280        let out = std::env::temp_dir()
281            .join("oximedia-workflow-cp-out.mp4")
282            .to_string_lossy()
283            .into_owned();
284        cp.set_state("output_path", &out);
285        assert_eq!(
286            cp.get_state("output_path").expect("should succeed in test"),
287            &out
288        );
289        assert!(cp.get_state("missing").is_none());
290    }
291
292    #[test]
293    fn test_checkpoint_task_completed() {
294        let mut cp = Checkpoint::new(1, "wf-001");
295        cp.completed_tasks.push("task-a".to_string());
296        assert!(cp.is_task_completed("task-a"));
297        assert!(!cp.is_task_completed("task-b"));
298    }
299
300    #[test]
301    fn test_checkpoint_task_count() {
302        let mut cp = Checkpoint::new(1, "wf-001");
303        cp.completed_tasks.push("a".to_string());
304        cp.running_tasks.push("b".to_string());
305        assert_eq!(cp.task_count(), 2);
306    }
307
308    // -- CheckpointManager --------------------------------------------------
309
310    #[test]
311    fn test_manager_default() {
312        let m = CheckpointManager::default();
313        assert_eq!(m.policy(), CheckpointPolicy::AfterEveryTask);
314        assert_eq!(m.count(), 0);
315    }
316
317    #[test]
318    fn test_manager_create_checkpoint() {
319        let mut m = CheckpointManager::new(CheckpointPolicy::AfterEveryTask);
320        m.create_checkpoint("wf-001");
321        assert_eq!(m.count(), 1);
322        assert_eq!(
323            m.latest().expect("should succeed in test").workflow_id,
324            "wf-001"
325        );
326    }
327
328    #[test]
329    fn test_manager_multiple_checkpoints() {
330        let mut m = CheckpointManager::new(CheckpointPolicy::AfterEveryTask);
331        m.create_checkpoint("wf-001");
332        m.create_checkpoint("wf-001");
333        m.create_checkpoint("wf-002");
334        assert_eq!(m.count(), 3);
335        assert_eq!(m.checkpoints_for("wf-001").len(), 2);
336        assert_eq!(m.checkpoints_for("wf-002").len(), 1);
337    }
338
339    #[test]
340    fn test_manager_retention_limit() {
341        let mut m = CheckpointManager::new(CheckpointPolicy::AfterEveryTask);
342        m.set_max_retained(2);
343        m.create_checkpoint("wf-001");
344        m.create_checkpoint("wf-001");
345        m.create_checkpoint("wf-001");
346        assert_eq!(m.count(), 2);
347        // The oldest one (id=1) should have been dropped.
348        assert_eq!(m.latest().expect("should succeed in test").id, 3);
349    }
350
351    #[test]
352    fn test_manager_clear() {
353        let mut m = CheckpointManager::new(CheckpointPolicy::AfterEveryTask);
354        m.create_checkpoint("wf-001");
355        m.clear();
356        assert_eq!(m.count(), 0);
357        assert!(m.latest().is_none());
358    }
359
360    #[test]
361    fn test_manager_with_interval() {
362        let m = CheckpointManager::with_interval(120);
363        assert_eq!(m.policy(), CheckpointPolicy::Interval);
364        assert_eq!(m.interval_secs(), 120);
365    }
366
367    #[test]
368    fn test_manager_latest_none() {
369        let m = CheckpointManager::new(CheckpointPolicy::Never);
370        assert!(m.latest().is_none());
371    }
372}