aion-rs 0.31.0

//! Concrete delegated signal router: record `SignalReceived`, then deliver to the mailbox.

use std::sync::Arc;

use aion_core::Payload;
use async_trait::async_trait;
use chrono::Utc;

use crate::runtime::UndeliveredWake;
use crate::{
    EngineError, HandleResidency, RuntimeHandle, SignalRouterError, WorkflowHandle,
    engine::delegated, signal::SignalResumeHandoff,
};

/// Delegated signal router for resident workflow processes.
///
/// Signals are first recorded through the target handle's single-writer recorder.
/// Only after that durable append succeeds does the router enqueue the signal
/// marker into the target runtime mailbox, preserving record-before-deliver
/// crash-safety.
#[derive(Clone)]
pub struct ConcreteSignalRouter {
    runtime: Arc<RuntimeHandle>,
    handoff: Arc<SignalResumeHandoff>,
}

impl ConcreteSignalRouter {
    /// Create a router that delivers recorded signals through `runtime` and defers through `handoff`.
    #[must_use]
    pub fn new(runtime: Arc<RuntimeHandle>, handoff: Arc<SignalResumeHandoff>) -> Self {
        Self { runtime, handoff }
    }
}

#[async_trait]
impl delegated::SignalRouter for ConcreteSignalRouter {
    async fn route(
        &self,
        target: &WorkflowHandle,
        name: String,
        payload: Payload,
    ) -> Result<(), EngineError> {
        let recorder = target.recorder();
        {
            let mut recorder = recorder.lock().await;
            // Terminal check and signal record are atomic under the recorder
            // lock: the exit monitor records terminal events through the same
            // recorder, and a terminal run must reject signals instead of
            // appending after its terminal event or deferring to a resume
            // queue that can never drain.
            let history = recorder.read_history().await.map_err(EngineError::from)?;
            if crate::engine::delegated::run_has_terminal_history(&history, target.run_id()) {
                return Err(SignalRouterError::Terminal {
                    workflow_id: target.workflow_id().clone(),
                    run_id: target.run_id().clone(),
                }
                .into());
            }
            recorder
                .record_signal_received(Utc::now(), name.clone(), payload.clone())
                .await?;
        }

        match target.residency() {
            HandleResidency::Resident => {
                let Err(error) = self
                    .runtime
                    .deliver_signal_received_async(target.pid())
                    .await
                else {
                    return Ok(());
                };
                // The signal is already durable; the marker is only a wake.
                // A process that completed between the record and this
                // delivery (completion racing the signal — including a surplus
                // wake letting the workflow resolve the just-recorded signal
                // from history and return before the marker lands) has
                // accepted the signal: its terminal is recorded by the exit
                // monitor, and a crashed run replays with the signal in
                // history. Only a RECORDED ending excuses the undelivered
                // marker; a pid that is merely absent from the scheduler's
                // process table is not one, so the classifier waits an exit
                // in flight out rather than guessing at it. That wait parks a
                // thread, so it is handed to the blocking pool instead of the
                // executor worker this task is running on.
                let delivery_reason = error.to_string();
                let classification = {
                    let runtime = Arc::clone(&self.runtime);
                    let pid = target.pid();
                    tokio::task::spawn_blocking(move || runtime.classify_undelivered_wake(pid))
                        .await
                };
                let reason = match classification {
                    Ok(Ok(UndeliveredWake::ProcessEnded)) => {
                        tracing::debug!(
                            workflow_id = %target.workflow_id(),
                            run_id = %target.run_id(),
                            signal_name = %name,
                            process = target.pid(),
                            "recorded signal raced process exit; durable record stands"
                        );
                        return Ok(());
                    }
                    Ok(Ok(UndeliveredWake::ProcessDidNotEnd { .. })) => delivery_reason,
                    // The exit was in flight and published nothing inside the
                    // readiness window. Name both facts: the caller is being
                    // told a durable signal was not woken, and only the pair
                    // explains why the engine could not settle it either way.
                    Ok(Ok(UndeliveredWake::ExitInFlight)) => format!(
                        "{delivery_reason}; the target left the scheduler process table with its \
                         exit still in flight, and no terminal was published on its exit record \
                         within the readiness window"
                    ),
                    // An unreadable exit registry cannot excuse an undelivered
                    // wake. Report the refusal and why it could not be
                    // classified together: either alone hides half the fault.
                    Ok(Err(classification)) => format!(
                        "{delivery_reason}; the target's ending could not be classified: \
                         {classification}"
                    ),
                    // The classification task itself did not run to an answer;
                    // that is a fault of this engine, not of the target.
                    Err(join) => format!(
                        "{delivery_reason}; the target's ending could not be classified because \
                         the classification task failed: {join}"
                    ),
                };
                tracing::warn!(
                    workflow_id = %target.workflow_id(),
                    run_id = %target.run_id(),
                    signal_name = %name,
                    process = target.pid(),
                    error = %reason,
                    "durably recorded signal could not be delivered to resident workflow mailbox"
                );
                Err(EngineError::from(SignalRouterError::DeliveryFailed {
                    workflow_id: target.workflow_id().clone(),
                    run_id: target.run_id().clone(),
                    process_id: target.pid(),
                    signal_name: name,
                    reason,
                }))
            }
            HandleResidency::Suspended => self
                .handoff
                .defer(target.workflow_id().clone(), name, payload)
                .map_err(|error| {
                    EngineError::from(SignalRouterError::Handoff {
                        reason: error.to_string(),
                    })
                }),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::time::{Duration, Instant};

    use aion_core::{Event, Payload, WorkflowStatus};
    use aion_package::ContentHash;
    use aion_store::{EventStore, InMemoryStore};
    use serde_json::json;

    use super::ConcreteSignalRouter;
    use crate::durability::Recorder;
    use crate::engine::delegated::SignalRouter;
    use crate::registry::{
        CompletionNotifier, HandleResidency, WorkflowHandle, WorkflowHandleParts,
    };
    use crate::runtime::{RuntimeConfig, RuntimeHandle, SignalDeliveryConfig, UndeliveredWake};
    use crate::signal::SignalResumeHandoff;
    use crate::{EngineError, SignalRouterError};

    type TestResult = Result<(), Box<dyn std::error::Error>>;

    fn payload(label: &str) -> Result<Payload, aion_core::PayloadError> {
        Payload::from_json(&json!({ "label": label }))
    }

    async fn started_workflow_handle(
        store: &Arc<dyn EventStore>,
        pid: u64,
    ) -> Result<WorkflowHandle, Box<dyn std::error::Error>> {
        let workflow_id = aion_core::WorkflowId::new_v4();
        let run_id = aion_core::RunId::new_v4();
        let mut recorder = Recorder::new(workflow_id.clone(), Arc::clone(store));
        recorder
            .record_workflow_started(
                chrono::Utc::now(),
                crate::durability::WorkflowStartRecord {
                    workflow_type: "checkout".to_owned(),
                    input: payload("input")?,
                    run_id: run_id.clone(),
                    parent_run_id: None,
                    parent_workflow_id: None,
                    package_version: aion_core::PackageVersion::new("a".repeat(64)),
                },
            )
            .await?;
        Ok(WorkflowHandle::new(WorkflowHandleParts {
            workflow_id,
            run_id,
            pid,
            workflow_type: "checkout".to_owned(),
            namespace: String::from("default"),
            loaded_version: ContentHash::from_bytes([3; 32]),
            cached_status: WorkflowStatus::Running,
            residency: HandleResidency::Resident,
            recorder,
            completion: CompletionNotifier::new(),
        }))
    }

    /// Build a runtime whose wake-readiness window is `ready_timeout`.
    ///
    /// That window is the bound the undelivered-wake classifier parks an exit
    /// in flight under, so these pins set it rather than inheriting a default
    /// they would then be measuring against.
    fn runtime_with_ready_timeout(
        ready_timeout: Duration,
    ) -> Result<Arc<RuntimeHandle>, Box<dyn std::error::Error>> {
        let delivery = SignalDeliveryConfig::new(
            ready_timeout,
            1,
            Duration::from_millis(1),
            Duration::from_millis(1),
        );
        Ok(Arc::new(RuntimeHandle::new(
            RuntimeConfig::new(Some(1), crate::runtime::config::TEST_STOP_DRAIN_TIMEOUT)
                .with_signal_delivery(delivery),
        )?))
    }

    fn test_runtime() -> Result<Arc<RuntimeHandle>, Box<dyn std::error::Error>> {
        runtime_with_ready_timeout(Duration::from_millis(50))
    }

    fn router_over(runtime: &Arc<RuntimeHandle>) -> ConcreteSignalRouter {
        ConcreteSignalRouter::new(Arc::clone(runtime), Arc::new(SignalResumeHandoff::new()))
    }

    fn router() -> Result<ConcreteSignalRouter, Box<dyn std::error::Error>> {
        Ok(router_over(&test_runtime()?))
    }

    /// How long a test waits for the exit drainer to publish a terminal.
    const RECORDED_ENDING_TIMEOUT: Duration = Duration::from_secs(10);

    /// Park until the engine holds a record that `pid` ended.
    ///
    /// The drainer publishes the terminal on its own thread, so the wait is
    /// the record's publication signal; it never asserts anything about the
    /// scheduler's process table.
    async fn wait_for_recorded_ending(runtime: &Arc<RuntimeHandle>, pid: u64) -> TestResult {
        let deadline = Instant::now() + RECORDED_ENDING_TIMEOUT;
        while Instant::now() < deadline {
            if runtime.process_ending_recorded(pid)? {
                return Ok(());
            }
            tokio::time::sleep(Duration::from_millis(5)).await;
        }
        Err("the exit registry never recorded the process' ending".into())
    }

    /// Blocking twin of [`wait_for_recorded_ending`] for the synchronous pins.
    fn blocking_wait_for_recorded_ending(runtime: &Arc<RuntimeHandle>, pid: u64) -> TestResult {
        let deadline = Instant::now() + RECORDED_ENDING_TIMEOUT;
        while Instant::now() < deadline {
            if runtime.process_ending_recorded(pid)? {
                return Ok(());
            }
            std::thread::sleep(Duration::from_millis(5));
        }
        Err("the exit registry never recorded the process' ending".into())
    }

    fn recorded_signal(history: &[Event]) -> Option<(&String, &Payload)> {
        history.iter().find_map(|event| match event {
            Event::SignalReceived { name, payload, .. } => Some((name, payload)),
            _ => None,
        })
    }

    /// A resident process that exited between the durable record and the
    /// marker delivery has accepted the signal: the record is the contract,
    /// the marker is only a wake. Before the fix this returned
    /// `DeliveryFailed` for an already-accepted signal, which surfaced as a
    /// flaky engine error whenever a completion (or a surplus wake letting
    /// the workflow resolve the just-recorded signal from history) raced the
    /// delivery.
    #[tokio::test]
    async fn a_recorded_ending_resolves_an_undelivered_signal_as_accepted() -> TestResult {
        let store: Arc<dyn EventStore> = Arc::new(InMemoryStore::default());
        let runtime = test_runtime()?;
        let pid = runtime.spawn_test_process()?;
        runtime.cancel_pid(pid)?;
        wait_for_recorded_ending(&runtime, pid).await?;
        let handle = started_workflow_handle(&store, pid).await?;
        let sent = payload("recorded")?;

        router_over(&runtime)
            .route(&handle, "wake".to_owned(), sent.clone())
            .await?;

        let history = store.read_history(handle.workflow_id()).await?;
        let recorded = recorded_signal(&history).ok_or("SignalReceived was not recorded")?;
        assert_eq!(recorded.0, "wake");
        assert_eq!(recorded.1, &sent);
        runtime.shutdown()?;
        Ok(())
    }

    /// A pid this runtime never spawned carries no record of an ending: it has
    /// no cleanup tombstone and no terminal in the exit registry. Its absence
    /// from the scheduler's process table is not an ending, so an undelivered
    /// wake marker must be reported as the delivery failure it is — while the
    /// already-durable `SignalReceived` still stands. Reading absence as an
    /// ending answered this route `Ok(())`, which told the caller a signal had
    /// been delivered to a process that never received it.
    #[tokio::test]
    async fn an_absent_pid_without_a_recorded_ending_is_a_delivery_failure() -> TestResult {
        let store: Arc<dyn EventStore> = Arc::new(InMemoryStore::default());
        let runtime = test_runtime()?;
        let pid = 424_242;
        assert!(
            !runtime.is_live(pid),
            "fixture control: the pid must be absent from the process table"
        );
        assert!(
            !runtime.process_ending_recorded(pid)?,
            "fixture control: nothing may record an ending for a pid never spawned here"
        );
        let handle = started_workflow_handle(&store, pid).await?;
        let sent = payload("recorded")?;

        let error = router_over(&runtime)
            .route(&handle, "wake".to_owned(), sent.clone())
            .await
            .err()
            .ok_or("an undeliverable wake to an absent pid reported success")?;

        assert!(
            matches!(
                error,
                EngineError::SignalRouter(SignalRouterError::DeliveryFailed {
                    process_id,
                    ref signal_name,
                    ..
                }) if process_id == pid && signal_name.as_str() == "wake"
            ),
            "an undelivered wake with no recorded ending must be DeliveryFailed, got: {error:?}"
        );

        let history = store.read_history(handle.workflow_id()).await?;
        let recorded = recorded_signal(&history).ok_or("SignalReceived was not recorded")?;
        assert_eq!(recorded.0, "wake");
        assert_eq!(recorded.1, &sent);
        runtime.shutdown()?;
        Ok(())
    }

    /// The exit registry may publish the terminal before the asynchronous Aion
    /// cleanup callback stamps its tombstone. That published terminal is a
    /// record of an ending on its own, and the classifier must read it without
    /// consulting the process table.
    #[test]
    fn a_recorded_terminal_before_cleanup_classifies_the_wake_as_completion() -> TestResult {
        let runtime = test_runtime()?;
        let pid = runtime.spawn_test_process()?;
        runtime.cancel_pid(pid)?;
        blocking_wait_for_recorded_ending(&runtime, pid)?;
        assert!(
            !runtime.process_cleanup_started(pid),
            "fixture control: Aion cleanup cannot have started without a monitor"
        );
        assert_eq!(
            runtime.classify_undelivered_wake(pid)?,
            UndeliveredWake::ProcessEnded,
            "the exit registry's published terminal is a recorded ending"
        );
        runtime.shutdown()?;
        Ok(())
    }

    /// The control for the arm above at the classifier itself: a pid the
    /// registry never registered is absent from the process table for a reason
    /// that has nothing to do with a workflow ending, and there is no record
    /// in flight to wait for either.
    #[test]
    fn an_absent_pid_without_a_recorded_ending_is_not_completion() -> TestResult {
        let runtime = test_runtime()?;
        let pid = 424_242;
        assert!(!runtime.is_live(pid), "fixture control: the pid is absent");
        assert!(
            !runtime.process_cleanup_started(pid),
            "fixture control: no cleanup tombstone exists for a pid never spawned here"
        );

        assert_eq!(
            runtime.classify_undelivered_wake(pid)?,
            UndeliveredWake::ProcessDidNotEnd { in_table: false },
            "absence from the process table alone must never classify as completion"
        );
        runtime.shutdown()?;
        Ok(())
    }

    /// An exit that is still in flight past the readiness window is reported,
    /// not guessed at: the caller is told a durable signal was not woken, and
    /// the reason carries both the refusal and why the engine could not settle
    /// it — the pid left the process table and nothing published on its record.
    #[tokio::test]
    async fn an_exit_in_flight_past_the_window_is_a_named_delivery_failure() -> TestResult {
        let store: Arc<dyn EventStore> = Arc::new(InMemoryStore::default());
        let runtime = runtime_with_ready_timeout(Duration::from_millis(20))?;
        runtime.pause_exit_drainer_for_test(RECORDED_ENDING_TIMEOUT)?;
        let pid = runtime.spawn_test_process()?;
        runtime.cancel_pid(pid)?;
        let deadline = Instant::now() + RECORDED_ENDING_TIMEOUT;
        while runtime.is_live(pid) {
            if Instant::now() >= deadline {
                return Err(format!("pid {pid} never left the scheduler process table").into());
            }
            tokio::time::sleep(Duration::from_millis(1)).await;
        }
        let handle = started_workflow_handle(&store, pid).await?;

        let routed = router_over(&runtime)
            .route(&handle, "wake".to_owned(), payload("recorded")?)
            .await;
        runtime.release_exit_drainer_for_test();

        let error = routed
            .err()
            .ok_or("an exit in flight past the window reported a delivered wake")?;
        let EngineError::SignalRouter(SignalRouterError::DeliveryFailed { reason, .. }) = error
        else {
            return Err(format!("expected a named DeliveryFailed, got: {error:?}").into());
        };
        assert!(
            reason.contains("exit still in flight")
                && reason.contains("within the readiness window"),
            "the failure must name both facts it rests on, got: {reason}"
        );
        runtime.shutdown()?;
        Ok(())
    }

    #[tokio::test]
    async fn terminal_run_rejects_signal_without_appending() -> TestResult {
        let store: Arc<dyn EventStore> = Arc::new(InMemoryStore::default());
        let handle = started_workflow_handle(&store, 424_243).await?;
        {
            let recorder = handle.recorder();
            let mut recorder = recorder.lock().await;
            recorder
                .record_workflow_failed(
                    chrono::Utc::now(),
                    aion_core::WorkflowError {
                        message: "killed".to_owned(),
                        details: None,
                    },
                )
                .await?;
        }
        let terminal_len = store.read_history(handle.workflow_id()).await?.len();

        let error = router()?
            .route(&handle, "wake".to_owned(), payload("rejected")?)
            .await
            .err()
            .ok_or("signal to terminal run unexpectedly succeeded")?;

        assert!(matches!(
            error,
            EngineError::SignalRouter(SignalRouterError::Terminal { .. })
        ));
        assert_eq!(
            store.read_history(handle.workflow_id()).await?.len(),
            terminal_len
        );
        Ok(())
    }

    /// A workflow that completed with its signal already durable is not an
    /// engine fault. Exit cleanup stamps its tombstone before beamr retires
    /// the pid, so the wake-marker failure arm must read cleanup, not
    /// scheduler liveness alone — judging on `is_live` by itself reported
    /// `DeliveryFailed` for a run that simply finished.
    ///
    /// The low-level beamr enqueue has no deterministic live-pid refusal
    /// seam, so this pins the classifier the failure arm consults, at the
    /// exact interleaving cleanup forces.
    #[test]
    fn signal_to_live_pid_with_cleanup_started_is_not_a_delivery_failure() -> TestResult {
        let runtime = test_runtime()?;
        let pid = runtime.spawn_test_process()?;
        runtime.nif_state().cleanup_process(pid);
        assert!(
            runtime.is_live(pid),
            "fixture control: cleanup must precede scheduler pid retirement"
        );
        assert!(
            runtime.process_cleanup_started(pid),
            "fixture control: cleanup must stamp the exit tombstone"
        );

        assert_eq!(
            runtime.classify_undelivered_wake(pid)?,
            UndeliveredWake::ProcessEnded,
            "a wake-marker failure after cleanup starts is the completed target's \
             already-durable signal, never a delivery failure"
        );
        runtime.shutdown()?;
        Ok(())
    }

    /// The control for the arm above: a live process that exit cleanup has
    /// never touched cannot excuse an undeliverable wake marker.
    #[test]
    fn signal_to_live_pid_without_cleanup_is_a_delivery_failure() -> TestResult {
        let runtime = test_runtime()?;
        let pid = runtime.spawn_test_process()?;
        assert!(runtime.is_live(pid));
        assert!(!runtime.process_cleanup_started(pid));

        assert_eq!(
            runtime.classify_undelivered_wake(pid)?,
            UndeliveredWake::ProcessDidNotEnd { in_table: true },
            "a live process with no cleanup underway must still report DeliveryFailed"
        );
        runtime.shutdown()?;
        Ok(())
    }
}