aion-server 0.31.0

//! aion#204 at small numbers: fan fourteen activities at a concurrency-four
//! worker and watch what the server says about it.
//!
//! # The shape this reproduces
//!
//! The live incident was a fan-out onto one gRPC worker running at its
//! configured concurrency. The worker stopped reading its task stream while it
//! waited for a permit, so it answered no liveness pings; the server withdrew
//! its dispatch eligibility (correct), the heartbeat sweep DEREGISTERED it
//! while it was connected and holding live work (wrong), the queue census —
//! now looking at an empty pool — reported that no worker was connected (a
//! false statement about a worker on an open stream), and the per-attempt
//! bound, which had been ticking through the whole park, discarded work the
//! worker had actually finished. Thirty-seven minutes of an empty worker pool
//! that was never empty.
//!
//! Nothing here is faked below the server's own seams: a REAL `aion-server`
//! worker service over TCP loopback, the REAL Rust `aion-worker` SDK, and the
//! REAL engine-seam dispatcher. Only the numbers are small enough to observe in
//! a test.
//!
//! # What is asserted, and why each half is needed
//!
//! - **Every activity completes exactly once.** The outage's seal was work that
//!   ran and was thrown away, so counting executions is not enough: the outputs
//!   have to come back, all fourteen of them, each once.
//! - **The worker is never deregistered.** Same worker id at the end as at the
//!   start. This is the arm the heartbeat sweep used to break.
//! - **The census never says `NO_LIVE_POLLERS`.** A worker on an open stream is
//!   connected, and a queue served by one is never an empty pool. This is the
//!   false statement the operator acted on.
//! - **The census DOES say `POLLERS_AT_CAPACITY`.** The negative above passes
//!   trivially if the fan never saturates the worker, so the positive control
//!   is what proves the run actually reached the condition under test.
//! - **Concurrency is respected.** The handler's peak overlap never exceeds
//!   four: the server does not push a worker past what it advertised.

use std::collections::BTreeSet;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use aion::{ActivityDispatch, ActivityDispatcher as _};
use aion_core::{ActivityId, WorkflowId};
use aion_server::ServerState;
use aion_server::api::worker_grpc::worker_service;
use aion_server::config::{
    AuthConfig, AuthoringConfig, DeployConfig, ListenConfig, MetricsConfig, NamespaceConfig,
    NamespaceMode, OpsConsoleAssetSource, OpsConsoleConfig, RuntimeConfig, WebSocketConfig,
    WorkerConfig,
};
use aion_server::worker::{
    ConnectedWorkerRegistry, QueueServiceReason, WorkerActivityDispatcher, WorkerId,
};
use aion_server::{NamespaceResolver, StaticScheduleNamespaces, StaticWorkflowNamespaces};
use aion_worker::{ReconnectConfig, Worker};
use serde::{Deserialize, Serialize};
use tokio::net::TcpListener;
use tokio_stream::wrappers::TcpListenerStream;

type TestError = Box<dyn std::error::Error>;

const NAMESPACE: &str = "default";
const TASK_QUEUE: &str = "default";
const ACTIVITY_TYPE: &str = "hold";

/// The worker's advertised concurrency — the number the server must not push
/// past, and the number the worker's own semaphore enforces.
const WORKER_CONCURRENCY: usize = 4;

/// How many activities are fanned at it. More than three times the capacity, so
/// the queue is deep enough that a run cannot pass by being scheduled leniently.
const FAN: usize = 14;

/// How long each activity holds its slot. Long enough that the fan genuinely
/// overlaps and the pool is observably full; short enough that fourteen of them
/// clear in a handful of seconds.
const HOLD: Duration = Duration::from_millis(120);

#[derive(Debug, Clone, Serialize, Deserialize)]
struct HoldInput {
    ordinal: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct HoldOutput {
    ordinal: usize,
}

/// Records how many activities ran at once, and how many ran in total.
#[derive(Default)]
struct Overlap {
    running: AtomicUsize,
    peak: AtomicUsize,
    executions: AtomicUsize,
}

impl Overlap {
    fn enter(&self) {
        let now = self.running.fetch_add(1, Ordering::SeqCst) + 1;
        self.executions.fetch_add(1, Ordering::SeqCst);
        self.peak.fetch_max(now, Ordering::SeqCst);
    }

    fn leave(&self) {
        self.running.fetch_sub(1, Ordering::SeqCst);
    }
}

fn hold_request(ordinal: usize) -> Result<ActivityDispatch, TestError> {
    Ok(ActivityDispatch {
        namespace: NAMESPACE.to_owned(),
        task_queue: TASK_QUEUE.to_owned(),
        node: None,
        workflow_id: WorkflowId::new_v4(),
        run_id: aion_core::RunId::new_v4(),
        activity_id: ActivityId::from_sequence_position(0),
        name: ACTIVITY_TYPE.to_owned(),
        input: serde_json::to_string(&HoldInput { ordinal })?,
        config: "{}".to_owned(),
        attempt: 1,
        labels: std::collections::BTreeMap::new(),
        advisory: false,
    })
}

fn runtime_config() -> RuntimeConfig {
    RuntimeConfig {
        listen: ListenConfig {
            grpc: SocketAddr::from(([127, 0, 0, 1], 0)),
            http: SocketAddr::from(([127, 0, 0, 1], 0)),
        },
        tls: None,
        auth: AuthConfig {
            enabled: false,
            jwks_url: None,
            jwks_refresh_seconds: 300,
        },
        ops_console: OpsConsoleConfig {
            source: OpsConsoleAssetSource::Embedded,
        },
        namespace: NamespaceConfig {
            mode: NamespaceMode::SharedEngine,
        },
        worker: WorkerConfig::default(),
        websocket: WebSocketConfig {
            outbound_buffer_bound: 32,
            event_broadcast_capacity: Some(64),
            cluster_broadcast_capacity: Some(64),
        },
        workflow_packages: Vec::new(),
        deploy: DeployConfig::default(),
        authoring: AuthoringConfig::default(),
        dev: aion_server::config::DevConfig::default(),
        outbox: aion_server::config::OutboxConfig::default(),
        observability: aion_server::config::ObservabilityConfig::with_flush_policy(64, 0),
        mcp: aion_server::config::ResolvedMcpConfig::default(),
        assistant: aion_server::config::ResolvedAssistantConfig::default(),
        scheduler_threads: 1,
        stop_drain_timeout: Some(std::time::Duration::from_secs(5)),
        jit_threshold: None,
        query_timeout: Some(Duration::from_secs(10)),
        workloop_sweep_interval: Some(std::time::Duration::from_millis(50)),
        default_namespace: NAMESPACE.to_owned(),
        auto_create: aion_server::config::AutoCreate::Open,
        max_in_flight_activities: aion_server::config::DEFAULT_MAX_IN_FLIGHT_ACTIVITIES,
        drain_timeout: Duration::from_secs(30),
        metrics: MetricsConfig { enabled: false },
        owned_shards: Vec::new(),
        cors_allowed_origins: Vec::new(),
    }
}

async fn registered_worker(registry: &ConnectedWorkerRegistry) -> Result<WorkerId, TestError> {
    let deadline = Instant::now() + Duration::from_secs(10);
    loop {
        if let Some(handle) = registry
            .workers_for(NAMESPACE, TASK_QUEUE, ACTIVITY_TYPE, None)?
            .first()
        {
            return Ok(handle.id());
        }
        if Instant::now() >= deadline {
            return Err("worker did not register with the server in time".into());
        }
        tokio::time::sleep(Duration::from_millis(10)).await;
    }
}

/// A real server, a real SDK worker registered against it, and the handles a
/// test needs to interrogate both.
struct Harness {
    state: ServerState,
    registry: ConnectedWorkerRegistry,
    overlap: Arc<Overlap>,
    worker_id: WorkerId,
    shutdown: tokio::sync::oneshot::Sender<()>,
    worker_run: tokio::task::JoinHandle<Result<(), aion_worker::WorkerError>>,
    server: tokio::task::JoinHandle<Result<(), tonic::transport::Error>>,
}

impl Harness {
    /// Boot the worker service on loopback and connect ONE real SDK worker
    /// advertising [`WORKER_CONCURRENCY`], returning once it is registered.
    async fn start() -> Result<Self, TestError> {
        let listener = TcpListener::bind("127.0.0.1:0").await?;
        let address = listener.local_addr()?;
        let registry = ConnectedWorkerRegistry::default();
        let resolver = NamespaceResolver::authorization_only(
            NamespaceMode::SharedEngine,
            StaticWorkflowNamespaces::default(),
            StaticScheduleNamespaces::default(),
        );
        let state =
            ServerState::from_parts_with_registry(resolver, runtime_config(), registry.clone());
        let server = tokio::spawn(
            tonic::transport::Server::builder()
                .add_service(worker_service(state.clone()))
                .serve_with_incoming(TcpListenerStream::new(listener)),
        );

        let overlap = Arc::new(Overlap::default());
        let worker_config = aion_worker::WorkerConfig::new(
            format!("http://{address}"),
            NAMESPACE,
            "capacity-e2e-worker",
            WORKER_CONCURRENCY,
            ReconnectConfig::new(Duration::from_millis(50), Duration::from_secs(2), 5),
            None,
        );
        let worker = Worker::builder(worker_config)
            .register_activity(ACTIVITY_TYPE, {
                let overlap = Arc::clone(&overlap);
                move |input: HoldInput, _context: &aion_worker::ActivityContext| {
                    let overlap = Arc::clone(&overlap);
                    Box::pin(async move {
                        overlap.enter();
                        tokio::time::sleep(HOLD).await;
                        overlap.leave();
                        Ok(HoldOutput {
                            ordinal: input.ordinal,
                        })
                    })
                }
            })?
            .build()?;
        let (shutdown, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
        let worker_run = tokio::spawn(worker.run_until(async move {
            let _ = shutdown_rx.await;
        }));
        let worker_id = registered_worker(&registry).await?;
        Ok(Self {
            state,
            registry,
            overlap,
            worker_id,
            shutdown,
            worker_run,
            server,
        })
    }

    /// The production engine-seam dispatcher over this harness's shared state.
    fn dispatcher(&self) -> WorkerActivityDispatcher {
        WorkerActivityDispatcher::new(
            self.registry.clone(),
            NAMESPACE,
            self.state.heartbeat_tracker().clone(),
        )
        .with_pending(self.state.pending_activities().clone())
        .with_drain_state(self.state.drain_state().clone())
        // The queue-service seams are what make a park VISIBLE: without them
        // the dispatcher still classifies and still logs, but nothing is
        // queryable, and this test's whole subject is what the server SAYS
        // about a busy queue.
        .with_queue_declarations(self.state.queue_declarations().clone())
        .with_queue_state(self.state.queue_service_state().clone())
    }

    async fn shutdown(self) -> Result<(), TestError> {
        let _ = self.shutdown.send(());
        self.worker_run.await??;
        self.server.abort();
        Ok(())
    }
}

/// Every unserved-queue reason the server published while the fan ran.
///
/// POLLED rather than sampled once at the end, because the interesting verdicts
/// are transient: the queue is only unserved while every worker is busy, and a
/// check that ran after the last activity finished would find nothing and pass
/// having observed nothing.
struct CensusWatch {
    observed: Arc<Mutex<BTreeSet<QueueServiceReason>>>,
    watching: Arc<std::sync::atomic::AtomicBool>,
    task: tokio::task::JoinHandle<()>,
}

impl CensusWatch {
    fn start(state: &ServerState) -> Self {
        let observed: Arc<Mutex<BTreeSet<QueueServiceReason>>> =
            Arc::new(Mutex::new(BTreeSet::new()));
        let watching = Arc::new(std::sync::atomic::AtomicBool::new(true));
        let task = tokio::spawn({
            let queue_state = state.queue_service_state().clone();
            let observed = Arc::clone(&observed);
            let watching = Arc::clone(&watching);
            async move {
                while watching.load(Ordering::SeqCst) {
                    if let Ok(unserved) = queue_state.unserved()
                        && let Ok(mut observed) = observed.lock()
                    {
                        observed.extend(unserved.iter().map(|queue| queue.reason));
                    }
                    tokio::time::sleep(Duration::from_millis(5)).await;
                }
            }
        });
        Self {
            observed,
            watching,
            task,
        }
    }

    async fn finish(self) -> Result<BTreeSet<QueueServiceReason>, TestError> {
        self.watching.store(false, Ordering::SeqCst);
        self.task.await?;
        let observed = self
            .observed
            .lock()
            .map_err(|_| "census observation lock poisoned")?
            .clone();
        Ok(observed)
    }
}

/// R4's PRODUCTION wiring: the lease signal fires when the worker ACCEPTS the
/// dispatch, not when the call is made and not when the activity finishes.
///
/// The engine half of this is pinned against a fake dispatcher
/// (`a_bound_starts_at_the_lease_not_at_the_call`), which proves the retry loop
/// starts its bound at the signal. Nothing proved the SERVER fires that signal
/// at the right instant — `worker_capacity_e2e`'s other test drives the
/// synchronous `dispatch`, which is handed `LeaseSignal::none()`, so the anchor
/// itself was untested in production wiring.
///
/// The activity here holds for [`HOLD`]. That is what makes the assertion
/// two-sided and not merely "the signal eventually fires": the lease must be
/// observed while the dispatch is STILL RUNNING. A signal fired at completion
/// would satisfy "it fired" and would move the per-attempt bound to the wrong
/// end of the work.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn the_lease_signal_fires_when_the_worker_accepts_not_when_it_finishes()
-> Result<(), TestError> {
    let harness = Harness::start().await?;
    let dispatcher = Arc::new(harness.dispatcher());

    let (signal, leased) = aion::LeaseSignal::channel();
    let request = hold_request(0)?;
    let leg = tokio::task::spawn(Arc::clone(&dispatcher).dispatch_async(request, signal));

    // The lease must arrive well inside the activity's own hold. Bounded so a
    // signal that never fires FAILS here rather than hanging the suite.
    tokio::time::timeout(HOLD * 4, leased.wait())
        .await
        .map_err(|_| "the lease signal never fired; the per-attempt bound would never start")?;

    // AND IT ARRIVED WHILE THE WORK IS STILL RUNNING. This is the half that
    // distinguishes "fired at acceptance" from "fired at completion".
    assert!(
        !leg.is_finished(),
        "the lease fired only as the dispatch completed; the per-attempt bound would then be \
         anchored at the END of the work it is supposed to measure"
    );

    // Joined, always: a leg still parked in a blocking dispatch when this body
    // unwinds is work the runtime cannot cancel, and `Runtime::drop` joins the
    // blocking pool — an early return here would hang the suite rather than
    // fail it.
    let payload = leg
        .await
        .map_err(|error| error.to_string())?
        .map_err(|reason| format!("the leased dispatch failed: {reason}"))?;
    let output: HoldOutput = serde_json::from_str(&payload)?;
    assert_eq!(
        output.ordinal, 0,
        "the leased dispatch must be the one served"
    );

    // THE "NOT BEFORE" HALF, and the reason the assertions above are not enough
    // on their own: the DEFAULT `dispatch_async` fires the lease before it even
    // spawns the work, so a server that had simply not overridden it would pass
    // everything above. Here the dispatch is for an activity type no worker
    // serves, so it parks — the call has been made and acceptance never comes.
    // The lease must stay silent, because there is no lease: nothing has
    // accepted this attempt, and starting its clock would charge the attempt for
    // time spent waiting for a worker to exist.
    let (parked_signal, parked_lease) = aion::LeaseSignal::channel();
    let mut unserved = hold_request(1)?;
    unserved.name = String::from("nothing-serves-this");
    let parked_leg =
        tokio::task::spawn(Arc::clone(&dispatcher).dispatch_async(unserved, parked_signal));
    // OBSERVED FIRST, ASSERTED LAST. The parked leg is a blocking dispatch the
    // runtime cannot cancel, so an assertion that panics while it is still
    // outstanding would hang this suite at `Runtime::drop` instead of failing
    // it — measured, not guessed: mutating the server to fire the lease at the
    // call made an earlier draft of this test hang rather than report. So the
    // observation is banked, the leg is ended and joined, and only then is the
    // verdict given.
    let fired_while_parked = tokio::time::timeout(HOLD * 2, parked_lease.wait())
        .await
        .is_ok();

    // Ended through the drain latch, which is what a park races.
    assert!(harness.state.drain_state().begin(), "drain must begin");
    let parked = tokio::time::timeout(HOLD * 20, parked_leg)
        .await
        .map_err(|_| "the parked dispatch never returned after the drain latch fired")?
        .map_err(|error| error.to_string())?;

    assert!(
        !fired_while_parked,
        "the lease fired for a dispatch no worker has accepted; the per-attempt bound would then \
         be charged for schedule-to-start, which is the anchor this landing moved"
    );
    assert!(
        parked.is_err(),
        "a dispatch nothing serves must not report success"
    );

    harness.shutdown().await
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn a_fan_onto_a_saturated_worker_is_served_not_declared_lost() -> Result<(), TestError> {
    let harness = Harness::start().await?;
    let census = CensusWatch::start(&harness.state);
    let dispatcher = Arc::new(harness.dispatcher());

    // Every leg is spawned before ANY leg is inspected, and every spawned leg is
    // joined before the first `?` or `assert!` can leave this body. That order is
    // load-bearing rather than stylistic: a fanned leg that is still parked in
    // `WorkerActivityDispatcher::dispatch` when the body unwinds is a blocking
    // task the runtime cannot cancel, and `Runtime::drop` joins the blocking pool
    // — so an early return here does not fail the test, it HANGS it, with the
    // real cause never printed. The park is legitimately unbounded while the
    // server accepts work (see `park_for_worker`), so nothing but joining every
    // leg makes this body safe to leave.
    let mut fanned = Vec::with_capacity(FAN);
    for ordinal in 0..FAN {
        let seam = Arc::clone(&dispatcher);
        let request = hold_request(ordinal)?;
        fanned.push(tokio::task::spawn_blocking(move || seam.dispatch(request)));
    }

    let mut joined = Vec::with_capacity(FAN);
    for leg in fanned {
        joined.push(leg.await.map_err(|error| error.to_string())?);
    }

    let mut delivered = BTreeSet::new();
    for result in joined {
        let payload = result.map_err(|reason| format!("a fanned activity failed: {reason}"))?;
        let output: HoldOutput = serde_json::from_str(&payload)?;
        assert!(
            delivered.insert(output.ordinal),
            "activity {} came back twice; a fan onto a busy worker must not duplicate work",
            output.ordinal
        );
    }
    let reasons = census.finish().await?;

    assert_eq!(
        delivered.len(),
        FAN,
        "every fanned activity must come back exactly once"
    );
    assert_eq!(
        harness.overlap.executions.load(Ordering::SeqCst),
        FAN,
        "each activity must EXECUTE exactly once: the outage's seal was work that ran and was \
         then discarded by a clock, which re-ran it"
    );
    assert!(
        harness.overlap.peak.load(Ordering::SeqCst) <= WORKER_CONCURRENCY,
        "the worker ran {} activities at once against an advertised concurrency of \
         {WORKER_CONCURRENCY}: the server pushed past what the worker said it would take",
        harness.overlap.peak.load(Ordering::SeqCst)
    );

    let still_registered =
        harness
            .registry
            .workers_for(NAMESPACE, TASK_QUEUE, ACTIVITY_TYPE, None)?;
    assert_eq!(
        still_registered
            .iter()
            .map(aion_server::worker::WorkerHandle::id)
            .collect::<Vec<_>>(),
        vec![harness.worker_id],
        "the worker must still be the SAME registration: a busy worker holding live work is not \
         a lost one, and deregistering it is what emptied a pool that was never empty"
    );

    assert!(
        !reasons.contains(&QueueServiceReason::NoLivePollers),
        "the census reported NO_LIVE_POLLERS about a worker on an open stream: {reasons:?}"
    );
    assert!(
        reasons.contains(&QueueServiceReason::PollersAtCapacity),
        "the fan never saturated the worker, so the negative above proves nothing. Observed: \
         {reasons:?}"
    );

    harness.shutdown().await
}