Skip to main content

sz_orm_core/
saga.rs

1//! v6.6.0 Saga 分布式事务协调器
2//!
3//! Saga 模式:正向执行 + 补偿编排,保证最终一致性。
4//!
5//! # 执行流程
6//! 1. 依次执行正向步骤
7//! 2. 步骤 i 失败 → 逆序执行已完成步骤的补偿事务
8//! 3. 全部成功 → Saga 完成
9//!
10//! # 超时处理
11//! - 全局超时:整个 Saga 的最大执行时间
12//! - 单步骤超时:每个步骤的最大执行时间
13
14use std::collections::HashMap;
15use std::sync::{Arc, Mutex};
16use std::time::{Duration, Instant};
17
18use futures::future::BoxFuture;
19use serde::{Deserialize, Serialize};
20
21/// Saga 实例 ID
22pub type SagaId = String;
23
24/// Saga 步骤索引
25pub type StepIndex = usize;
26
27/// Saga 执行状态
28#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
29pub enum SagaStatus {
30    /// 正在执行
31    Running,
32    /// 全部正向步骤成功
33    Completed,
34    /// 补偿完成(某步骤失败后补偿成功)
35    Compensated,
36    /// 补偿失败
37    Failed,
38    /// 超时
39    TimedOut,
40}
41
42/// Saga 步骤
43pub struct SagaStep {
44    /// 步骤名称
45    pub name: String,
46    /// 正向操作
47    pub action: Box<dyn FnOnce() -> BoxFuture<'static, Result<(), String>> + Send>,
48    /// �)偿操作
49    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/// Saga 执行结果
62#[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/// Saga 超时配置
73#[derive(Debug, Clone)]
74pub struct TimeoutConfig {
75    /// 全局超时(None = 无限制)
76    pub overall: Option<Duration>,
77    /// 单步骤超时(None = 无限制)
78    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
90/// Saga 状态存储(内存实现)
91pub 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
162/// Saga 协调器
163pub 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    /// 执行 Saga
181    ///
182    /// 依次执行正向步骤,失败时逆序补偿。
183    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}