aion-rs 0.31.0

//! Continue-as-new lifecycle transition.

use std::collections::HashSet;
use std::sync::Arc;

use aion_core::{ActivityId, Event, Payload, RunId, SearchAttributeSchema, WorkflowId};
use aion_store::EventStore;
use aion_store::visibility::VisibilityStore;

use crate::EngineError;
use crate::lifecycle::continuation::{
    self, ContinuationOrigin, ContinuationOutcome, ContinuationRequest,
};
use crate::lifecycle::start::StartWorkflowContext;
use crate::loader::WorkflowCatalog;
use crate::registry::{Registry, TerminalOutcome, WorkflowHandle};
use crate::runtime::RuntimeHandle;
use crate::supervision::SupervisionTree;

/// Dependencies required to continue a workflow as a new run.
pub struct ContinueAsNewContext<'a> {
    /// Durable event store used to scan history and start the replacement run.
    pub store: Arc<dyn EventStore>,
    /// Visibility store for workflow visibility projections.
    pub visibility_store: Arc<dyn VisibilityStore>,
    /// Shared workflow catalog resolving types to loaded package versions.
    pub catalog: Arc<WorkflowCatalog>,
    /// Runtime boundary used to spawn the replacement workflow process.
    pub runtime: &'a Arc<RuntimeHandle>,
    /// Structural supervision tree recording the per-type supervisor placement.
    pub supervision: Arc<SupervisionTree>,
    /// Active execution registry keyed by workflow/run identifiers.
    pub registry: &'a Arc<Registry>,
    /// Schema validating initial search attributes on the replacement run.
    pub search_attribute_schema: Arc<SearchAttributeSchema>,
}

/// Request payload carried into the replacement run.
#[derive(Clone, Debug, PartialEq)]
pub struct ContinueAsNewRequest {
    /// Opaque workflow input payload for the replacement run.
    pub input: Payload,
    /// Optional workflow type override for the replacement run.
    pub workflow_type: Option<String>,
}

/// Continues a live workflow run as a new run under the same workflow id.
///
/// The transition is ONE durable batch through the predecessor's own recorder
/// — see [`crate::lifecycle::continuation`] for why a second recorder is not
/// an option — and the registry's single handle for the workflow is re-keyed
/// onto the successor run before the recorder lock drops.
///
/// # Errors
///
/// Returns [`EngineError::WorkflowNotFound`] when the `(workflow, run)` pair is
/// not registered. Returns [`EngineError::Runtime`] when pending activities or
/// child workflows remain unresolved, when the run already recorded a terminal,
/// or when the requested replacement type is not the run's own. Recorder,
/// runtime, supervision, and registry failures surface as their typed
/// [`EngineError`] variants.
pub async fn continue_as_new(
    context: ContinueAsNewContext<'_>,
    id: &WorkflowId,
    run: &RunId,
    request: ContinueAsNewRequest,
) -> Result<WorkflowHandle, EngineError> {
    let handle = registered_handle(context.registry, id, run)?;

    let workflow_type = request
        .workflow_type
        .as_deref()
        .unwrap_or(handle.workflow_type());
    if workflow_type != handle.workflow_type() {
        return Err(EngineError::Runtime {
            reason: format!(
                "continue_as_new must restart the same workflow type: current={}, requested={workflow_type}",
                handle.workflow_type()
            ),
        });
    }
    validate_replacement_workflow_type(&context.catalog, workflow_type)?;
    let workflow_type = workflow_type.to_owned();

    let outcome = continuation::open_successor_generation(
        &StartWorkflowContext {
            store: context.store,
            visibility_store: context.visibility_store,
            catalog: Arc::clone(&context.catalog),
            runtime: Arc::clone(context.runtime),
            supervision: context.supervision,
            registry: Arc::clone(context.registry),
            signal_handoff: None,
            search_attribute_schema: context.search_attribute_schema,
            monitor_tokio_handle: tokio::runtime::Handle::current(),
        },
        &handle,
        ContinuationRequest {
            predecessor_run: run.clone(),
            origin: ContinuationOrigin::RecordsTheTerminal {
                input: request.input.clone(),
                workflow_type: request.workflow_type.clone(),
            },
            workflow_type,
            input: request.input.clone(),
        },
    )
    .await?;

    let new_handle = match outcome {
        ContinuationOutcome::Opened(handle) => *handle,
        // Only reachable if the workflow's own `continue_as_new` NIF (or the
        // exit monitor acting on it) opened a successor for this run between
        // this caller's registry read and the transition's lock. The run DID
        // continue — just not by this call — and the operator asked for a
        // continuation, not for this particular one to be the winner. Reported
        // rather than silently returning the incumbent handle, because the
        // caller is owed the successor's identity and this path does not know
        // it without a second lookup that could race again.
        ContinuationOutcome::AlreadyOpen => {
            return Err(EngineError::Runtime {
                reason: format!(
                    "continue_as_new rejected: workflow {id} run {run} already has a successor \
                     generation, started by the run\'s own continue_as_new"
                ),
            });
        }
    };

    // The doorbell rings AFTER the transition is durable and published: a
    // result waiter released earlier would observe the predecessor as
    // continued while the successor was not yet resolvable.
    handle.completion().notify(TerminalOutcome::ContinuedAsNew {
        input: request.input,
        workflow_type: request.workflow_type,
        parent_run_id: run.clone(),
    });

    Ok(new_handle)
}

pub(crate) fn guard_no_pending_work(events: &[Event]) -> Result<(), EngineError> {
    let mut pending_activities = HashSet::<ActivityId>::new();
    let mut pending_children = HashSet::<WorkflowId>::new();

    for event in events {
        match event {
            Event::ActivityScheduled { activity_id, .. }
            | Event::ActivityStarted { activity_id, .. }
            // A leased attempt is a worker holding the work: still pending.
            | Event::ActivityLeased { activity_id, .. }
            | Event::ActivityAdoptionOffered { activity_id, .. } => {
                pending_activities.insert(activity_id.clone());
            }
            Event::ActivityCompleted { activity_id, .. }
            | Event::ActivityFailed { activity_id, .. }
            | Event::ActivityCancelled { activity_id, .. } => {
                pending_activities.remove(activity_id);
            }
            Event::ChildWorkflowStarted {
                child_workflow_id, ..
            } => {
                pending_children.insert(child_workflow_id.clone());
            }
            Event::ChildWorkflowCompleted {
                child_workflow_id, ..
            }
            | Event::ChildWorkflowFailed {
                child_workflow_id, ..
            }
            | Event::ChildWorkflowCancelled {
                child_workflow_id, ..
            } => {
                pending_children.remove(child_workflow_id);
            }
            Event::WorkflowStarted { .. }
            | Event::WorkflowCompleted { .. }
            | Event::WorkflowFailed { .. }
            | Event::WorkflowCancelled { .. }
            | Event::WorkflowTimedOut { .. }
            | Event::WorkflowContinuedAsNew { .. }
            | Event::WorkflowReopened { .. }
            | Event::WorkflowPaused { .. }
            | Event::WorkflowResumed { .. }
            | Event::SearchAttributesUpdated { .. }
            // A warning, not a settlement: the advisory activity's own
            // terminal ActivityFailed is what clears it from pending work.
            | Event::ActivityAdvisoryExhausted { .. }
            | Event::ActivityFallbackRouted { .. }
            | Event::TimerStarted { .. }
            | Event::TimerFired { .. }
            | Event::TimerCancelled { .. }
            | Event::WithTimeoutCompleted { .. }
            | Event::SignalReceived { .. }
            | Event::SignalSent { .. }
            | Event::ScheduleCreated { .. }
            | Event::ScheduleUpdated { .. }
            | Event::SchedulePaused { .. }
            | Event::ScheduleResumed { .. }
            | Event::ScheduleDeleted { .. }
            | Event::ScheduleTriggered { .. }
            // Workloop bookkeeping is never pending work: a hatched workflow
            // is DETACHED (no awaited terminal ties it to this run), and
            // cadence/iteration/retirement/alarm records settle nothing.
            | Event::CadenceFired { .. }
            | Event::IterationClosed { .. }
            | Event::LoopRetired { .. }
            | Event::WorkflowHatched { .. }
            | Event::InvariantUnconfirmed { .. } => {}
        }
    }

    if pending_activities.is_empty() && pending_children.is_empty() {
        return Ok(());
    }

    Err(EngineError::Runtime {
        reason: format!(
            "cannot continue as new while pending work exists: {} activities ({:?}), {} child workflows ({:?})",
            pending_activities.len(),
            pending_activities,
            pending_children.len(),
            pending_children
        ),
    })
}

fn registered_handle(
    registry: &Registry,
    id: &WorkflowId,
    run: &RunId,
) -> Result<WorkflowHandle, EngineError> {
    registry
        .get(id, run)?
        .ok_or_else(|| EngineError::WorkflowNotFound {
            workflow_type: format!("{id}/{run}"),
        })
}

fn validate_replacement_workflow_type(
    catalog: &WorkflowCatalog,
    workflow_type: &str,
) -> Result<(), EngineError> {
    catalog
        .routed(workflow_type)?
        .ok_or_else(|| EngineError::WorkflowNotFound {
            workflow_type: workflow_type.to_owned(),
        })
        .map(|_| ())
}

#[cfg(test)]
#[path = "continue_as_new_tests.rs"]
mod tests;