use std::sync::{Arc, Weak};
use aion_core::{Event, RunId, TimerCancelCause, WorkflowId};
use aion_store::EventStore;
use aion_store::visibility::VisibilityStore;
use chrono::Utc;
use crate::durability::Recorder;
use crate::registry::{Registry, TerminalOutcome, WorkflowHandle};
use crate::runtime::RuntimeHandle;
use crate::time::timer_service::live_timers_in_active_segment;
use crate::time::{DeadlineHandler, DeadlineHandlerError, WORKFLOW_TIMEOUT_DESCRIPTOR};
use super::completion::terminal_outcome_from_history;
use super::visibility::upsert_workflow_visibility;
enum DeadlineDisposition {
Appended,
ResumeTeardown,
LoseCleanly,
}
pub struct WorkflowDeadlineHandler {
runtime: Weak<RuntimeHandle>,
store: Arc<dyn EventStore>,
visibility_store: Arc<dyn VisibilityStore>,
registry: Arc<Registry>,
}
impl WorkflowDeadlineHandler {
#[must_use]
pub fn new(
runtime: Weak<RuntimeHandle>,
store: Arc<dyn EventStore>,
visibility_store: Arc<dyn VisibilityStore>,
registry: Arc<Registry>,
) -> Self {
Self {
runtime,
store,
visibility_store,
registry,
}
}
async fn drive_timed_out(
&self,
workflow_id: WorkflowId,
run_id: RunId,
) -> Result<(), crate::EngineError> {
let Some(handle) = self.registry.get(&workflow_id, &run_id)? else {
return self
.finalize_timed_out_without_handle(&workflow_id, &run_id)
.await;
};
let disposition = self
.decide_disposition(&handle, &workflow_id, &run_id)
.await?;
match disposition {
DeadlineDisposition::LoseCleanly => Ok(()),
DeadlineDisposition::Appended | DeadlineDisposition::ResumeTeardown => {
self.tear_down(&handle, &workflow_id, &run_id).await
}
}
}
async fn decide_disposition(
&self,
handle: &WorkflowHandle,
workflow_id: &WorkflowId,
run_id: &RunId,
) -> Result<DeadlineDisposition, crate::EngineError> {
let recorder = handle.recorder();
let mut recorder = recorder.lock().await;
let history = self.store.read_history(workflow_id).await?;
match terminal_outcome_from_history(&history, run_id) {
Some(TerminalOutcome::TimedOut(_)) => {
tracing::debug!(
%workflow_id,
%run_id,
"workflow deadline re-fired after its WorkflowTimedOut was recorded; resuming teardown"
);
Ok(DeadlineDisposition::ResumeTeardown)
}
Some(_) => {
tracing::debug!(
%workflow_id,
%run_id,
"workflow deadline elapsed but another terminal was already recorded; retiring the deadline and losing"
);
crate::time::retire_run_deadline(&mut recorder, &history, run_id).await?;
Ok(DeadlineDisposition::LoseCleanly)
}
None => {
if crate::time::outstanding_deadline_timer(&history, run_id).is_none() {
tracing::debug!(
%workflow_id,
%run_id,
"workflow deadline elapsed but its timer was already retired; deadline loses"
);
return Ok(DeadlineDisposition::LoseCleanly);
}
recorder
.record_workflow_timed_out(Utc::now(), WORKFLOW_TIMEOUT_DESCRIPTOR.to_owned())
.await?;
Ok(DeadlineDisposition::Appended)
}
}
}
async fn tear_down(
&self,
handle: &WorkflowHandle,
workflow_id: &WorkflowId,
run_id: &RunId,
) -> Result<(), crate::EngineError> {
self.retire_ordinary_timers(handle, workflow_id, run_id)
.await?;
match self.runtime.upgrade() {
Some(runtime) => {
if let Err(error) = runtime.cancel_pid(handle.pid()) {
tracing::debug!(
%workflow_id,
%run_id,
%error,
"workflow process already exited during deadline teardown"
);
}
}
None => {
return Err(crate::EngineError::Runtime {
reason: format!(
"runtime dropped during deadline teardown of {workflow_id}/{run_id}; a later re-fire resumes teardown"
),
});
}
}
upsert_workflow_visibility(
Arc::clone(&self.store),
Arc::clone(&self.visibility_store),
workflow_id,
run_id,
)
.await?;
handle.completion().notify(TerminalOutcome::TimedOut(
WORKFLOW_TIMEOUT_DESCRIPTOR.to_owned(),
));
self.retire_deadline(handle, workflow_id, run_id).await?;
self.registry.remove(workflow_id, run_id)?;
Ok(())
}
async fn retire_ordinary_timers(
&self,
handle: &WorkflowHandle,
workflow_id: &WorkflowId,
run_id: &RunId,
) -> Result<(), crate::EngineError> {
let recorder = handle.recorder();
let mut recorder = recorder.lock().await;
let history = self.store.read_history(workflow_id).await?;
record_ordinary_timer_retirements(&mut recorder, &history, run_id).await?;
Ok(())
}
async fn finalize_timed_out_without_handle(
&self,
workflow_id: &WorkflowId,
run_id: &RunId,
) -> Result<(), crate::EngineError> {
let history = self.store.read_history(workflow_id).await?;
if !matches!(
terminal_outcome_from_history(&history, run_id),
Some(TerminalOutcome::TimedOut(_))
) {
tracing::debug!(
%workflow_id,
%run_id,
"unregistered deadline elapsed for a run that is not TimedOut; nothing to finalize"
);
return Ok(());
}
let head = history.iter().map(Event::seq).max().unwrap_or_default();
let mut recorder = Recorder::resume_at(workflow_id.clone(), Arc::clone(&self.store), head);
record_ordinary_timer_retirements(&mut recorder, &history, run_id).await?;
upsert_workflow_visibility(
Arc::clone(&self.store),
Arc::clone(&self.visibility_store),
workflow_id,
run_id,
)
.await?;
crate::time::retire_run_deadline(&mut recorder, &history, run_id).await?;
Ok(())
}
async fn retire_deadline(
&self,
handle: &WorkflowHandle,
workflow_id: &WorkflowId,
run_id: &RunId,
) -> Result<(), crate::EngineError> {
let recorder = handle.recorder();
let mut recorder = recorder.lock().await;
let history = self.store.read_history(workflow_id).await?;
crate::time::retire_run_deadline(&mut recorder, &history, run_id).await?;
Ok(())
}
}
async fn record_ordinary_timer_retirements(
recorder: &mut Recorder,
history: &[Event],
run_id: &RunId,
) -> Result<(), crate::durability::DurabilityError> {
let deadline = crate::time::outstanding_deadline_timer(history, run_id);
for timer_id in live_timers_in_active_segment(history) {
if deadline.as_ref() == Some(&timer_id) {
continue;
}
recorder
.record_timer_cancelled(Utc::now(), timer_id, TimerCancelCause::WorkflowIntent)
.await?;
}
Ok(())
}
#[async_trait::async_trait]
impl DeadlineHandler for WorkflowDeadlineHandler {
async fn on_deadline_elapsed(
&self,
workflow_id: WorkflowId,
run_id: RunId,
) -> Result<(), DeadlineHandlerError> {
self.drive_timed_out(workflow_id, run_id)
.await
.map_err(|error| DeadlineHandlerError(error.to_string()))
}
}
#[cfg(test)]
#[path = "deadline_tests.rs"]
mod tests;