aion_server/worker/

task_delivery.rs

//! One awaited call: hand an already-built task to an already-chosen worker,
//! over whichever transport that worker registered on.
//!
//! # Why this seam exists
//!
//! Before it, the dispatcher's delivery line assumed gRPC. A worker whose
//! `sender()` was `None` — which is every liminal-delivered worker — was
//! **deregistered like a closed stream**, on the premise that its presence could
//! only be a leak. That premise held while the outbox chose one transport for
//! the whole server, and stops holding the moment a row is routed by the
//! worker's own delivery ([aion#52]).
//!
//! # The contract, and why the blocking shape wins
//!
//! "Enqueued onto the stream" is the weakest truth a transport can offer: it
//! says a message was accepted by a channel, not that a worker took the work.
//! The liminal path already offers a stronger one — a correlated reply, with the
//! run id resolved **before** an owner is bound — and that is the truth the
//! completion fences actually want. So the seam is a single awaited call
//! returning a typed outcome, and gRPC is the fast-completing case of it rather
//! than a different shape. Latency is not the dispatcher's concern; it awaits
//! either.
//!
//! # The outcome carries the deregistration decision as a TYPE
//!
//! The distinction that used to live in a comment at the delivery line is now
//! [`Undeliverable`]'s two variants, so a caller cannot act on it by accident:
//! a worker the transport says is **gone** is deregistered, and a delivery that
//! failed while the worker is **alive** leaves the registration standing.
//!
//! [aion#52]: https://github.com/ablative-io/aion/issues/52

use async_trait::async_trait;

use aion_proto::ProtoActivityTask;

use super::delivery_intent::SharedDeliveryIntent;
use super::registry::WorkerHandle;

/// How a delivered task's liveness is tracked, as a value the transport must
/// interpret rather than a bare number it must remember the meaning of.
///
/// Distinct from [`TaskLiveness`](super::heartbeat::TaskLiveness), which is
/// the tracker's RECORD of an in-flight activity. This type says whether a
/// task is tracked that way at all.
///
/// # Why this is a type and not `0`
///
/// [`ProtoActivityTask`] carries no heartbeat window, so a transport that needs
/// one — liminal's `DispatchRequest` has the field — cannot derive it and must
/// **assign** it. Today's assignment is `0`, and it is correct only while liminal
/// dispatches stay outside the server's per-task liveness tracker: the outbox
/// retry loop is their backstop, so the worker pumps no beats for them.
///
/// Written as an exhaustive enum so that the day a liminal dispatch *is* tracked
/// per task, adding the variant makes every `match` on it non-exhaustive and the
/// compiler names each site that has to decide again. A literal `0` would simply
/// keep compiling, silently claiming "no window" for a dispatch that now needs
/// one.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub enum LivenessTracking {
    /// This transport does not track the task's liveness per task; the caller's
    /// own retry path is the backstop.
    NotTrackedPerTask,
}

impl LivenessTracking {
    /// The heartbeat window, in milliseconds, to stamp on the wire.
    ///
    /// Zero means "pump no beats for this task" — which every consumer already
    /// reads that way, and which is only ever correct for
    /// [`LivenessTracking::NotTrackedPerTask`].
    #[must_use]
    pub const fn heartbeat_window_ms(self) -> u64 {
        match self {
            Self::NotTrackedPerTask => 0,
        }
    }
}

/// Why one delivery attempt did not place its task with a worker.
///
/// The two variants are the whole reason this is a type rather than a `bool`:
/// they carry **opposite obligations** for the worker's registration, and
/// conflating them either destroys a live worker's registration or keeps a dead
/// one's forever.
#[derive(Debug)]
pub enum Undeliverable {
    /// The transport reports the worker is **gone** — a closed gRPC stream, a
    /// liminal connection the supervisor no longer owns.
    ///
    /// The caller deregisters it. This is the pre-existing behaviour of the
    /// closed-stream path, now reached for a reason that is true of both
    /// transports rather than inferred from a missing sender.
    WorkerUnreachable {
        /// Operator-facing reason, named by the transport that observed it.
        reason: String,
    },
    /// The delivery failed while the worker is **alive** — a liminal reply that
    /// did not arrive in time, a full channel, a drain that closed the gate
    /// mid-push.
    ///
    /// 🔴 The caller **keeps the registration** and withdraws only its own
    /// token. Deregistering here would destroy a healthy worker because one
    /// delivery was slow, which is exactly the failure the typed outcome exists
    /// to prevent.
    DeliveryFailed {
        /// Operator-facing reason, named by the transport that observed it.
        reason: String,
    },
}

impl Undeliverable {
    /// Whether the caller should remove this worker from the registry.
    ///
    /// `true` only for [`Undeliverable::WorkerUnreachable`]. Expressed as a
    /// method so the decision has one definition rather than a `match` repeated
    /// at every call site, where the two arms could drift apart.
    #[must_use]
    pub fn deregisters_worker(&self) -> bool {
        matches!(self, Self::WorkerUnreachable { .. })
    }

    /// The operator-facing reason, whichever variant this is.
    #[must_use]
    pub fn reason(&self) -> &str {
        match self {
            Self::WorkerUnreachable { reason } | Self::DeliveryFailed { reason } => reason,
        }
    }
}

/// The outcome of one delivery attempt to one already-chosen worker.
#[derive(Debug)]
pub enum TaskDelivery {
    /// The worker took the task. For liminal this means a correlated reply
    /// arrived; for gRPC it means the stream accepted the message.
    Delivered,
    /// The task was not placed with this worker, and [`Undeliverable`] says what
    /// the caller owes the worker's registration.
    Undeliverable(Undeliverable),
}

impl TaskDelivery {
    /// Convenience for the common "gone" refusal.
    #[must_use]
    pub fn unreachable(reason: impl Into<String>) -> Self {
        Self::Undeliverable(Undeliverable::WorkerUnreachable {
            reason: reason.into(),
        })
    }

    /// Convenience for the "alive but this attempt failed" refusal.
    #[must_use]
    pub fn failed(reason: impl Into<String>) -> Self {
        Self::Undeliverable(Undeliverable::DeliveryFailed {
            reason: reason.into(),
        })
    }
}

/// What a transport arm calls at the instant the worker HOLDS the task — the
/// gRPC stream accepted the frame; the liminal push was acknowledged — and
/// before it waits for any reply.
///
/// The one implementation in production records the attempt's lease
/// ([`super::lease_record::LeaseHandoff`]). It is a trait rather than that type
/// so the transport arms depend on the moment, not on what is done with it,
/// and so a test can hand in a recorder of its own.
#[async_trait]
pub trait DeliveryAccepted: Send + Sync {
    /// The worker holds the task.
    async fn accepted(&self);
}

/// A transport that can place one task with one already-selected worker.
///
/// 🔴 **Selection is NOT this trait's job, and that is the point.** The worker
/// arrives already chosen by the dispatcher's single selection — with its
/// `Prefer`/`Pinned` tier walk and placement cache already applied — so a
/// transport cannot select again. Two selections per placement would let the
/// spill resolve differently from the delivery, which is the defect a composite
/// dispatcher would have reintroduced.
#[async_trait]
pub trait WorkerTaskDelivery: Send + Sync + 'static {
    /// Deliver `task` to `worker`, awaiting whatever that transport's notion of
    /// delivery is.
    ///
    /// `intent` is the **caller's** abandonment condition, re-asked while a
    /// blocking transport waits for its reply. It is an argument rather than
    /// transport state because one of its terms — whether the outbox pass still
    /// holds the row's claim — is unanswerable from inside a transport, and
    /// answering it wrongly abandons every dispatcher-originated delivery at its
    /// first poll. See [`DeliveryIntent`](super::delivery_intent::DeliveryIntent)
    /// for the fact that decides it.
    ///
    /// It arrives behind an [`Arc`](std::sync::Arc) rather than as a bare
    /// `&dyn` because the blocking transport re-asks it **on a blocking thread**:
    /// the wait runs under `spawn_blocking`, which requires `'static`, and a
    /// borrow cannot cross that boundary. This is a mechanical requirement of
    /// where the predicate is evaluated, not a widening of what it may capture.
    ///
    /// `accepted` is called exactly once, at the instant the transport knows
    /// the worker HOLDS the task and before it waits for any reply — the
    /// gRPC stream accepted the frame; the liminal push was acknowledged. It is
    /// never called for an undeliverable outcome. Production records the
    /// attempt's lease there (WA-010 R3).
    ///
    /// Returns [`TaskDelivery::Delivered`] when the worker took the task, and
    /// [`TaskDelivery::Undeliverable`] otherwise — never an error for an
    /// ordinary refusal, because the caller's next move differs by variant and
    /// an error type would flatten that distinction back into a comment.
    async fn deliver(
        &self,
        worker: &WorkerHandle,
        task: &ProtoActivityTask,
        intent: &SharedDeliveryIntent,
        accepted: &dyn DeliveryAccepted,
    ) -> TaskDelivery;
}