1#![allow(dead_code)]
8
9use serde::{Deserialize, Serialize};
10use std::collections::HashMap;
11use std::time::{Duration, SystemTime, UNIX_EPOCH};
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
19pub enum CheckpointPolicy {
20 Never,
22 AfterEveryTask,
24 Interval,
26 Explicit,
28}
29
30impl CheckpointPolicy {
31 #[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 #[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#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
66pub struct Checkpoint {
67 pub id: u64,
69 pub workflow_id: String,
71 pub timestamp_secs: u64,
73 pub completed_tasks: Vec<String>,
75 pub running_tasks: Vec<String>,
77 pub state_data: HashMap<String, String>,
79 pub label: Option<String>,
81}
82
83impl Checkpoint {
84 #[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 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 #[must_use]
109 pub fn get_state(&self, key: &str) -> Option<&String> {
110 self.state_data.get(key)
111 }
112
113 #[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 #[must_use]
121 pub fn task_count(&self) -> usize {
122 self.completed_tasks.len() + self.running_tasks.len()
123 }
124}
125
126#[derive(Debug, Clone)]
132pub struct CheckpointManager {
133 policy: CheckpointPolicy,
135 interval_secs: u64,
137 checkpoints: Vec<Checkpoint>,
139 next_id: u64,
141 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 #[must_use]
160 pub fn new(policy: CheckpointPolicy) -> Self {
161 Self {
162 policy,
163 ..Default::default()
164 }
165 }
166
167 #[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 pub fn set_max_retained(&mut self, max: usize) {
179 self.max_retained = max;
180 }
181
182 #[must_use]
184 pub fn policy(&self) -> CheckpointPolicy {
185 self.policy
186 }
187
188 #[must_use]
190 pub fn count(&self) -> usize {
191 self.checkpoints.len()
192 }
193
194 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 #[must_use]
208 pub fn latest(&self) -> Option<&Checkpoint> {
209 self.checkpoints.last()
210 }
211
212 #[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 pub fn clear(&mut self) {
223 self.checkpoints.clear();
224 }
225
226 #[must_use]
228 pub fn interval_secs(&self) -> u64 {
229 self.interval_secs
230 }
231
232 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#[cfg(test)]
246mod tests {
247 use super::*;
248
249 #[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 #[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 cp.set_state("output_path", "/tmp/out.mp4");
281 assert_eq!(
282 cp.get_state("output_path").expect("should succeed in test"),
283 "/tmp/out.mp4"
284 );
285 assert!(cp.get_state("missing").is_none());
286 }
287
288 #[test]
289 fn test_checkpoint_task_completed() {
290 let mut cp = Checkpoint::new(1, "wf-001");
291 cp.completed_tasks.push("task-a".to_string());
292 assert!(cp.is_task_completed("task-a"));
293 assert!(!cp.is_task_completed("task-b"));
294 }
295
296 #[test]
297 fn test_checkpoint_task_count() {
298 let mut cp = Checkpoint::new(1, "wf-001");
299 cp.completed_tasks.push("a".to_string());
300 cp.running_tasks.push("b".to_string());
301 assert_eq!(cp.task_count(), 2);
302 }
303
304 #[test]
307 fn test_manager_default() {
308 let m = CheckpointManager::default();
309 assert_eq!(m.policy(), CheckpointPolicy::AfterEveryTask);
310 assert_eq!(m.count(), 0);
311 }
312
313 #[test]
314 fn test_manager_create_checkpoint() {
315 let mut m = CheckpointManager::new(CheckpointPolicy::AfterEveryTask);
316 m.create_checkpoint("wf-001");
317 assert_eq!(m.count(), 1);
318 assert_eq!(
319 m.latest().expect("should succeed in test").workflow_id,
320 "wf-001"
321 );
322 }
323
324 #[test]
325 fn test_manager_multiple_checkpoints() {
326 let mut m = CheckpointManager::new(CheckpointPolicy::AfterEveryTask);
327 m.create_checkpoint("wf-001");
328 m.create_checkpoint("wf-001");
329 m.create_checkpoint("wf-002");
330 assert_eq!(m.count(), 3);
331 assert_eq!(m.checkpoints_for("wf-001").len(), 2);
332 assert_eq!(m.checkpoints_for("wf-002").len(), 1);
333 }
334
335 #[test]
336 fn test_manager_retention_limit() {
337 let mut m = CheckpointManager::new(CheckpointPolicy::AfterEveryTask);
338 m.set_max_retained(2);
339 m.create_checkpoint("wf-001");
340 m.create_checkpoint("wf-001");
341 m.create_checkpoint("wf-001");
342 assert_eq!(m.count(), 2);
343 assert_eq!(m.latest().expect("should succeed in test").id, 3);
345 }
346
347 #[test]
348 fn test_manager_clear() {
349 let mut m = CheckpointManager::new(CheckpointPolicy::AfterEveryTask);
350 m.create_checkpoint("wf-001");
351 m.clear();
352 assert_eq!(m.count(), 0);
353 assert!(m.latest().is_none());
354 }
355
356 #[test]
357 fn test_manager_with_interval() {
358 let m = CheckpointManager::with_interval(120);
359 assert_eq!(m.policy(), CheckpointPolicy::Interval);
360 assert_eq!(m.interval_secs(), 120);
361 }
362
363 #[test]
364 fn test_manager_latest_none() {
365 let m = CheckpointManager::new(CheckpointPolicy::Never);
366 assert!(m.latest().is_none());
367 }
368}