use std::sync::Arc;
use aion_core::{Event, Payload, RunId, WorkflowId, WorkflowStatus};
use aion_store::EventStore;
use aion_store::visibility::VisibilityStore;
use chrono::Utc;
use super::error::WorkloopError;
use crate::durability::{Recorder, WorkflowStartRecord};
use crate::loader::WorkflowCatalog;
use crate::registry::{
CompletionNotifier, HandleResidency, Registry, TerminalOutcome, WorkflowHandle,
WorkflowHandleParts,
};
use crate::runtime::{RuntimeHandle, RuntimeInput};
pub const RETIRE_ENTRY: &str = "retire";
pub const RETIRE_ARITY: u32 = 1;
pub struct RetireInvocation<'a> {
pub loop_id: &'a WorkflowId,
pub runtime: &'a Arc<RuntimeHandle>,
pub catalog: &'a WorkflowCatalog,
pub registry: &'a Arc<Registry>,
pub store: &'a Arc<dyn EventStore>,
pub visibility_store: &'a Arc<dyn VisibilityStore>,
}
#[must_use]
pub fn missing_entry_refusal(loop_id: &WorkflowId, module: &str) -> String {
format!(
"workloop {loop_id} declares a retire body but its deployed module `{module}` exports no \
`{RETIRE_ENTRY}/{RETIRE_ARITY}`. Refusing to retire: the engine cannot distinguish a \
module compiled before the retire entry existed from a loop that declared no cleanup at \
all, and silently skipping a declared retirement is how a lease is lost or a queue is \
stranded. Nothing was recorded and the loop is still running. Redeploy `{module}` from a \
toolchain that emits `{RETIRE_ENTRY}/{RETIRE_ARITY}`, then retire again"
)
}
pub async fn run_retire_body(invocation: &RetireInvocation<'_>) -> Result<(), WorkloopError> {
let history = invocation.store.read_history(invocation.loop_id).await?;
refuse_if_terminal(invocation.loop_id, &history)?;
let generation = current_generation(invocation, &history)?;
if !invocation
.runtime
.module_exports_function(&generation.module, RETIRE_ENTRY)
{
return Err(WorkloopError::Engine {
reason: missing_entry_refusal(invocation.loop_id, &generation.module),
});
}
let (retirement_run, recorder) = open_retirement_generation(invocation, &generation).await?;
stand_down_resident_generation(
invocation.registry,
invocation.runtime,
invocation.loop_id,
&generation.carry,
)?;
let input =
RuntimeInput::from_payload(&generation.carry).map_err(|error| WorkloopError::Engine {
reason: format!("encoding the retire body's carry argument failed: {error}"),
})?;
let module = generation.module.clone();
let pid = invocation
.runtime
.spawn_workflow(&module, RETIRE_ENTRY, input)
.map_err(|error| WorkloopError::Engine {
reason: format!("spawning `{module}:{RETIRE_ENTRY}/{RETIRE_ARITY}` failed: {error}"),
})?;
if let Err(error) = publish_retire_body(invocation, &retirement_run, pid, &generation, recorder)
{
cancel_orphaned_body(invocation.runtime, invocation.loop_id, pid, &error);
return Err(error);
}
let outcome = await_retire_body(invocation.runtime, invocation.loop_id, pid).await;
if let Err(error) = invocation
.registry
.remove(invocation.loop_id, &retirement_run)
{
tracing::warn!(
loop_id = %invocation.loop_id,
error = %error,
"removing the retire body's registry publication failed; the terminal batch below \
appends through a fresh one-shot recorder and a stale publication would make that \
a second writer"
);
}
outcome
}
struct CurrentGeneration {
run_id: RunId,
workflow_type: String,
package_version: aion_core::PackageVersion,
loaded_version: aion_package::ContentHash,
module: String,
carry: Payload,
}
pub fn refuse_if_terminal(loop_id: &WorkflowId, history: &[Event]) -> Result<(), WorkloopError> {
if aion_core::current_lease_terminal(history).is_some() {
return Err(WorkloopError::Engine {
reason: format!(
"cannot retire workloop {loop_id}: its run already recorded a terminal, so it \
is not running and has nothing left to retire. No retire body was invoked — a \
declared cleanup that ran a second time would release an already-released \
lease and re-drain a drained queue"
),
});
}
Ok(())
}
fn stand_down_resident_generation(
registry: &Arc<Registry>,
runtime: &Arc<RuntimeHandle>,
loop_id: &WorkflowId,
carry: &Payload,
) -> Result<(), WorkloopError> {
let handles = registry.list().map_err(|error| WorkloopError::Engine {
reason: format!("listing the registry to stand down {loop_id} failed: {error}"),
})?;
for handle in handles
.into_iter()
.filter(|handle| handle.workflow_id() == loop_id)
{
let run_id = handle.run_id().clone();
let pid = handle.pid();
tracing::info!(
%loop_id,
run_id = %run_id,
pid,
"retirement is standing down the loop's resident generation before running its \
declared retire body; the generation's own work stops here"
);
handle.completion().notify(TerminalOutcome::ContinuedAsNew {
input: carry.clone(),
workflow_type: None,
parent_run_id: run_id.clone(),
});
registry
.remove(loop_id, &run_id)
.map_err(|error| WorkloopError::Engine {
reason: format!(
"removing the resident generation's handle for {loop_id} run {run_id} \
failed: {error}"
),
})?;
runtime
.cancel_pid(pid)
.map_err(|error| WorkloopError::Engine {
reason: format!(
"ending the resident generation's process {pid} for {loop_id} failed: \
{error}. Refusing to run the retire body: that process is still runnable \
on the loop's history and would be a second writer alongside the body"
),
})?;
}
Ok(())
}
async fn open_retirement_generation(
invocation: &RetireInvocation<'_>,
generation: &CurrentGeneration,
) -> Result<(RunId, Recorder), WorkloopError> {
let retirement_run = RunId::new_v4();
let start = WorkflowStartRecord {
workflow_type: generation.workflow_type.clone(),
input: generation.carry.clone(),
run_id: retirement_run.clone(),
parent_run_id: Some(generation.run_id.clone()),
parent_workflow_id: None,
package_version: generation.package_version.clone(),
};
if let Some(handle) = live_handle(invocation)? {
let recorder = handle.recorder();
let mut recorder = recorder.lock().await;
let history = invocation.store.read_history(invocation.loop_id).await?;
refuse_if_terminal(invocation.loop_id, &history)?;
recorder
.record_workloop_retirement_generation(
Utc::now(),
generation.carry.clone(),
generation.run_id.clone(),
start,
)
.await?;
let head = recorder.head();
drop(recorder);
return Ok((
retirement_run.clone(),
Recorder::resume_at(
invocation.loop_id.clone(),
Arc::clone(invocation.store),
head,
)
.with_visibility(retirement_run, Arc::clone(invocation.visibility_store)),
));
}
let history = invocation.store.read_history(invocation.loop_id).await?;
refuse_if_terminal(invocation.loop_id, &history)?;
let head = history.iter().map(Event::seq).max().unwrap_or_default();
let mut recorder = Recorder::resume_at(
invocation.loop_id.clone(),
Arc::clone(invocation.store),
head,
)
.with_visibility(
retirement_run.clone(),
Arc::clone(invocation.visibility_store),
);
recorder
.record_workloop_retirement_generation(
Utc::now(),
generation.carry.clone(),
generation.run_id.clone(),
start,
)
.await?;
Ok((retirement_run, recorder))
}
fn live_handle(invocation: &RetireInvocation<'_>) -> Result<Option<WorkflowHandle>, WorkloopError> {
Ok(invocation
.registry
.list()
.map_err(|error| WorkloopError::Engine {
reason: format!(
"listing the registry to find {}'s live writer failed: {error}",
invocation.loop_id
),
})?
.into_iter()
.find(|handle| handle.workflow_id() == invocation.loop_id))
}
fn publish_retire_body(
invocation: &RetireInvocation<'_>,
retirement_run: &RunId,
pid: crate::Pid,
generation: &CurrentGeneration,
recorder: Recorder,
) -> Result<(), WorkloopError> {
let handle = WorkflowHandle::new(WorkflowHandleParts {
workflow_id: invocation.loop_id.clone(),
run_id: retirement_run.clone(),
pid,
workflow_type: generation.workflow_type.clone(),
namespace: String::from("default"),
loaded_version: generation.loaded_version.clone(),
cached_status: WorkflowStatus::Running,
residency: HandleResidency::Resident,
recorder,
completion: CompletionNotifier::new(),
});
invocation
.registry
.insert_sole_workflow_writer((invocation.loop_id.clone(), retirement_run.clone()), handle)
.map_err(|error| WorkloopError::Engine {
reason: format!("publishing the retire body's handle failed: {error}"),
})
}
fn cancel_orphaned_body(
runtime: &Arc<RuntimeHandle>,
loop_id: &WorkflowId,
pid: crate::Pid,
cause: &WorkloopError,
) {
if let Err(error) = runtime.cancel_pid(pid) {
tracing::error!(
%loop_id,
pid,
cause = %cause,
error = %error,
"the retire body was spawned but could neither be published nor cancelled; it is \
running unmonitored on the loop's module and nothing will observe its exit"
);
}
}
async fn await_retire_body(
runtime: &Arc<RuntimeHandle>,
loop_id: &WorkflowId,
pid: crate::Pid,
) -> Result<(), WorkloopError> {
let (sender, receiver) = tokio::sync::oneshot::channel();
if let Err(error) = runtime.monitor_process(pid, move |outcome| {
let _ = sender.send(outcome);
}) {
let failure = WorkloopError::Engine {
reason: format!("monitoring the retire body's process {pid} failed: {error}"),
};
cancel_orphaned_body(runtime, loop_id, pid, &failure);
return Err(failure);
}
let outcome = receiver.await.map_err(|_| WorkloopError::Engine {
reason: format!(
"the retire body's process {pid} exit was never reported; refusing to record a \
retirement whose cleanup cannot be shown to have completed"
),
})?;
let outcome = outcome.map_err(|error| WorkloopError::Engine {
reason: format!("observing the retire body's exit failed: {error}"),
})?;
classify(pid, &outcome)
}
fn classify(
pid: crate::Pid,
outcome: &crate::runtime::outcome::WorkflowProcessOutcome,
) -> Result<(), WorkloopError> {
use crate::runtime::outcome::WorkflowProcessOutcome;
match outcome {
WorkflowProcessOutcome::Completed(_) => Ok(()),
WorkflowProcessOutcome::Failed(error) => Err(WorkloopError::Engine {
reason: format!(
"the retire body (process {pid}) failed: {message}. No `LoopRetired` terminal \
is recorded — that would claim a cleanup which did not finish — so the loop is \
NOT retired and stays registered on the sweep set. What the body did before \
failing IS in the loop's history, inside the retirement generation opened for \
it; a further retirement attempt opens a fresh generation and runs the body \
again from the top, so any effect it already performed will be performed twice \
unless the body is idempotent",
message = error.message
),
}),
}
}
fn current_generation(
invocation: &RetireInvocation<'_>,
history: &[Event],
) -> Result<CurrentGeneration, WorkloopError> {
let loop_id = invocation.loop_id;
let (run_id, workflow_type, package_version, carry) = history
.iter()
.rev()
.find_map(|event| match event {
Event::WorkflowStarted {
run_id,
workflow_type,
package_version,
input,
..
} => Some((
run_id.clone(),
workflow_type.clone(),
package_version.clone(),
input.clone(),
)),
_ => None,
})
.ok_or_else(|| WorkloopError::Engine {
reason: format!("workloop {loop_id} has no recorded generation to retire"),
})?;
let loaded_version = crate::loader::parse_package_version(&workflow_type, &package_version)
.map_err(|error| WorkloopError::Engine {
reason: format!("resolving the retiring loop's package version failed: {error}"),
})?;
let loaded = invocation
.catalog
.get(&workflow_type, &loaded_version)
.map_err(|error| WorkloopError::Engine {
reason: format!("resolving the retiring loop's package failed: {error}"),
})?
.ok_or_else(|| WorkloopError::Engine {
reason: format!(
"workloop {loop_id} is pinned to package version {loaded_version} of \
`{workflow_type}`, which is not loaded on this engine, so its retire body \
cannot be reached"
),
})?;
let module = loaded.deployed_entry_module().to_owned();
Ok(CurrentGeneration {
run_id,
workflow_type,
package_version,
loaded_version,
module,
carry,
})
}