aion-rs 0.29.0

//! Round-12 regressions for reservations, child outcome release, and callback isolation.

use std::sync::{Arc, mpsc};
use std::time::{Duration, Instant};

use beamr::scheduler::EXIT_EVENT_CAPACITY;

use crate::EngineError;
use crate::runtime::{RuntimeConfig, RuntimeHandle, RuntimeInput, SignalDeliveryConfig};

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

/// The bound these pins wait under: small enough that a blocked callback's
/// bounded failure lands well inside the two-second assertion below, large
/// enough that a callback that merely runs is never mistaken for one that
/// blocks (AE-017: the bound restarts on every completed job).
const ROUND12_STOP_DRAIN_TIMEOUT: Duration = Duration::from_millis(500);

fn fixture_workflow_beam() -> &'static [u8] {
    include_bytes!("../../tests/fixtures/aion_fixture_workflow.beam")
}

fn wait_until(deadline: Instant, predicate: impl Fn() -> bool) -> TestResult {
    while Instant::now() < deadline {
        if predicate() {
            return Ok(());
        }
        std::thread::sleep(Duration::from_millis(1));
    }
    Err("round-12 process-exit condition missed its deadline".into())
}

#[test]
fn lag_recovery_waits_for_spawn_reservation_before_snapshot() -> TestResult {
    let runtime = Arc::new(RuntimeHandle::new(RuntimeConfig::new(
        Some(1),
        ROUND12_STOP_DRAIN_TIMEOUT,
    ))?);
    runtime.process_exits.pause_for_test();
    let first_pid = runtime.spawn_test_process()?;
    runtime.cancel_pid(first_pid)?;
    runtime
        .process_exits
        .wait_for_pause_for_test(Duration::from_secs(10))?;

    for _ in 0..=EXIT_EVENT_CAPACITY {
        let pid = runtime.spawn_test_process()?;
        runtime.cancel_pid(pid)?;
    }

    runtime.process_exits.pause_next_registration();
    let spawn_runtime = Arc::clone(&runtime);
    let spawn = std::thread::spawn(move || spawn_runtime.spawn_test_process());
    let fast_pid = runtime
        .process_exits
        .wait_for_registration_pause(Duration::from_secs(10))?;
    runtime.cancel_pid(fast_pid)?;
    let lag_recoveries = runtime.process_exits.lag_recoveries_for_test();
    runtime.process_exits.release_for_test();
    wait_until(Instant::now() + Duration::from_secs(10), || {
        runtime.process_exits.lag_recoveries_for_test() > lag_recoveries
    })
    .map_err(
        |_| "exit drainer did not enter lag resynchronization while the spawn reservation was held",
    )?;
    runtime.process_exits.release_registration();
    let registered_pid = spawn
        .join()
        .map_err(|_| "reserved spawn thread terminated unexpectedly")??;
    assert_eq!(registered_pid, fast_pid);

    let (sender, receiver) = mpsc::channel();
    runtime.monitor_process_for_test(fast_pid, move |outcome| {
        let _ = sender.send(outcome.is_ok());
    })?;
    assert!(receiver.recv_timeout(Duration::from_secs(10))?);
    assert!(receiver.try_recv().is_err());
    wait_until(Instant::now() + Duration::from_secs(10), || {
        !runtime.process_exits.contains(fast_pid)
    })?;

    let shutdown_started = Instant::now();
    runtime.shutdown()?;
    assert!(shutdown_started.elapsed() < Duration::from_secs(10));
    Ok(())
}

#[test]
fn beam_spawn_children_release_outcomes_to_baseline() -> TestResult {
    let runtime = RuntimeHandle::new(RuntimeConfig::new(Some(2), ROUND12_STOP_DRAIN_TIMEOUT))?;
    runtime.register_module("aion_fixture_workflow", fixture_workflow_beam())?;
    let baseline = runtime.process_exits.unobserved_children_for_test()?;

    let parent = runtime.spawn_workflow(
        "aion_fixture_workflow",
        "spawn_children",
        RuntimeInput::default(),
    )?;
    let _ = runtime.process_exit_for_test(parent)?;
    wait_until(Instant::now() + Duration::from_secs(10), || {
        matches!(
            runtime.process_exits.unobserved_children_for_test(),
            Ok(count) if count == baseline
        )
    })?;

    assert_eq!(
        runtime.process_exits.unobserved_children_for_test()?,
        baseline
    );
    runtime.shutdown()?;
    Ok(())
}

#[test]
fn lag_recovery_releases_beam_spawn_children_to_baseline() -> TestResult {
    let runtime = RuntimeHandle::new(RuntimeConfig::new(Some(2), ROUND12_STOP_DRAIN_TIMEOUT))?;
    runtime.register_module("aion_fixture_workflow", fixture_workflow_beam())?;
    let baseline = runtime.process_exits.unobserved_children_for_test()?;
    runtime.process_exits.pause_for_test();

    let parent = runtime.spawn_workflow(
        "aion_fixture_workflow",
        "overflow_children",
        RuntimeInput::default(),
    )?;
    runtime
        .process_exits
        .wait_for_pause_for_test(Duration::from_secs(10))?;
    wait_until(Instant::now() + Duration::from_secs(20), || {
        !runtime.is_live(parent)
    })
    .map_err(|_| "overflow parent remained live")?;
    let last_child = parent + 1_100;
    wait_until(Instant::now() + Duration::from_secs(20), || {
        (parent + 1..=last_child).all(|pid| !runtime.is_live(pid))
    })
    .map_err(|_| "overflow children remained live")?;
    runtime.process_exits.release_for_test();
    let _ = runtime.process_exit_for_test(parent)?;
    wait_until(Instant::now() + Duration::from_secs(20), || {
        runtime.process_exits.lag_recoveries_for_test() > 0
            && matches!(
                runtime.process_exits.unobserved_children_for_test(),
                Ok(count) if count == baseline
            )
    })
    .map_err(|_| {
        format!(
            "overflow recovery stalled; lag recoveries={}",
            runtime.process_exits.lag_recoveries_for_test()
        )
    })?;

    assert_eq!(
        runtime.process_exits.unobserved_children_for_test()?,
        baseline
    );
    runtime.shutdown()?;
    Ok(())
}

#[test]
fn blocking_callback_cannot_backpressure_drainer_and_shutdown_is_retryable() -> TestResult {
    let delivery = SignalDeliveryConfig::new(
        Duration::from_millis(50),
        1,
        Duration::from_millis(1),
        Duration::from_millis(1),
    );
    let runtime = Arc::new(RuntimeHandle::new(
        RuntimeConfig::new(Some(1), ROUND12_STOP_DRAIN_TIMEOUT).with_signal_delivery(delivery),
    )?);
    let blocked_pid = runtime.spawn_test_process()?;
    let observed_pid = runtime.spawn_test_process()?;
    let (entered_sender, entered_receiver) = mpsc::channel();
    let (release_sender, release_receiver) = mpsc::channel();
    runtime.monitor_process_for_test(blocked_pid, move |_| {
        let _ = entered_sender.send(());
        let _ = release_receiver.recv();
    })?;
    let (observed_sender, observed_receiver) = mpsc::channel();
    runtime.monitor_process_for_test(observed_pid, move |outcome| {
        let _ = observed_sender.send(outcome.is_ok());
    })?;

    runtime.cancel_pid(blocked_pid)?;
    entered_receiver.recv_timeout(Duration::from_secs(10))?;
    runtime.cancel_pid(observed_pid)?;
    wait_until(Instant::now() + Duration::from_secs(10), || {
        matches!(runtime.process_exits.has_terminal(observed_pid), Ok(true))
    })?;

    let shutdown_started = Instant::now();
    let first_shutdown = runtime
        .shutdown()
        .err()
        .ok_or("blocked callback did not produce bounded shutdown failure")?;
    assert!(matches!(
        first_shutdown,
        EngineError::ProcessExitCallbackDispatcherShutdownTimedOut { .. }
    ));
    assert!(shutdown_started.elapsed() < Duration::from_secs(2));
    release_sender.send(())?;
    assert!(observed_receiver.recv_timeout(Duration::from_secs(10))?);
    assert!(observed_receiver.try_recv().is_err());
    runtime.shutdown()?;
    Ok(())
}

#[test]
fn panicking_callback_cannot_terminate_singleton_drainer() -> TestResult {
    let runtime = Arc::new(RuntimeHandle::new(RuntimeConfig::new(
        Some(1),
        ROUND12_STOP_DRAIN_TIMEOUT,
    ))?);
    let panicking_pid = runtime.spawn_test_process()?;
    let observed_pid = runtime.spawn_test_process()?;
    let (panic_entered_sender, panic_entered_receiver) = mpsc::channel();
    runtime.monitor_process_for_test(panicking_pid, move |_| {
        let _ = panic_entered_sender.send(());
        std::panic::resume_unwind(Box::new(String::from("intentional callback panic")));
    })?;
    let (observed_sender, observed_receiver) = mpsc::channel();
    runtime.monitor_process_for_test(observed_pid, move |outcome| {
        let _ = observed_sender.send(outcome.is_ok());
    })?;

    runtime.cancel_pid(panicking_pid)?;
    panic_entered_receiver.recv_timeout(Duration::from_secs(10))?;
    runtime.cancel_pid(observed_pid)?;
    assert!(observed_receiver.recv_timeout(Duration::from_secs(10))?);
    assert!(observed_receiver.try_recv().is_err());
    assert!(runtime.process_exits.has_terminal(observed_pid)?);
    runtime.shutdown()?;
    Ok(())
}

/// AE-017: a drain is a NO-PROGRESS bound. Four callbacks that each run
/// longer than half the bound — a queue that takes several bounds in total —
/// are waited out, because the clock restarts on every completed job. Under
/// the old total bound this stop failed; under the new one it succeeds.
#[test]
fn a_slow_but_progressing_dispatcher_is_waited_out() -> TestResult {
    let bound = Duration::from_millis(300);
    let runtime = Arc::new(RuntimeHandle::new(RuntimeConfig::new(Some(1), bound))?);
    let (done_sender, done_receiver) = mpsc::channel();
    let mut pids = Vec::new();
    for _ in 0..4 {
        let pid = runtime.spawn_test_process()?;
        let done = done_sender.clone();
        runtime.monitor_process_for_test(pid, move |_| {
            std::thread::sleep(Duration::from_millis(180));
            let _ = done.send(());
        })?;
        pids.push(pid);
    }
    for pid in &pids {
        runtime.cancel_pid(*pid)?;
    }
    // Let the first callback be IN FLIGHT when the stop arrives, so the drain
    // is measured with work outstanding rather than on an idle worker.
    std::thread::sleep(Duration::from_millis(50));
    let stop_started = Instant::now();
    runtime.shutdown()?;
    let elapsed = stop_started.elapsed();
    assert!(
        elapsed >= Duration::from_millis(400),
        "the stop waited for the whole queue ({elapsed:?}), which is longer than one bound"
    );
    let mut completed = 0;
    while done_receiver.try_recv().is_ok() {
        completed += 1;
    }
    assert_eq!(
        completed, 4,
        "every queued callback ran to completion before the stop returned"
    );
    Ok(())
}

/// AE-017: while a drain waits on a callback that has not completed, it says
/// so at WARN — the worker's name, the in-flight age, the queue depth — and
/// the timeout error carries the same facts. The evidence of a wedge is in
/// the log before the stop gives up, not only in the error after.
#[test]
fn a_waiting_drain_names_the_in_flight_callback_at_warn() -> TestResult {
    let bound = Duration::from_millis(400);
    let (logs, subscriber) = crate::log_capture::LogCapture::new()?;
    let guard = tracing::subscriber::set_default(subscriber);
    let runtime = Arc::new(RuntimeHandle::new(RuntimeConfig::new(Some(1), bound))?);
    let blocked_pid = runtime.spawn_test_process()?;
    let (entered_sender, entered_receiver) = mpsc::channel();
    let (release_sender, release_receiver) = mpsc::channel();
    runtime.monitor_process_for_test(blocked_pid, move |_| {
        let _ = entered_sender.send(());
        let _ = release_receiver.recv();
    })?;
    runtime.cancel_pid(blocked_pid)?;
    entered_receiver.recv_timeout(Duration::from_secs(10))?;

    let failure = runtime
        .shutdown()
        .err()
        .ok_or("a blocked callback must produce a bounded stop failure")?;
    match failure {
        EngineError::ProcessExitCallbackDispatcherShutdownTimedOut {
            timeout_millis,
            since_progress_millis,
            queued,
        } => {
            assert_eq!(timeout_millis, 400);
            assert!(
                since_progress_millis >= 400,
                "bounded from last progress: {since_progress_millis}"
            );
            assert_eq!(queued, 0, "the blocked job is in flight, not queued");
        }
        other => {
            return Err(format!("expected the dispatcher's bounded failure, got {other:?}").into());
        }
    }
    drop(guard);
    let warnings = logs.at_level("WARN")?;
    let progress = warnings
        .iter()
        .find(|line| line.mentions("engine stop still draining"))
        .ok_or("the drain must say at WARN that it is still waiting")?;
    assert_eq!(progress.field("worker"), Some("aion-process-exit-callback"));
    assert!(
        progress.field("in_flight_ms").is_some(),
        "the in-flight age is named"
    );
    assert!(
        progress.field("queued").is_some(),
        "the queue depth is named"
    );

    release_sender.send(())?;
    runtime.shutdown()?;
    Ok(())
}