aion-rs 0.31.0

//! Activity dispatch bridge for the `aion_flow_ffi` activity NIFs.
//!
//! The bridge decouples the raw NIF function pointer (which cannot capture
//! state) from the engine's activity execution path. A concrete dispatcher
//! is installed on the engine's NIF state during build; the NIF recovers it
//! from its calling context's NIF private data.

use std::collections::BTreeMap;
use std::sync::Arc;

use aion_core::{ActivityId, RunId, WorkflowId};
use futures::future::BoxFuture;

/// A fully-resolved activity dispatch crossing the engine→transport seam.
///
/// The engine builds this once the durability layer has assigned the
/// activity its ordinal, so it carries the *real* owning [`WorkflowId`] and
/// the *real* per-workflow [`ActivityId`] recorded in history — not a
/// transport-fabricated correlation token. A worker that logs these ids can
/// therefore be correlated directly against the event store, and a result
/// re-reported from a previous worker session keys to the exact execution it
/// belongs to.
///
/// `input` and `config` are the JSON strings the Gleam SDK sends through the
/// `aion_flow_ffi:dispatch_activity/3` binding; `attempt` is the one-based
/// delivery attempt the engine stamps onto the transport (the worker wire's
/// `ActivityTask.attempt`), so consumers distinguish retries without guessing.
#[derive(Clone, Debug)]
pub struct ActivityDispatch {
    /// Namespace selected for worker matching — the correctness/isolation
    /// boundary the activity may dispatch within.
    pub namespace: String,
    /// Task queue (pool/flavour) selected within the namespace. The worker-pool
    /// address is `(namespace, task_queue)`. The engine resolves the activity
    /// override, workflow default, or recorded start-time queue before this seam.
    pub task_queue: String,
    /// OPTIONAL node affinity (NODE-4): the concrete worker host this dispatch is
    /// pinned to within the `(namespace, task_queue)` pool. `None` = no affinity
    /// (any worker in the pool). Resolved once at the schedule seam from the
    /// SDK's per-activity `node` selection; there is no workflow-level default.
    pub node: Option<String>,
    /// Real owning workflow id, recorded in history at `WorkflowStarted`.
    pub workflow_id: WorkflowId,
    /// Concrete workflow run. Required for a run-scoped idempotency key and
    /// echoed by the worker on completion.
    pub run_id: RunId,
    /// Real per-workflow activity ordinal, recorded at `ActivityScheduled`.
    pub activity_id: ActivityId,
    /// Registered activity-type name to match against worker registrations.
    pub name: String,
    /// JSON-encoded activity input.
    pub input: String,
    /// JSON-encoded dispatch config (retry/timeout/heartbeat policy).
    pub config: String,
    /// One-based delivery attempt for this dispatch.
    pub attempt: u32,
    /// Human-meaningful display labels the workflow attached to the activity
    /// (for example `brief=IP-001`, `repo=ablative-io/yggdrasil`). The engine
    /// never interprets these; they ride to the worker purely so its logs and
    /// the ops console can show what a dispatch is working on. `BTreeMap` keeps
    /// the rendered order stable.
    pub labels: BTreeMap<String, String>,
    /// Whether the DECLARATION classes this activity as advisory: a side
    /// channel whose exhaustion warns on the run and never faults the calling
    /// step (RUNTIME-OPERATIONS.md R5).
    ///
    /// Resolved at the schedule seam from the package contract this run is
    /// pinned to (`nif_activity_advisory`), never from the dispatch config the
    /// SDK builds — the class is declaration-owned and a call site cannot
    /// forge it. Dispatchers ignore it; it exists so the retry loop knows to
    /// record the warning when the attempt budget is spent.
    pub advisory: bool,
}

/// The one-shot "a worker has taken this attempt" event, carried across the
/// engine→transport seam.
///
/// # Why an event and not an instant
///
/// The per-attempt bound is authored in the document and lives in the engine;
/// the lease happens in the server, on the far side of a `spawn_blocking`
/// boundary. Handing the lease INSTANT back would be handing back a value that
/// arrives after the thing it dates, and the timer would still have to decide
/// what to do with the gap. Handing back the EVENT lets the clock start where
/// the lease happens, by construction.
///
/// # What it fixes
///
/// The per-attempt bound used to wrap the whole dispatch future, and the first
/// thing inside that future is an unbounded park waiting for a worker to exist.
/// So schedule-to-start time was charged to a bound the document author wrote
/// to describe EXECUTION: an activity that waited eleven minutes for a worker
/// and then ran for four seconds could exceed a five-minute bound without ever
/// having been slow. Worse, the expiry DISCARDS a result the worker really
/// produced — a completed measurement thrown away by its own clock.
///
/// # The park is now bounded ONLY if an operator bounds it
///
/// Say this plainly, because it is a behaviour change and not a refinement. An
/// authored per-attempt timeout used to end a dispatch that was still waiting
/// for a worker; it no longer does. It bounds the ATTEMPT, from the lease, and
/// nothing else.
///
/// What bounds the wait instead is the queue-service policy — the
/// service-availability and schedule-to-start clocks — and both of those are
/// `Option<Duration>` that default to `None`. So on a server whose operator has
/// set neither, a dispatch to a queue with no eligible worker waits
/// indefinitely where it used to fail at the per-attempt bound.
///
/// That is deliberate and, on balance, the better behaviour: the old expiry
/// could not cancel the blocking dispatch it gave up on, so each one stacked
/// another parked thread on top of the first — the bound amplified threads
/// rather than freeing them. But no default is invented here to replace it.
/// Deciding how long work may wait for a worker that does not exist is an
/// operator's call about their own fleet, and the honest thing is to say the
/// clock is theirs to set rather than to pick one for them.
/// Cloning shares one signal: the transport that fires it and the caller that
/// waits on it hold the same notification, and a clone that outlives the
/// dispatch simply never fires.
#[derive(Clone, Debug)]
pub struct LeaseSignal {
    leased: Arc<tokio::sync::Notify>,
}

impl LeaseSignal {
    /// A signal and the future that waits for it.
    ///
    /// The wait retains a fire that lands before it is awaited (the notify
    /// stores one permit), so a lease that happens while the caller is still
    /// setting up cannot be missed.
    #[must_use]
    pub fn channel() -> (Self, Leased) {
        let leased = Arc::new(tokio::sync::Notify::new());
        (
            Self {
                leased: Arc::clone(&leased),
            },
            Leased { leased },
        )
    }

    /// A signal nobody is waiting on.
    ///
    /// For the paths that dispatch WITHOUT a bound to start — the synchronous
    /// [`ActivityDispatcher::dispatch`] entry, which has no timer above it.
    /// Firing it is a no-op rather than an error: the signal's meaning is "a
    /// worker took this attempt", which is true whether or not anybody is
    /// timing it.
    #[must_use]
    pub fn none() -> Self {
        Self {
            leased: Arc::new(tokio::sync::Notify::new()),
        }
    }

    /// A worker holds this attempt. The per-attempt clock starts HERE.
    ///
    /// Idempotent and cheap: `Notify::notify_one` stores AT MOST ONE permit, so
    /// a second fire on an already-fired signal is a no-op. Deliberately takes `&self` so it can be fired from inside the
    /// transport's own accept hook, which is the instant that also produces the
    /// durable `ActivityLeased` record — one anchor for both, by construction
    /// rather than by comment.
    pub fn fire(&self) {
        self.leased.notify_one();
    }
}

/// The waiting half of a [`LeaseSignal`].
#[derive(Debug)]
pub struct Leased {
    leased: Arc<tokio::sync::Notify>,
}

impl Leased {
    /// Wait for the attempt to be leased.
    ///
    /// This future never completes on its own if no lease ever happens, and
    /// that is correct: its only caller races it against the dispatch itself,
    /// so a dispatch that fails before any worker takes it resolves through the
    /// other arm. A bound must never start on a dispatch that never began.
    pub async fn wait(self) {
        self.leased.notified().await;
    }
}

/// Executes an activity request originating from workflow code.
///
/// The return value is the JSON-encoded activity result or a prefixed error
/// string matching the SDK's error-decoding convention.
pub trait ActivityDispatcher: Send + Sync + 'static {
    /// Dispatch the activity and block until completion.
    ///
    /// Returns `Ok(encoded_output)` on success or `Err(error_string)` on
    /// failure. Both sides are strings matching the Gleam SDK's
    /// `Result(String, String)` FFI contract.
    ///
    /// # Errors
    ///
    /// Returns the error string surfaced by the activity execution path —
    /// worker rejection, decode failure, timeout, or activity body error.
    fn dispatch(&self, request: ActivityDispatch) -> Result<String, String>;

    /// Dispatch the activity from a Tokio task.
    ///
    /// The default runs the synchronous [`Self::dispatch`] on the runtime's
    /// blocking pool via [`tokio::task::spawn_blocking`], so a dispatcher that
    /// blocks its calling thread cannot wedge the engine's async workers (a
    /// single-threaded engine runtime keeps servicing queries and completions
    /// while the dispatch waits). Nonblocking dispatchers can override this.
    ///
    /// `lease` is fired at the instant a worker HOLDS this attempt. It is what
    /// starts the authored per-attempt bound, so an implementation that parks
    /// waiting for a worker keeps that park outside the bound — the park has
    /// its own clocks (service-availability and schedule-to-start), and
    /// charging schedule time to an execution bound is what let a completed
    /// activity be discarded by its own timer.
    ///
    /// Must be awaited inside a Tokio runtime context; the engine's
    /// completion task guarantees that.
    ///
    /// # Errors
    ///
    /// Returns the same errors as [`Self::dispatch`], plus a dispatch-failure
    /// reason when the blocking task itself is cancelled or panics.
    fn dispatch_async(
        self: Arc<Self>,
        request: ActivityDispatch,
        lease: LeaseSignal,
    ) -> BoxFuture<'static, Result<String, String>> {
        Box::pin(async move {
            // A dispatcher on this default has NO lease step: it runs the
            // synchronous entry above, which begins executing the activity the
            // moment it is called. So the attempt begins when the call begins,
            // and the signal fires HERE — the honest anchor for a dispatcher
            // that never parks waiting for a worker.
            //
            // A dispatcher that DOES park — one that selects a worker, waits
            // for one to exist, and hands the work over — overrides this method
            // and fires at its own handover instead. Leaving the signal unfired
            // is not an option for either: the bound above would then never
            // start, and an authored per-attempt timeout would silently stop
            // bounding anything.
            lease.fire();
            let blocking = tokio::task::spawn_blocking(move || self.dispatch(request));
            match blocking.await {
                Ok(result) => result,
                Err(join_error) => Err(format!("activity dispatch task failed: {join_error}")),
            }
        })
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use std::collections::BTreeMap;

    use aion_core::{ActivityId, RunId, WorkflowId};

    use super::{ActivityDispatch, ActivityDispatcher};
    use crate::runtime::EngineNifState;

    struct Echo;

    impl ActivityDispatcher for Echo {
        fn dispatch(&self, request: ActivityDispatch) -> Result<String, String> {
            Ok(request.input)
        }
    }

    fn echo_request(input: &str) -> ActivityDispatch {
        ActivityDispatch {
            namespace: "default".to_owned(),
            task_queue: "default".to_owned(),
            node: None,
            workflow_id: WorkflowId::new_v4(),
            run_id: RunId::new_v4(),
            activity_id: ActivityId::from_sequence_position(0),
            name: "test".to_owned(),
            input: input.to_owned(),
            config: "{}".to_owned(),
            attempt: 1,
            labels: BTreeMap::new(),
            advisory: false,
        }
    }

    #[test]
    fn dispatcher_is_accessible_after_install_on_engine_state() {
        let state = EngineNifState::default();
        state.set_activity_dispatcher(Arc::new(Echo));
        let dispatcher = state.activity_dispatcher();
        assert!(dispatcher.is_some());
        assert_eq!(
            dispatcher
                .as_ref()
                .and_then(|d| d.dispatch(echo_request("hello")).ok()),
            Some("hello".to_owned())
        );
    }
}