aion-server 0.31.0

//! Walking a candidate list: claiming a slot on each worker before pushing to
//! it, handing a delivered dispatch to the liveness tracker, and deciding what a
//! pass that placed nothing actually means.
//!
//! Split out of `dispatch.rs` because the capacity contract added a whole
//! concern to this path — reserve, push, hand over, release — and that file was
//! already past the size limit. Everything here is the fan-out/outbox side of
//! the same `ActivityDispatcher`; the child module can reach its private fields,
//! so nothing was widened to move.

use std::sync::Arc;

use crate::error::ServerError;
use crate::worker::delivery_intent::{DispatcherPass, OutboxClaim, SharedDeliveryIntent};
use crate::worker::envelope::CompletionToken;
use crate::worker::lease_record::{LeaseHandoff, LeaseKey, attribution_for};
use crate::worker::queue_service::taxonomy::QueueServiceReason;
use crate::worker::task_delivery::TaskDelivery;

use super::{ActivityDispatcher, DispatchOrigin, ScheduledActivity, unservable_outbox_row_reason};

/// What one walk of a candidate list did.
///
/// Three answers rather than two, because the capacity filter made "placed
/// nothing" ambiguous. A pass that skipped every candidate for want of a free
/// slot and a pass whose every candidate's stream had closed both used to
/// return the same "nothing was placed", and the caller turned that into "all
/// matching worker streams closed" — a delivery failure that spends an outbox
/// row's attempt and dead-letters it once the budget runs out. A busy pool must
/// not be priced as a dead one.
#[derive(Clone, Debug, Eq, PartialEq)]
pub(super) enum CandidateOutcome {
    /// The task was placed with a worker.
    Delivered,
    /// No candidate had a free slot: the pool is working, not gone.
    AllFull,
    /// Every candidate's stream was already closed.
    NoLiveStream,
    /// The transport proved the task can never be placed — its frame exceeds
    /// the bound every worker connection shares — so the walk stopped at the
    /// first candidate that said so. Terminal for the dispatch: no further
    /// candidate, no park, no retry ([`Undeliverable::Unservable`]).
    Unservable {
        /// The transport's own reason, naming the frame, the bound and the key.
        reason: String,
    },
}

impl CandidateOutcome {
    /// Whether the task was placed, for the tier walk, which cares about
    /// nothing else — a tier that could not place moves to the next tier
    /// whichever way it failed.
    pub(super) const fn is_delivered(&self) -> bool {
        matches!(self, Self::Delivered)
    }
}

impl ActivityDispatcher {
    /// How an outbox row's selection miss resolves — and, crucially, whether it
    /// costs the row one of its attempts.
    ///
    /// # A busy pool is not an unservable one
    ///
    /// `workers_for` now excludes a worker that is at its advertised capacity or has
    /// not yet announced one, so an empty candidate list no longer means what it
    /// used to. Before that filter it meant "nobody serves this"; it now also means
    /// "everybody who serves this is busy". Returning the unservable reason for the
    /// second case spends one of the row's attempts, backs it off, and dead-letters
    /// it if the pool stays busy across the budget — work discarded because a worker
    /// was working, which is the exact harm this landing exists to remove and the
    /// precise opposite of what [`QueueServiceReason::PollersAtCapacity`]'s own
    /// sentence promises an operator.
    ///
    /// So the busy classes resolve as [`ServerError::WorkerBusy`], which the outbox
    /// dispatcher already treats as ATTEMPT-NEUTRAL: the row is re-armed pending
    /// with its budget untouched and no dead-letter eligibility, carrying the
    /// queue-service reason so an operator reads why it is waiting rather than
    /// merely that it is. No park (an outbox row must not park in the dispatch call
    /// — that call has to return, or the row's own mechanism never runs), no new
    /// clock, and no default: if the pool stays busy the row stays pending, exactly
    /// as a parked bridge dispatch does.
    ///
    /// Everything else — no declaration, no poller, incompatible pollers, lost
    /// reachability — keeps today's behaviour and spends the attempt, because none
    /// of those resolve by waiting for a slot.
    pub(super) fn outbox_row_selection_refusal(
        activity: &ScheduledActivity,
        reported: Option<QueueServiceReason>,
    ) -> ServerError {
        let busy = match reported {
            Some(QueueServiceReason::PollersAtCapacity) => Some("at its advertised concurrency"),
            Some(QueueServiceReason::PollersCapacityUnannounced) => {
                Some("connected but has not announced its capacity")
            }
            _ => None,
        };
        match busy {
            Some(condition) => ServerError::worker_busy(
                activity.task_queue.clone(),
                format!(
                    "every worker serving activity `{}` in task queue `{}` is {condition}; the row \
                 is re-queued without spending an attempt",
                    activity.activity_type, activity.task_queue,
                ),
            ),
            None => ServerError::worker_dispatch(
                activity.namespace.clone(),
                activity.activity_type.clone(),
                unservable_outbox_row_reason(reported, &activity.task_queue),
            ),
        }
    }

    /// Try each candidate in order, pushing the task to the first live stream.
    /// Returns `Ok(Some(()))` on a delivered task, `Ok(None)` when every candidate
    /// stream was already closed (deregistered as it went). An empty candidate
    /// list returns `Ok(None)` so callers can treat it as "no live worker here".
    pub(super) async fn send_to_candidates(
        &self,
        activity: &ScheduledActivity,
        candidates: Vec<crate::worker::registry::WorkerHandle>,
        span_fields: &tracing::Span,
    ) -> Result<CandidateOutcome, ServerError> {
        let run_id = activity.require_run_id()?;
        // The run and the attempt are what tell a REDELIVERY of this attempt
        // apart from a genuine retry or a new execution generation: a
        // redelivery adds a sibling authorization beside the one the first
        // worker is still holding, instead of replacing it.
        let completion_token = self.completion_fences.issue(
            &activity.workflow_id,
            run_id,
            &activity.activity_id,
            activity.attempt,
        )?;
        let task = activity.to_task(&completion_token)?;
        let lease_key = LeaseKey {
            workflow_id: activity.workflow_id.clone(),
            run_id: run_id.clone(),
            activity_id: activity.activity_id.clone(),
            attempt: activity.attempt,
        };
        // Set when a candidate is passed over because it has no free slot. It is
        // what stops "every candidate is BUSY" being reported as "every
        // candidate's stream is CLOSED" — two conditions with opposite
        // consequences for an outbox row's attempt budget.
        let mut full_candidates = false;
        for worker in candidates {
            if let Err(error) = self
                .drain_state
                .ensure_accepting(&activity.namespace, &activity.activity_type)
            {
                // Withdraw the authorization THIS pass minted, and only that
                // one: a sibling token held by a worker already executing the
                // same attempt must survive the drain refusal.
                self.completion_fences.revoke(
                    &activity.workflow_id,
                    &activity.activity_id,
                    &completion_token,
                )?;
                return Err(error);
            }
            // CLAIM THIS WORKER'S SLOT BEFORE PUSHING TO IT.
            //
            // The candidate list was computed earlier, by `workers_for`, and a
            // list is a snapshot: by the time this loop reaches a given
            // candidate that worker may already be full — filled by a sibling
            // leg of this very fan. Claiming here, per candidate, is what makes
            // the snapshot safe to act on. A worker with no slot left is skipped
            // rather than pushed past its own admission, and the reservation is
            // held across the push so a concurrent pass cannot claim the same
            // slot.
            //
            // On a failed push it is dropped and the slot goes straight back, so
            // the next candidate is considered against a true count.
            let Some(reservation) = self.registry.reserve_worker(worker.id())? else {
                full_candidates = true;
                tracing::debug!(
                    namespace = %activity.namespace,
                    task_queue = %activity.task_queue,
                    activity_type = %activity.activity_type,
                    worker_id = ?worker.id(),
                    "candidate worker has no free slot; trying the next candidate"
                );
                continue;
            };
            span_fields.record("worker_id", format!("{:?}", worker.id()));
            // The abandonment condition is chosen by the DECLARED origin — see
            // `delivery_intent_for`.
            let intent = self.delivery_intent_for(activity, &worker);
            // Route by the SELECTED WORKER'S delivery. Selection already
            // happened — this loop walks candidates the caller chose — so no
            // transport re-selects, and the spill cannot resolve differently
            // from the delivery (#52 R1).
            // The lease this candidate's acceptance will record, armed on the
            // fences BEFORE the push so a completion racing the append waits
            // for it; settled by the handoff whether the push lands or not.
            let handoff = LeaseHandoff::arm(
                self.lease_recorder.clone(),
                lease_key.clone(),
                attribution_for(&worker),
                self.completion_fences.clone(),
                completion_token.clone(),
                // No per-attempt timer sits above this leg. These dispatches
                // come off the durable outbox, whose own attempt budget,
                // backoff and dead-letter bound them; the engine's authored
                // per-attempt bound belongs to the retry loop on the engine
                // seam, which is a different caller. Passing an unwatched
                // signal states that, rather than handing over a clock nobody
                // started.
                aion::LeaseSignal::none(),
            )?;
            // TRACKED BEFORE THE PUSH — the bridge's `track_then_send` ordering,
            // and it is load-bearing on this path for a reason the bridge never
            // had to face.
            //
            // `deliver_to` does NOT mean the same thing on both transports.
            // `GrpcTaskDelivery::deliver` returns `Delivered` once the frame is
            // queued and the result comes back later on the worker's stream. The
            // LIMINAL delivery pushes, waits for the worker's reply INLINE, and
            // routes the completion before it returns. So tracking after a
            // delivery meant, on liminal, creating a tracker entry for work that
            // had already finished — and the completion that would have removed
            // it had already gone by. Nothing else on that path removes one, so
            // it was a permanent +1 per delivery: a worker advertising N went
            // unselectable after N rows, the fan then re-armed forever under the
            // attempt-neutral busy arm instead of failing, and `in_flight_count`
            // never returned to zero so graceful drain sat out its timeout for
            // the life of the process.
            //
            // It also closes a narrower gRPC race with the same shape: the
            // acceptance appends the durable lease before this returned, so a
            // fast activity's result could be processed in that window, find
            // nothing to untrack, and then be tracked afterwards.
            //
            // Tracking first means the entry always exists before anything can
            // complete it, on either transport, and the completion path is the
            // only thing that removes it.
            // ONE ACT. The reservation goes INTO the tracker, which commits it
            // and skips its own increment — the count goes 1 → 1 rather than
            // 1 → 2 → 1. It used to be two: track (+1 under one lock) then drop
            // the reservation (-1 under another), and a concurrent leg landing
            // between them was refused a slot the worker demonstrably had.
            //
            // On failure the tracker has already released it, so there is
            // nothing left here to hold or drop.
            self.track_before_push(&worker, activity, &completion_token, reservation);
            let outcome = self.deliver_to(&worker, &task, &intent, &handoff).await;
            match outcome {
                TaskDelivery::Delivered => return Ok(CandidateOutcome::Delivered),
                TaskDelivery::Undeliverable(undeliverable) => {
                    tracing::warn!(
                        namespace = %activity.namespace,
                        task_queue = %activity.task_queue,
                        activity_type = %activity.activity_type,
                        workflow_id = %activity.workflow_id,
                        activity_id = %activity.activity_id,
                        worker_id = ?worker.id(),
                        deregistered = undeliverable.deregisters_worker(),
                        reason = %undeliverable.reason(),
                        "activity delivery to selected worker did not place the task"
                    );
                    // The decision has ONE definition, on the type. A worker the
                    // transport says is GONE is removed; a delivery that failed
                    // while the worker is ALIVE leaves the registration standing
                    // — which is the whole of #52: this line used to run
                    // unconditionally, destroying a correctly-selected liminal
                    // worker for the crime of not carrying a gRPC sender.
                    // NOTHING WAS PLACED, so the tracker entry this pass made
                    // must go back too. The reservation is released a few lines
                    // down and the tracked entry gets the same treatment: an
                    // entry left behind for a delivery that never landed is a
                    // slot the server never gets back, which is exactly the leak
                    // the ordering above exists to prevent — just on the failure
                    // arm instead of the success one.
                    self.untrack_undelivered(&worker, activity);
                    // TERMINAL, and it ends the walk here: every connection
                    // this server holds was built by the one supervisor with
                    // the one bound the frame just exceeded, so the next
                    // candidate would refuse identically — after recording one
                    // more lease. The pass withdraws its own token exactly as
                    // the exhausted walk below does, and hands the class up.
                    if undeliverable.is_unservable() {
                        return self.settle_unservable(
                            activity,
                            &completion_token,
                            undeliverable.reason(),
                        );
                    }
                    // The slot goes back before the next candidate is
                    // considered — explicitly, so the release is visible beside
                    // the failure it belongs to rather than left to scope exit.
                    // The slot itself is already back: either the tracker took
                    // it over and `untrack_undelivered` above returned it, or
                    // the handover failed and the tracker released it there.
                    if undeliverable.deregisters_worker() {
                        self.registry.deregister(worker.id())?;
                    }
                }
            }
        }
        // Every candidate stream was closed, so this pass placed nothing and
        // withdraws its OWN token. It must not remove the execution site's
        // generation outright: when this pass was a redelivery, the first
        // worker is still alive, still executing, and still holding the token
        // it was given — and its finished result is the truth.
        self.completion_fences.revoke(
            &activity.workflow_id,
            &activity.activity_id,
            &completion_token,
        )?;
        // BUSY OUTRANKS CLOSED. If any candidate was passed over for want of a
        // slot, this pass placed nothing because the pool is working, not
        // because it is gone — and the caller must not price that as a delivery
        // failure. Where both happened, the closed stream has already
        // deregistered its worker, so the next pass sees a smaller pool; saying
        // "busy" costs one attempt-neutral re-queue and never a dead-letter.
        Ok(if full_candidates {
            CandidateOutcome::AllFull
        } else {
            CandidateOutcome::NoLiveStream
        })
    }

    /// The abandonment condition for one candidate, chosen by the DECLARED
    /// origin rather than by the presence of a field. An engine-scheduled
    /// activity holds no row claim and must not consult one — a key that was
    /// never begun is indistinguishable from a released one, so asking would
    /// abandon every engine dispatch at its first poll. An outbox row's pass
    /// does hold a claim, and losing it must stop the wait.
    fn delivery_intent_for(
        &self,
        activity: &ScheduledActivity,
        worker: &crate::worker::registry::WorkerHandle,
    ) -> SharedDeliveryIntent {
        match &activity.origin {
            DispatchOrigin::Engine => Arc::new(DispatcherPass::new(
                self.delivery_gate.clone(),
                self.registry.clone(),
                worker.id(),
            )),
            DispatchOrigin::OutboxRow { dispatch_key } => Arc::new(OutboxClaim::new(
                self.delivery_gate.clone(),
                dispatch_key.clone(),
                self.registry.clone(),
                worker.id(),
            )),
        }
    }

    /// End the pass on the first UNSERVABLE candidate: withdraw the token this
    /// pass minted — exactly as the exhausted walk does — and hand the class up
    /// as [`CandidateOutcome::Unservable`] so the dispatcher settles it
    /// terminally instead of walking on or parking.
    fn settle_unservable(
        &self,
        activity: &ScheduledActivity,
        completion_token: &CompletionToken,
        reason: &str,
    ) -> Result<CandidateOutcome, ServerError> {
        self.completion_fences.revoke(
            &activity.workflow_id,
            &activity.activity_id,
            completion_token,
        )?;
        Ok(CandidateOutcome::Unservable {
            reason: reason.to_owned(),
        })
    }

    /// Remove the tracker entry this pass made, when the push did not land.
    ///
    /// The mirror of [`Self::track_before_push`]. Tracking happens first so the
    /// entry always exists before anything can complete it; that makes it this
    /// arm's job to take it back when nothing was placed, exactly as the
    /// reservation is dropped on the same arm.
    ///
    /// Idempotent by key: `clear_completed_task_tracking` reports whether it
    /// actually retired an entry, so a completion that raced this and got there
    /// first simply makes this a no-op rather than a second decrement.
    fn untrack_undelivered(
        &self,
        worker: &crate::worker::registry::WorkerHandle,
        activity: &ScheduledActivity,
    ) {
        let Some(tracker) = &self.heartbeat_tracker else {
            return;
        };
        let _ = crate::worker::bridge::clear_completed_task_tracking(
            tracker,
            &self.registry,
            worker.id(),
            &activity.workflow_id,
            &activity.activity_id,
        );
    }

    /// Track a dispatch BEFORE it is pushed, which is what makes it COUNT
    /// against the worker's advertised capacity.
    ///
    /// Consumes the reservation, which is the point: the tracker commits that
    /// slot and skips its own increment, so the worker's count never passes
    /// through a value higher than the work it is really holding. Every exit
    /// from here has either committed the reservation or released it.
    ///
    /// Bookkeeping, so it never fails the delivery: the task is placed and the
    /// worker is executing it whatever happens here. But it is never silent
    /// either — an untracked dispatch is a slot the server does not know it has
    /// spent, and the next selection will over-push on the strength of it.
    ///
    /// A dispatcher with no tracker wired says so on every delivery. That is an
    /// isolated-test wiring, and the alternative — counting nothing, quietly —
    /// is exactly the shape of the defect this arm was added to fix.
    fn track_before_push(
        &self,
        worker: &crate::worker::registry::WorkerHandle,
        activity: &ScheduledActivity,
        completion_token: &CompletionToken,
        reservation: crate::worker::registry::DispatchReservation,
    ) {
        let Some(tracker) = &self.heartbeat_tracker else {
            // The reservation is the caller's only hold on this slot and this
            // dispatcher counts nothing, so give it straight back rather than
            // holding a slot no tracker will ever release.
            drop(reservation);
            tracing::error!(
                namespace = %activity.namespace,
                task_queue = %activity.task_queue,
                activity_type = %activity.activity_type,
                workflow_id = %activity.workflow_id,
                activity_id = %activity.activity_id,
                worker_id = ?worker.id(),
                "outbox dispatcher has no heartbeat tracker wired: this delivery is NOT counted \
                 against the worker's advertised capacity, and selection may push past it"
            );
            return;
        };
        if let Err(error) = tracker.track_task(
            worker.id(),
            crate::worker::heartbeat::InFlightActivity {
                workflow_id: activity.workflow_id.clone(),
                activity_id: activity.activity_id.clone(),
                attempt: activity.attempt,
                completion_token: completion_token.clone(),
            },
            std::time::Instant::now(),
            &self.registry,
            // HANDED OVER, not re-counted: the tracker commits this slot and
            // skips its own increment, so the count never passes through a
            // value higher than the work the worker holds.
            Some(reservation),
        ) {
            tracing::error!(
                namespace = %activity.namespace,
                activity_type = %activity.activity_type,
                workflow_id = %activity.workflow_id,
                activity_id = %activity.activity_id,
                worker_id = ?worker.id(),
                %error,
                "dispatch could not be tracked; the worker will execute it but the server's \
                 capacity count and expiry sweep will not see it"
            );
        }
    }
}