1use std::collections::HashMap;
15use std::sync::{Arc, Mutex};
16use std::time::{Duration, Instant};
17
18use futures::future::BoxFuture;
19use serde::{Deserialize, Serialize};
20
21pub type SagaId = String;
23
24pub type StepIndex = usize;
26
27#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
29pub enum SagaStatus {
30 Running,
32 Completed,
34 Compensated,
36 Failed,
38 TimedOut,
40}
41
42pub struct SagaStep {
44 pub name: String,
46 pub action: Box<dyn FnOnce() -> BoxFuture<'static, Result<(), String>> + Send>,
48 pub compensate: Option<Box<dyn FnOnce() -> BoxFuture<'static, Result<(), String>> + Send>>,
50}
51
52impl std::fmt::Debug for SagaStep {
53 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54 f.debug_struct("SagaStep")
55 .field("name", &self.name)
56 .field("has_compensate", &self.compensate.is_some())
57 .finish()
58 }
59}
60
61#[derive(Debug, Clone)]
63pub struct SagaResult {
64 pub saga_id: SagaId,
65 pub status: SagaStatus,
66 pub completed_steps: Vec<StepIndex>,
67 pub failed_step: Option<StepIndex>,
68 pub error: Option<String>,
69 pub elapsed: Duration,
70}
71
72#[derive(Debug, Clone)]
74pub struct TimeoutConfig {
75 pub overall: Option<Duration>,
77 pub per_step: Option<Duration>,
79}
80
81impl Default for TimeoutConfig {
82 fn default() -> Self {
83 Self {
84 overall: Some(Duration::from_secs(30)),
85 per_step: Some(Duration::from_secs(10)),
86 }
87 }
88}
89
90pub struct SagaStore {
92 instances: Mutex<HashMap<SagaId, StoredInstance>>,
93}
94
95struct StoredInstance {
96 status: SagaStatus,
97 completed_steps: Vec<StepIndex>,
98 failed_step: Option<StepIndex>,
99}
100
101impl SagaStore {
102 pub fn new() -> Self {
103 Self {
104 instances: Mutex::new(HashMap::new()),
105 }
106 }
107
108 pub fn create(&self, saga_id: &str) {
109 let mut instances = self.instances.lock().unwrap();
110 instances.insert(
111 saga_id.to_string(),
112 StoredInstance {
113 status: SagaStatus::Running,
114 completed_steps: Vec::new(),
115 failed_step: None,
116 },
117 );
118 }
119
120 pub fn mark_step_complete(&self, saga_id: &str, step: StepIndex) {
121 let mut instances = self.instances.lock().unwrap();
122 if let Some(inst) = instances.get_mut(saga_id) {
123 inst.completed_steps.push(step);
124 }
125 }
126
127 pub fn mark_failed(&self, saga_id: &str, step: StepIndex) {
128 let mut instances = self.instances.lock().unwrap();
129 if let Some(inst) = instances.get_mut(saga_id) {
130 inst.failed_step = Some(step);
131 inst.status = SagaStatus::Failed;
132 }
133 }
134
135 pub fn update_status(&self, saga_id: &str, status: SagaStatus) {
136 let mut instances = self.instances.lock().unwrap();
137 if let Some(inst) = instances.get_mut(saga_id) {
138 inst.status = status;
139 }
140 }
141
142 pub fn get_status(&self, saga_id: &str) -> Option<SagaStatus> {
143 let instances = self.instances.lock().unwrap();
144 instances.get(saga_id).map(|i| i.status.clone())
145 }
146
147 pub fn get_completed_steps(&self, saga_id: &str) -> Vec<StepIndex> {
148 let instances = self.instances.lock().unwrap();
149 instances
150 .get(saga_id)
151 .map(|i| i.completed_steps.clone())
152 .unwrap_or_default()
153 }
154}
155
156impl Default for SagaStore {
157 fn default() -> Self {
158 Self::new()
159 }
160}
161
162pub struct SagaCoordinator {
164 store: Arc<SagaStore>,
165 timeout_config: TimeoutConfig,
166}
167
168impl SagaCoordinator {
169 pub fn new(store: Arc<SagaStore>, timeout_config: TimeoutConfig) -> Self {
170 Self {
171 store,
172 timeout_config,
173 }
174 }
175
176 pub fn store(&self) -> &Arc<SagaStore> {
177 &self.store
178 }
179
180 pub async fn execute(&self, saga_id: &str, steps: Vec<SagaStep>) -> SagaResult {
184 let start = Instant::now();
185 let _n = steps.len();
186 self.store.create(saga_id);
187
188 let mut completed: Vec<StepIndex> = Vec::new();
189 let mut compensates: Vec<(
190 StepIndex,
191 Box<dyn FnOnce() -> BoxFuture<'static, Result<(), String>> + Send>,
192 )> = Vec::new();
193
194 for (i, step) in steps.into_iter().enumerate() {
195 if let Some(overall) = self.timeout_config.overall {
196 if start.elapsed() >= overall {
197 self.store.update_status(saga_id, SagaStatus::TimedOut);
198 return SagaResult {
199 saga_id: saga_id.to_string(),
200 status: SagaStatus::TimedOut,
201 completed_steps: completed,
202 failed_step: None,
203 error: Some("overall timeout".into()),
204 elapsed: start.elapsed(),
205 };
206 }
207 }
208
209 let action = step.action;
210 if let Some(c) = step.compensate {
211 compensates.push((i, c));
212 }
213
214 let step_result = if let Some(per_step) = self.timeout_config.per_step {
215 match tokio::time::timeout(per_step, action()).await {
216 Ok(r) => r,
217 Err(_) => Err(format!("step {} timed out", step.name)),
218 }
219 } else {
220 action().await
221 };
222
223 match step_result {
224 Ok(()) => {
225 completed.push(i);
226 self.store.mark_step_complete(saga_id, i);
227 }
228 Err(e) => {
229 self.store.mark_failed(saga_id, i);
230 let _ = self.compensate(saga_id, &mut compensates).await;
231 let status = self.store.get_status(saga_id).unwrap_or(SagaStatus::Failed);
232 return SagaResult {
233 saga_id: saga_id.to_string(),
234 status,
235 completed_steps: completed,
236 failed_step: Some(i),
237 error: Some(e),
238 elapsed: start.elapsed(),
239 };
240 }
241 }
242 }
243
244 self.store.update_status(saga_id, SagaStatus::Completed);
245 SagaResult {
246 saga_id: saga_id.to_string(),
247 status: SagaStatus::Completed,
248 completed_steps: completed,
249 failed_step: None,
250 error: None,
251 elapsed: start.elapsed(),
252 }
253 }
254
255 async fn compensate(
256 &self,
257 saga_id: &str,
258 compensates: &mut Vec<(
259 StepIndex,
260 Box<dyn FnOnce() -> BoxFuture<'static, Result<(), String>> + Send>,
261 )>,
262 ) -> Result<(), String> {
263 let mut all_ok = true;
264 while let Some((idx, comp)) = compensates.pop() {
265 let result = if let Some(per_step) = self.timeout_config.per_step {
266 match tokio::time::timeout(per_step, comp()).await {
267 Ok(r) => r,
268 Err(_) => Err(format!("compensate step {} timed out", idx)),
269 }
270 } else {
271 comp().await
272 };
273 if result.is_err() {
274 all_ok = false;
275 }
276 }
277 if all_ok {
278 self.store.update_status(saga_id, SagaStatus::Compensated);
279 Ok(())
280 } else {
281 self.store.update_status(saga_id, SagaStatus::Failed);
282 Err("compensation failed".into())
283 }
284 }
285}
286
287#[cfg(test)]
288mod tests {
289 use super::*;
290
291 fn ok_action() -> Box<dyn FnOnce() -> BoxFuture<'static, Result<(), String>> + Send> {
292 Box::new(|| {
293 let f: BoxFuture<'static, Result<(), String>> = Box::pin(async { Ok(()) });
294 f
295 })
296 }
297
298 fn fail_action(
299 msg: &str,
300 ) -> Box<dyn FnOnce() -> BoxFuture<'static, Result<(), String>> + Send> {
301 let msg = msg.to_string();
302 Box::new(move || {
303 let f: BoxFuture<'static, Result<(), String>> = Box::pin(async move { Err(msg) });
304 f
305 })
306 }
307
308 fn ok_compensate() -> Box<dyn FnOnce() -> BoxFuture<'static, Result<(), String>> + Send> {
309 Box::new(|| {
310 let f: BoxFuture<'static, Result<(), String>> = Box::pin(async { Ok(()) });
311 f
312 })
313 }
314
315 #[tokio::test]
316 async fn test_saga_all_success() {
317 let store = Arc::new(SagaStore::new());
318 let coord = SagaCoordinator::new(store.clone(), TimeoutConfig::default());
319 let steps = vec![
320 SagaStep {
321 name: "step1".into(),
322 action: ok_action(),
323 compensate: Some(ok_compensate()),
324 },
325 SagaStep {
326 name: "step2".into(),
327 action: ok_action(),
328 compensate: Some(ok_compensate()),
329 },
330 ];
331 let result = coord.execute("saga1", steps).await;
332 assert_eq!(result.status, SagaStatus::Completed);
333 assert_eq!(result.completed_steps, vec![0, 1]);
334 assert!(result.error.is_none());
335 }
336
337 #[tokio::test]
338 async fn test_saga_step2_fails_compensate() {
339 let store = Arc::new(SagaStore::new());
340 let coord = SagaCoordinator::new(store.clone(), TimeoutConfig::default());
341 let steps = vec![
342 SagaStep {
343 name: "step1".into(),
344 action: ok_action(),
345 compensate: Some(ok_compensate()),
346 },
347 SagaStep {
348 name: "step2".into(),
349 action: fail_action("step2 failed"),
350 compensate: Some(ok_compensate()),
351 },
352 ];
353 let result = coord.execute("saga2", steps).await;
354 assert_eq!(result.status, SagaStatus::Compensated);
355 assert_eq!(result.completed_steps, vec![0]);
356 assert_eq!(result.failed_step, Some(1));
357 assert!(result.error.is_some());
358 }
359
360 #[tokio::test]
361 async fn test_saga_no_compensate_on_fail() {
362 let store = Arc::new(SagaStore::new());
363 let coord = SagaCoordinator::new(store.clone(), TimeoutConfig::default());
364 let steps = vec![
365 SagaStep {
366 name: "step1".into(),
367 action: ok_action(),
368 compensate: None,
369 },
370 SagaStep {
371 name: "step2".into(),
372 action: fail_action("fail"),
373 compensate: None,
374 },
375 ];
376 let result = coord.execute("saga2", steps).await;
377 assert_eq!(result.status, SagaStatus::Compensated);
378 }
379
380 #[tokio::test]
381 async fn test_saga_empty_steps() {
382 let store = Arc::new(SagaStore::new());
383 let coord = SagaCoordinator::new(store.clone(), TimeoutConfig::default());
384 let result = coord.execute("saga3", vec![]).await;
385 assert_eq!(result.status, SagaStatus::Completed);
386 assert_eq!(result.completed_steps, Vec::<usize>::new());
387 }
388
389 #[tokio::test]
390 async fn test_saga_per_step_timeout() {
391 let store = Arc::new(SagaStore::new());
392 let config = TimeoutConfig {
393 overall: Some(Duration::from_secs(5)),
394 per_step: Some(Duration::from_millis(10)),
395 };
396 let coord = SagaCoordinator::new(store.clone(), config);
397 let slow: Box<dyn FnOnce() -> BoxFuture<'static, Result<(), String>> + Send> =
398 Box::new(|| {
399 let f: BoxFuture<'static, Result<(), String>> = Box::pin(async {
400 tokio::time::sleep(Duration::from_secs(1)).await;
401 Ok(())
402 });
403 f
404 });
405 let steps = vec![SagaStep {
406 name: "slow".into(),
407 action: slow,
408 compensate: None,
409 }];
410 let result = coord.execute("saga4", steps).await;
411 assert_eq!(result.status, SagaStatus::Compensated);
412 assert!(result.error.is_some());
413 }
414
415 #[tokio::test]
416 async fn test_saga_store_persistence() {
417 let store = SagaStore::new();
418 store.create("s1");
419 assert_eq!(store.get_status("s1"), Some(SagaStatus::Running));
420 store.mark_step_complete("s1", 0);
421 assert_eq!(store.get_completed_steps("s1"), vec![0]);
422 store.update_status("s1", SagaStatus::Completed);
423 assert_eq!(store.get_status("s1"), Some(SagaStatus::Completed));
424 }
425
426 #[tokio::test]
427 async fn test_saga_compensate_failed() {
428 let store = Arc::new(SagaStore::new());
429 let coord = SagaCoordinator::new(store.clone(), TimeoutConfig::default());
430 let bad_comp: Box<dyn FnOnce() -> BoxFuture<'static, Result<(), String>> + Send> =
431 Box::new(|| {
432 let f: BoxFuture<'static, Result<(), String>> =
433 Box::pin(async { Err("comp failed".into()) });
434 f
435 });
436 let steps = vec![
437 SagaStep {
438 name: "s1".into(),
439 action: ok_action(),
440 compensate: Some(bad_comp),
441 },
442 SagaStep {
443 name: "s2".into(),
444 action: fail_action("s2 fail"),
445 compensate: None,
446 },
447 ];
448 let result = coord.execute("saga5", steps).await;
449 assert_eq!(result.status, SagaStatus::Failed);
450 }
451
452 #[tokio::test]
453 async fn test_saga_overall_timeout() {
454 let store = Arc::new(SagaStore::new());
455 let config = TimeoutConfig {
456 overall: Some(Duration::from_millis(5)),
457 per_step: None,
458 };
459 let coord = SagaCoordinator::new(store.clone(), config);
460 let slow: Box<dyn FnOnce() -> BoxFuture<'static, Result<(), String>> + Send> =
461 Box::new(|| {
462 let f: BoxFuture<'static, Result<(), String>> = Box::pin(async {
463 tokio::time::sleep(Duration::from_millis(20)).await;
464 Ok(())
465 });
466 f
467 });
468 let steps = vec![
469 SagaStep {
470 name: "slow1".into(),
471 action: slow,
472 compensate: None,
473 },
474 SagaStep {
475 name: "slow2".into(),
476 action: ok_action(),
477 compensate: None,
478 },
479 ];
480 let result = coord.execute("saga6", steps).await;
481 assert_eq!(result.status, SagaStatus::TimedOut);
482 }
483
484 #[tokio::test]
485 async fn test_saga_three_steps_middle_fails() {
486 let store = Arc::new(SagaStore::new());
487 let coord = SagaCoordinator::new(store.clone(), TimeoutConfig::default());
488 let steps = vec![
489 SagaStep {
490 name: "s1".into(),
491 action: ok_action(),
492 compensate: Some(ok_compensate()),
493 },
494 SagaStep {
495 name: "s2".into(),
496 action: fail_action("s2 fail"),
497 compensate: Some(ok_compensate()),
498 },
499 SagaStep {
500 name: "s3".into(),
501 action: ok_action(),
502 compensate: Some(ok_compensate()),
503 },
504 ];
505 let result = coord.execute("saga7", steps).await;
506 assert_eq!(result.status, SagaStatus::Compensated);
507 assert_eq!(result.completed_steps, vec![0]);
508 assert_eq!(result.failed_step, Some(1));
509 }
510}