1use aion_core::{Event, WorkflowId, WorkflowStatus, status_from_events};
11use aion_store::EventStore;
12use aion_store::visibility::VisibilityStore;
13use async_trait::async_trait;
14use std::sync::Arc;
15
16use super::error::WorkloopError;
17use super::service::WorkloopWaker;
18use crate::durability::Recorder;
19use crate::lifecycle::reopen::{self, ReopenWorkflowContext};
20use crate::loader::WorkflowCatalog;
21use crate::registry::Registry;
22use crate::runtime::RuntimeHandle;
23use crate::supervision::SupervisionTree;
24
25pub struct EngineWorkloopWaker {
27 store: Arc<dyn EventStore>,
28 visibility_store: Arc<dyn VisibilityStore>,
29 catalog: Arc<WorkflowCatalog>,
30 runtime: Arc<RuntimeHandle>,
31 supervision: Arc<SupervisionTree>,
32 registry: Arc<Registry>,
33 search_attribute_schema: Arc<aion_core::SearchAttributeSchema>,
34}
35
36impl EngineWorkloopWaker {
37 #[must_use]
39 pub fn new(
40 store: Arc<dyn EventStore>,
41 visibility_store: Arc<dyn VisibilityStore>,
42 catalog: Arc<WorkflowCatalog>,
43 runtime: Arc<RuntimeHandle>,
44 supervision: Arc<SupervisionTree>,
45 registry: Arc<Registry>,
46 search_attribute_schema: Arc<aion_core::SearchAttributeSchema>,
47 ) -> Self {
48 Self {
49 store,
50 visibility_store,
51 catalog,
52 runtime,
53 supervision,
54 registry,
55 search_attribute_schema,
56 }
57 }
58}
59
60#[async_trait]
61impl WorkloopWaker for EngineWorkloopWaker {
62 async fn wake(&self, loop_id: &WorkflowId) -> Result<(), WorkloopError> {
63 let history = self.store.read_history(loop_id).await?;
64 let Some(run_id) = history.iter().rev().find_map(|event| match event {
65 Event::WorkflowStarted { run_id, .. } => Some(run_id.clone()),
66 _ => None,
67 }) else {
68 return Err(WorkloopError::Engine {
69 reason: format!("workloop {loop_id} has no recorded generation to wake"),
70 });
71 };
72
73 if self
76 .registry
77 .get(loop_id, &run_id)
78 .map_err(|error| WorkloopError::Engine {
79 reason: error.to_string(),
80 })?
81 .is_some()
82 {
83 return Ok(());
84 }
85
86 let status = status_from_events(&history);
87 if status != WorkflowStatus::Running {
88 return Err(WorkloopError::Engine {
89 reason: format!("workloop {loop_id} is {status:?}, not Running; refusing to wake"),
90 });
91 }
92
93 let head = history.iter().map(Event::seq).max().unwrap_or_default();
97 let recorder = Recorder::resume_at(loop_id.clone(), Arc::clone(&self.store), head)
98 .with_visibility(run_id.clone(), Arc::clone(&self.visibility_store));
99
100 let context = ReopenWorkflowContext {
101 store: Arc::clone(&self.store),
102 visibility_store: Arc::clone(&self.visibility_store),
103 catalog: Arc::clone(&self.catalog),
104 runtime: &self.runtime,
105 supervision: Arc::clone(&self.supervision),
106 registry: &self.registry,
107 search_attribute_schema: Arc::clone(&self.search_attribute_schema),
108 };
109 let segment = aion_core::run_segment(&history, &run_id);
110 let rearm = reopen::rearmable_timers(segment);
111 let handle = reopen::respawn_and_register(&context, loop_id, &run_id, &history, recorder)
112 .await
113 .map_err(|error| WorkloopError::Engine {
114 reason: format!("waking workloop {loop_id} failed: {error}"),
115 })?;
116 reopen::rearm_reopened_timers(&context, loop_id, handle.pid(), &rearm)
117 .await
118 .map_err(|error| WorkloopError::Engine {
119 reason: format!("re-arming woken workloop {loop_id} timers failed: {error}"),
120 })?;
121 Ok(())
122 }
123}