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            self.ensure_runtime_build_available(run_id, &snapshot.spec)?;
33            if let Some(existing) = &snapshot.cancellation {
34                if existing.request != request {
35                    return Err(FlowError::RunConflict {
36                        run_id: run_id.to_string(),
37                        reason: "cancellation request differs from the durable request".to_string(),
38                    });
39                }
40                match self.drive(run_id).await {
41                    Ok(snapshot) => return Ok(snapshot),
42                    Err(err) if is_event_conflict(&err) => continue,
43                    Err(err) => return Err(err),
44                }
45            }
46            match self
47                .record_event_at(
48                    run_id,
49                    snapshot.last_sequence,
50                    FlowEvent::RunCancellationRequested {
51                        request: request.clone(),
52                    },
53                )
54                .await
55            {
56                Ok(_) => match self.drive(run_id).await {
57                    Ok(snapshot) => return Ok(snapshot),
58                    Err(err) if is_event_conflict(&err) => continue,
59                    Err(err) => return Err(err),
60                },
61                Err(err) if is_event_conflict(&err) => continue,
62                Err(err) => return Err(err),
63            }
64        }
65
66        Err(FlowError::ReplayLimitExceeded(self.max_replay_iterations))
67    }
68
69    /// Immediately terminate a run as cancelled without replaying cleanup.
70    pub async fn force_cancel(&self, run_id: &str, reason: Option<String>) -> Result<()> {
71        self.terminate_run(run_id, FlowEvent::RunCancelled { reason })
72            .await
73    }
74
75    /// Backward-compatible immediate cancellation API.
76    ///
77    /// New cleanup-aware workflows should call [`Self::request_cancellation`].
78    pub async fn cancel(&self, run_id: &str, reason: Option<String>) -> Result<()> {
79        self.force_cancel(run_id, reason).await
80    }
81
82    /// Immediately terminate a run with a typed timeout outcome.
83    pub async fn terminate_for_timeout(
84        &self,
85        run_id: &str,
86        deadline: DateTime<Utc>,
87        reason: Option<String>,
88    ) -> Result<()> {
89        self.terminate_run(run_id, FlowEvent::RunTimedOut { deadline, reason })
90            .await
91    }
92
93    /// Explicitly abandon a run under a non-resumable host-shutdown policy.
94    ///
95    /// Ordinary process shutdown must not call this method: durable runs should
96    /// normally remain non-terminal and resume on a replacement host.
97    pub async fn terminate_for_host_shutdown(
98        &self,
99        run_id: &str,
100        reason: Option<String>,
101    ) -> Result<()> {
102        self.terminate_run(run_id, FlowEvent::RunHostShutdown { reason })
103            .await
104    }
105
106    /// Persist a host-reported progress update exactly once by `progress_id`.
107    pub async fn record_progress(&self, run_id: &str, progress: WorkflowProgress) -> Result<()> {
108        progress.validate()?;
109        for _ in 0..self.max_replay_iterations {
110            let snapshot = self.snapshot(run_id).await?;
111            if snapshot.status.is_terminal() {
112                return Err(FlowError::RunTerminal(run_id.to_string()));
113            }
114            if let Some(existing) = snapshot.progress(&progress.progress_id) {
115                ensure_progress_matches(run_id, existing, &progress)?;
116                return Ok(());
117            }
118            match self
119                .record_event_at(
120                    run_id,
121                    snapshot.last_sequence,
122                    FlowEvent::RunProgressRecorded {
123                        progress: progress.clone(),
124                    },
125                )
126                .await
127            {
128                Ok(_) => return Ok(()),
129                Err(err) if is_event_conflict(&err) => continue,
130                Err(err) => return Err(err),
131            }
132        }
133        Err(FlowError::ReplayLimitExceeded(self.max_replay_iterations))
134    }
135
136    /// Persist a parent-to-child operation reference exactly once by id.
137    pub async fn link_child_operation(
138        &self,
139        run_id: &str,
140        child: ChildOperationReference,
141    ) -> Result<()> {
142        child.validate()?;
143        for _ in 0..self.max_replay_iterations {
144            let snapshot = self.snapshot(run_id).await?;
145            if snapshot.status.is_terminal() {
146                return Err(FlowError::RunTerminal(run_id.to_string()));
147            }
148            if let Some(existing) = snapshot.child_operation(&child.reference_id) {
149                ensure_child_operation_matches(run_id, existing, &child)?;
150                return Ok(());
151            }
152            match self
153                .record_event_at(
154                    run_id,
155                    snapshot.last_sequence,
156                    FlowEvent::ChildOperationLinked {
157                        child: child.clone(),
158                    },
159                )
160                .await
161            {
162                Ok(_) => return Ok(()),
163                Err(err) if is_event_conflict(&err) => continue,
164                Err(err) => return Err(err),
165            }
166        }
167        Err(FlowError::ReplayLimitExceeded(self.max_replay_iterations))
168    }
169}