Skip to main content

a3s_flow/engine/
operations.rs

1use chrono::{DateTime, Utc};
2
3use crate::error::{FlowError, Result};
4use crate::model::{
5    CancellationRequest, ChildOperationReference, FlowEvent, WorkflowProgress, WorkflowRunSnapshot,
6};
7
8use super::validation::{
9    ensure_child_operation_matches, ensure_progress_matches, is_event_conflict,
10};
11use super::FlowEngine;
12
13impl FlowEngine {
14    /// Request cleanup-aware cancellation and replay the workflow.
15    ///
16    /// The request atomically makes waits, hooks, and retrying/running steps
17    /// that existed before it non-actionable. Workflow code observes the
18    /// request through [`WorkflowContext::cancellation_request`](crate::WorkflowContext::cancellation_request),
19    /// performs host-owned cleanup with stable step identities, and returns
20    /// [`RuntimeCommand::Cancel`](crate::RuntimeCommand::Cancel). Repeating the
21    /// same request is idempotent.
22    pub async fn request_cancellation(
23        &self,
24        run_id: &str,
25        request: CancellationRequest,
26    ) -> Result<WorkflowRunSnapshot> {
27        for _ in 0..self.max_replay_iterations {
28            let snapshot = self.snapshot(run_id).await?;
29            if snapshot.status.is_terminal() {
30                return Ok(snapshot);
31            }
32            if let Some(existing) = &snapshot.cancellation {
33                if existing.request != request {
34                    return Err(FlowError::RunConflict {
35                        run_id: run_id.to_string(),
36                        reason: "cancellation request differs from the durable request".to_string(),
37                    });
38                }
39                match self.drive(run_id).await {
40                    Ok(snapshot) => return Ok(snapshot),
41                    Err(err) if is_event_conflict(&err) => continue,
42                    Err(err) => return Err(err),
43                }
44            }
45            match self
46                .record_event_at(
47                    run_id,
48                    snapshot.last_sequence,
49                    FlowEvent::RunCancellationRequested {
50                        request: request.clone(),
51                    },
52                )
53                .await
54            {
55                Ok(_) => match self.drive(run_id).await {
56                    Ok(snapshot) => return Ok(snapshot),
57                    Err(err) if is_event_conflict(&err) => continue,
58                    Err(err) => return Err(err),
59                },
60                Err(err) if is_event_conflict(&err) => continue,
61                Err(err) => return Err(err),
62            }
63        }
64
65        Err(FlowError::ReplayLimitExceeded(self.max_replay_iterations))
66    }
67
68    /// Immediately terminate a run as cancelled without replaying cleanup.
69    pub async fn force_cancel(&self, run_id: &str, reason: Option<String>) -> Result<()> {
70        self.terminate_run(run_id, FlowEvent::RunCancelled { reason })
71            .await
72    }
73
74    /// Backward-compatible immediate cancellation API.
75    ///
76    /// New cleanup-aware workflows should call [`Self::request_cancellation`].
77    pub async fn cancel(&self, run_id: &str, reason: Option<String>) -> Result<()> {
78        self.force_cancel(run_id, reason).await
79    }
80
81    /// Immediately terminate a run with a typed timeout outcome.
82    pub async fn terminate_for_timeout(
83        &self,
84        run_id: &str,
85        deadline: DateTime<Utc>,
86        reason: Option<String>,
87    ) -> Result<()> {
88        self.terminate_run(run_id, FlowEvent::RunTimedOut { deadline, reason })
89            .await
90    }
91
92    /// Explicitly abandon a run under a non-resumable host-shutdown policy.
93    ///
94    /// Ordinary process shutdown must not call this method: durable runs should
95    /// normally remain non-terminal and resume on a replacement host.
96    pub async fn terminate_for_host_shutdown(
97        &self,
98        run_id: &str,
99        reason: Option<String>,
100    ) -> Result<()> {
101        self.terminate_run(run_id, FlowEvent::RunHostShutdown { reason })
102            .await
103    }
104
105    /// Persist a host-reported progress update exactly once by `progress_id`.
106    pub async fn record_progress(&self, run_id: &str, progress: WorkflowProgress) -> Result<()> {
107        progress.validate()?;
108        for _ in 0..self.max_replay_iterations {
109            let snapshot = self.snapshot(run_id).await?;
110            if snapshot.status.is_terminal() {
111                return Err(FlowError::RunTerminal(run_id.to_string()));
112            }
113            if let Some(existing) = snapshot.progress(&progress.progress_id) {
114                ensure_progress_matches(run_id, existing, &progress)?;
115                return Ok(());
116            }
117            match self
118                .record_event_at(
119                    run_id,
120                    snapshot.last_sequence,
121                    FlowEvent::RunProgressRecorded {
122                        progress: progress.clone(),
123                    },
124                )
125                .await
126            {
127                Ok(_) => return Ok(()),
128                Err(err) if is_event_conflict(&err) => continue,
129                Err(err) => return Err(err),
130            }
131        }
132        Err(FlowError::ReplayLimitExceeded(self.max_replay_iterations))
133    }
134
135    /// Persist a parent-to-child operation reference exactly once by id.
136    pub async fn link_child_operation(
137        &self,
138        run_id: &str,
139        child: ChildOperationReference,
140    ) -> Result<()> {
141        child.validate()?;
142        for _ in 0..self.max_replay_iterations {
143            let snapshot = self.snapshot(run_id).await?;
144            if snapshot.status.is_terminal() {
145                return Err(FlowError::RunTerminal(run_id.to_string()));
146            }
147            if let Some(existing) = snapshot.child_operation(&child.reference_id) {
148                ensure_child_operation_matches(run_id, existing, &child)?;
149                return Ok(());
150            }
151            match self
152                .record_event_at(
153                    run_id,
154                    snapshot.last_sequence,
155                    FlowEvent::ChildOperationLinked {
156                        child: child.clone(),
157                    },
158                )
159                .await
160            {
161                Ok(_) => return Ok(()),
162                Err(err) if is_event_conflict(&err) => continue,
163                Err(err) => return Err(err),
164            }
165        }
166        Err(FlowError::ReplayLimitExceeded(self.max_replay_iterations))
167    }
168}