Skip to main content

beam_core/workflow_runtime/
completion.rs

1use anyhow::Result;
2use serde_json::Value;
3
4use crate::workflow_orchestrator::OrchestratorAction;
5use crate::{EventDraft, EventLog, WorkflowActor};
6
7use super::WorkflowDispatchOutcome;
8use super::helpers::write_json_blob;
9
10pub async fn complete_node_succeeded(
11    log: &mut EventLog,
12    action: &crate::OrchestratorAction,
13) -> Result<()> {
14    if let OrchestratorAction::CompleteNodeSucceeded {
15        node_id,
16        last_activity_id,
17        ..
18    } = action
19    {
20        let _ = log.append(EventDraft {
21            event_type: "nodeSucceeded".to_string(),
22            actor: WorkflowActor::Scheduler,
23            payload: serde_json::json!({
24                "nodeId": node_id,
25                "lastActivityId": last_activity_id,
26            }),
27            timestamp: None,
28            payload_hash: None,
29        })?;
30        Ok(())
31    } else {
32        anyhow::bail!("complete_node_succeeded called with wrong action")
33    }
34}
35
36pub async fn complete_node_failed(
37    log: &mut EventLog,
38    action: &crate::OrchestratorAction,
39) -> Result<()> {
40    if let OrchestratorAction::CompleteNodeFailed {
41        node_id,
42        last_activity_id,
43        error_class,
44    } = action
45    {
46        let _ = log.append(EventDraft {
47            event_type: "nodeFailed".to_string(),
48            actor: WorkflowActor::Scheduler,
49            payload: serde_json::json!({
50                "nodeId": node_id,
51                "lastActivityId": last_activity_id,
52                "errorClass": error_class,
53            }),
54            timestamp: None,
55            payload_hash: None,
56        })?;
57        Ok(())
58    } else {
59        anyhow::bail!("complete_node_failed called with wrong action")
60    }
61}
62
63pub async fn complete_run_succeeded(
64    log: &mut EventLog,
65    action: &crate::OrchestratorAction,
66) -> Result<()> {
67    if let OrchestratorAction::CompleteRunSucceeded { output_ref, .. } = action {
68        let _ = log.append(EventDraft {
69            event_type: "runSucceeded".to_string(),
70            actor: WorkflowActor::Scheduler,
71            payload: serde_json::json!({
72                "outputRef": output_ref,
73            }),
74            timestamp: None,
75            payload_hash: None,
76        })?;
77        Ok(())
78    } else {
79        anyhow::bail!("complete_run_succeeded called with wrong action")
80    }
81}
82
83pub async fn complete_run_failed(
84    log: &mut EventLog,
85    action: &crate::OrchestratorAction,
86) -> Result<()> {
87    if let OrchestratorAction::CompleteRunFailed { failed_node_id } = action {
88        let root_cause_event_id = find_root_cause_event_id(log, failed_node_id).await?;
89        let _ = log.append(EventDraft {
90            event_type: "runFailed".to_string(),
91            actor: WorkflowActor::Scheduler,
92            payload: serde_json::json!({
93                "failedNodeId": failed_node_id,
94                "rootCauseEventId": root_cause_event_id,
95            }),
96            timestamp: None,
97            payload_hash: None,
98        })?;
99        Ok(())
100    } else {
101        anyhow::bail!("complete_run_failed called with wrong action")
102    }
103}
104
105pub(crate) async fn settle_work_result(
106    log: &mut EventLog,
107    activity_id: &str,
108    attempt_id: &str,
109    result: WorkflowDispatchOutcome,
110) -> Result<WorkflowDispatchOutcome> {
111    match &result {
112        WorkflowDispatchOutcome::Succeeded { output, .. } => {
113            let output_ref = write_json_blob(log, output.clone())?;
114            let _ = log.append(EventDraft {
115                event_type: "activitySucceeded".to_string(),
116                actor: WorkflowActor::Worker,
117                payload: serde_json::json!({
118                    "activityId": activity_id,
119                    "attemptId": attempt_id,
120                    "outputRef": output_ref,
121                }),
122                timestamp: None,
123                payload_hash: None,
124            })?;
125        }
126        WorkflowDispatchOutcome::Failed {
127            error_code,
128            error_class,
129            error_message,
130            ..
131        } => {
132            let _ = log.append(EventDraft {
133                event_type: "activityFailed".to_string(),
134                actor: WorkflowActor::Worker,
135                payload: serde_json::json!({
136                    "activityId": activity_id,
137                    "attemptId": attempt_id,
138                    "error": {
139                        "errorCode": error_code,
140                        "errorClass": error_class,
141                        "errorMessage": error_message,
142                    }
143                }),
144                timestamp: None,
145                payload_hash: None,
146            })?;
147        }
148        WorkflowDispatchOutcome::Cancelled {
149            cancel_origin_event_id,
150            ..
151        } => {
152            let _ = log.append(EventDraft {
153                event_type: "activityCanceled".to_string(),
154                actor: WorkflowActor::Worker,
155                payload: serde_json::json!({
156                    "activityId": activity_id,
157                    "attemptId": attempt_id,
158                    "cancelOriginEventId": cancel_origin_event_id,
159                }),
160                timestamp: None,
161                payload_hash: None,
162            })?;
163        }
164    }
165    Ok(result)
166}
167
168async fn find_root_cause_event_id(log: &EventLog, node_id: &str) -> Result<String> {
169    let events = log.read_all()?;
170    let mut node_failed_event_id: Option<String> = None;
171    let mut activity_failed_event_id: Option<String> = None;
172    let mut loop_finished_event_id: Option<String> = None;
173    let mut node_activities = std::collections::BTreeSet::new();
174    for e in &events {
175        match e.event_type.as_str() {
176            "attemptCreated" => {
177                if e.payload.get("nodeId").and_then(Value::as_str) == Some(node_id)
178                    && let Some(activity_id) = e.payload.get("activityId").and_then(Value::as_str)
179                {
180                    node_activities.insert(activity_id.to_string());
181                }
182            }
183            "activityFailed" => {
184                if let Some(activity_id) = e.payload.get("activityId").and_then(Value::as_str)
185                    && node_activities.contains(activity_id)
186                {
187                    activity_failed_event_id = Some(e.event_id.clone());
188                }
189            }
190            "nodeFailed" => {
191                if e.payload.get("nodeId").and_then(Value::as_str) == Some(node_id) {
192                    node_failed_event_id = Some(e.event_id.clone());
193                }
194            }
195            "loopFinished" => {
196                if e.payload.get("loopId").and_then(Value::as_str) == Some(node_id)
197                    && e.payload.get("resolution").and_then(Value::as_str) != Some("approved")
198                {
199                    loop_finished_event_id = Some(e.event_id.clone());
200                }
201            }
202            _ => {}
203        }
204    }
205    Ok(activity_failed_event_id
206        .or(node_failed_event_id)
207        .or(loop_finished_event_id)
208        .unwrap_or_else(|| {
209            events
210                .first()
211                .map(|e| e.event_id.clone())
212                .unwrap_or_default()
213        }))
214}