1#![allow(dead_code)]
2use std::collections::HashMap;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
12pub struct SnapshotSeq(u64);
13
14impl SnapshotSeq {
15 #[must_use]
17 pub fn new(value: u64) -> Self {
18 Self(value)
19 }
20
21 #[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#[derive(Debug, Clone, PartialEq, Eq)]
36pub enum CapturedTaskState {
37 Pending,
39 Running,
41 Completed,
43 Failed(String),
45 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#[derive(Debug, Clone)]
63pub struct TaskSnapshot {
64 pub task_id: String,
66 pub state: CapturedTaskState,
68 pub retry_count: u32,
70 pub output: Option<String>,
72}
73
74impl TaskSnapshot {
75 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 #[must_use]
87 pub fn with_retries(mut self, count: u32) -> Self {
88 self.retry_count = count;
89 self
90 }
91
92 pub fn with_output(mut self, output: impl Into<String>) -> Self {
94 self.output = Some(output.into());
95 self
96 }
97
98 #[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#[derive(Debug, Clone)]
112pub struct WorkflowSnapshot {
113 pub seq: SnapshotSeq,
115 pub workflow_id: String,
117 pub timestamp_secs: u64,
119 pub label: String,
121 pub tasks: Vec<TaskSnapshot>,
123 pub metadata: HashMap<String, String>,
125}
126
127impl WorkflowSnapshot {
128 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 pub fn add_task(&mut self, task: TaskSnapshot) {
147 self.tasks.push(task);
148 }
149
150 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 #[must_use]
157 pub fn task_count(&self) -> usize {
158 self.tasks.len()
159 }
160
161 #[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 #[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 #[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#[derive(Debug)]
187pub struct SnapshotStore {
188 workflow_id: String,
190 snapshots: Vec<WorkflowSnapshot>,
192 next_seq: u64,
194 max_snapshots: usize,
196}
197
198impl SnapshotStore {
199 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 #[must_use]
211 pub fn with_max_snapshots(mut self, max: usize) -> Self {
212 self.max_snapshots = max;
213 self
214 }
215
216 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 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 #[must_use]
245 pub fn len(&self) -> usize {
246 self.snapshots.len()
247 }
248
249 #[must_use]
251 pub fn is_empty(&self) -> bool {
252 self.snapshots.is_empty()
253 }
254
255 #[must_use]
257 pub fn latest(&self) -> Option<&WorkflowSnapshot> {
258 self.snapshots.last()
259 }
260
261 #[must_use]
263 pub fn get(&self, seq: SnapshotSeq) -> Option<&WorkflowSnapshot> {
264 self.snapshots.iter().find(|s| s.seq == seq)
265 }
266
267 #[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 #[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 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 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 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}