use std::sync::Arc;
use aion_core::{Event, RunId, WorkflowId};
use aion_store::EventStore;
use aion_store::visibility::VisibilityStore;
use aion_store::workloop::WorkloopStore;
use chrono::Utc;
use super::iteration::{self, WorkloopIterationClose};
use super::service::{WorkloopService, window_context};
use crate::durability::Recorder;
use crate::error::EngineError;
use crate::registry::{Registry, TerminalOutcome};
#[derive(Clone)]
pub struct IterationCloseContext {
pub workloop_store: Arc<dyn WorkloopStore>,
pub service: Arc<WorkloopService>,
pub store: Arc<dyn EventStore>,
pub visibility_store: Arc<dyn VisibilityStore>,
pub registry: Arc<Registry>,
}
pub async fn close_iteration(
context: &IterationCloseContext,
loop_id: &WorkflowId,
close: WorkloopIterationClose,
) -> Result<RunId, EngineError> {
let record = context
.workloop_store
.get_workloop(loop_id)
.await
.map_err(EngineError::from)?
.ok_or_else(|| EngineError::InvalidState {
reason: format!("workflow {loop_id} is not a registered workloop"),
})?;
let window_seq = window_context(&record);
let spec = record.spec.clone();
let retention = record.spec.retention();
let carry_for_notify = close.carry.clone();
let close_for_recorder = close;
let outcome = with_loop_recorder(context, loop_id, move |recorder, history| {
let spec = spec.clone();
let close = close_for_recorder.clone();
let loop_id = loop_id.clone();
Box::pin(async move {
let run_id = active_run_id(history).ok_or_else(|| {
crate::durability::DurabilityError::HistoryShape {
reason: format!("workloop {loop_id} has no recorded generation"),
}
})?;
iteration::close_iteration(
recorder,
iteration::IterationContext {
history,
run_id: &run_id,
loop_id: &loop_id,
spec: &spec,
window_seq,
recorded_at: Utc::now(),
},
close,
)
.await
.map_err(|error| crate::durability::DurabilityError::HistoryShape {
reason: error.to_string(),
})
})
})
.await?;
if let Some(handle) = registry_handle(context, loop_id)? {
let closed_run = handle.run_id().clone();
handle.completion().notify(TerminalOutcome::ContinuedAsNew {
input: carry_for_notify,
workflow_type: None,
parent_run_id: closed_run.clone(),
});
context.registry.remove(loop_id, &closed_run)?;
}
let cutoff = retention_cutoff(Utc::now(), retention)?;
for state_record in outcome.records.clone() {
context
.workloop_store
.put_invariant_record(state_record, cutoff)
.await
.map_err(EngineError::from)?;
}
context
.service
.note_iteration_closed(loop_id, &outcome.samples)
.await
.map_err(EngineError::from)?;
Ok(outcome.next_run_id)
}
fn retention_cutoff(
now: chrono::DateTime<Utc>,
retention: std::time::Duration,
) -> Result<chrono::DateTime<Utc>, EngineError> {
let window =
chrono::Duration::from_std(retention).map_err(|error| EngineError::InvalidState {
reason: format!(
"workloop retention window of {seconds}s cannot be expressed as a calendar \
duration ({error}), so no retention cutoff exists; refusing rather than \
pruning against a fallback that would delete every prior generation",
seconds = retention.as_secs()
),
})?;
now.checked_sub_signed(window)
.ok_or_else(|| EngineError::InvalidState {
reason: format!(
"subtracting the declared workloop retention window of {seconds}s from \
{now} left the representable calendar range, so no retention cutoff exists",
seconds = retention.as_secs()
),
})
}
fn registry_handle(
context: &IterationCloseContext,
loop_id: &WorkflowId,
) -> Result<Option<crate::registry::WorkflowHandle>, EngineError> {
context.registry.sole_handle(loop_id)
}
async fn with_loop_recorder<T>(
context: &IterationCloseContext,
loop_id: &WorkflowId,
record: impl for<'a> FnOnce(
&'a mut Recorder,
&'a [Event],
) -> std::pin::Pin<
Box<
dyn std::future::Future<Output = Result<T, crate::durability::DurabilityError>>
+ Send
+ 'a,
>,
>,
) -> Result<T, EngineError> {
if let Some(handle) = registry_handle(context, loop_id)? {
let recorder = handle.recorder();
let mut recorder = recorder.lock().await;
let history = context.store.read_history(loop_id).await?;
let value = record(&mut recorder, &history).await?;
return Ok(value);
}
let history = context.store.read_history(loop_id).await?;
let head = history.iter().map(Event::seq).max().unwrap_or_default();
let mut recorder = Recorder::resume_at(loop_id.clone(), Arc::clone(&context.store), head);
if let Some(run_id) = active_run_id(&history) {
recorder = recorder.with_visibility(run_id, Arc::clone(&context.visibility_store));
}
let value = record(&mut recorder, &history).await?;
Ok(value)
}
pub(crate) fn active_run_id(history: &[Event]) -> Option<RunId> {
history.iter().rev().find_map(|event| match event {
Event::WorkflowStarted { run_id, .. } => Some(run_id.clone()),
_ => None,
})
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use chrono::TimeZone;
use super::retention_cutoff;
use crate::error::EngineError;
fn now() -> Result<chrono::DateTime<chrono::Utc>, Box<dyn std::error::Error>> {
chrono::Utc
.with_ymd_and_hms(2026, 8, 26, 12, 0, 0)
.single()
.ok_or_else(|| "test instant must be valid".into())
}
#[test]
fn a_declared_window_subtracts_to_its_own_past() -> Result<(), Box<dyn std::error::Error>> {
let now = now()?;
let cutoff = retention_cutoff(now, Duration::from_secs(14 * 86_400))?;
assert_eq!(cutoff, now - chrono::Duration::days(14));
Ok(())
}
#[test]
fn an_unrepresentable_window_refuses_instead_of_pruning_everything()
-> Result<(), Box<dyn std::error::Error>> {
let now = now()?;
let outcome = retention_cutoff(now, Duration::from_secs(u64::MAX / 1_000));
assert!(
!matches!(&outcome, Ok(cutoff) if *cutoff == now),
"an unrepresentable retention must never yield a cutoff of `now`: that prunes \
every prior generation, which is the opposite of what it declares"
);
let refusal = outcome
.err()
.ok_or("an unrepresentable retention window must be refused")?;
assert!(
matches!(&refusal, EngineError::InvalidState { reason }
if reason.contains("retention")),
"the refusal must name the retention window: {refusal}"
);
Ok(())
}
#[test]
fn a_window_that_underflows_the_calendar_refuses() {
let early = chrono::DateTime::<chrono::Utc>::MIN_UTC + chrono::Duration::days(1);
let outcome = retention_cutoff(early, Duration::from_secs(1_000 * 365 * 86_400));
assert!(
outcome.is_err(),
"subtracting past the representable range must refuse, not saturate: {outcome:?}"
);
}
}