aion_server/worker/registry/

reservation.rs

//! Claiming a worker's capacity slot for one dispatch, and the RAII holder
//! that gives it back.
//!
//! Split out of `registry.rs`: choosing a worker and claiming the slot that
//! justifies the next choice is one critical section with one owner, and
//! keeping it beside the registration and routing surfaces pushed that file
//! past the per-file length budget.

use crate::error::ServerError;

use super::capacity::worker_is_at_capacity;
use super::{
    ActivityKey, ConnectedWorkerRegistry, PoolAddress, RegistryState, WorkerHandle, WorkerId,
};

impl ConnectedWorkerRegistry {
    /// Select a worker for a dispatch AND hold one of its capacity slots, both
    /// under one acquisition of the registry lock.
    ///
    /// # Why selection alone was not enough
    ///
    /// Selection filters out a worker already at its advertised
    /// concurrency, but the increment that makes a worker *reach* that
    /// concurrency used to happen much later — at
    /// [`HeartbeatTracker::track_task`](crate::worker::heartbeat::HeartbeatTracker::track_task),
    /// once the pending entry, the superseded-attempt release and the lease
    /// record were all done. Everything between was a window in which the
    /// projection still read zero. A fan of N simultaneous dispatches onto one
    /// worker therefore had all N threads read the same pre-dispatch count, all
    /// N pass the capacity filter, and all N push — the filter was real but it
    /// was reading a number nobody had claimed yet. The worker then refused the
    /// surplus, and a refusal resolves in the TRANSPORT domain, whose ledger is
    /// deliberately bounded: routine over-dispatch would spend that budget on a
    /// worker that was merely busy and dead-letter the surplus. The saturation
    /// the fix exists to survive would have produced the outage the fix exists
    /// to prevent.
    ///
    /// Claiming the slot here closes the window by construction: the choice and
    /// the count that justifies the next choice are one critical section, so the
    /// second caller of a full worker sees a full worker.
    ///
    /// The returned [`DispatchReservation`] releases the slot when it is
    /// dropped, which is what makes every early exit between here and tracking
    /// safe without any of them knowing about capacity. Hold it across
    /// `track_task` — which takes the durable count over — and drop it after;
    /// the overlap is one extra held slot for the width of that call, which can
    /// only park a dispatch that would otherwise have raced, never admit one.
    ///
    /// # Errors
    ///
    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
    pub fn select_and_reserve(
        &self,
        namespace: &str,
        task_queue: &str,
        activity_type: &str,
        node: Option<&str>,
    ) -> Result<Option<(WorkerHandle, DispatchReservation)>, ServerError> {
        let mut state = self.state()?;
        let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
        let Some(worker) = eligible_candidates_in_rotation(&mut state, key, node)
            .into_iter()
            .next()
        else {
            return Ok(None);
        };
        let worker_id = worker.id();
        let held = state.in_flight.entry(worker_id).or_insert(0);
        *held = held.saturating_add(1);
        drop(state);
        Ok(Some((
            worker,
            DispatchReservation {
                registry: self.clone(),
                worker_id,
                committed: false,
            },
        )))
    }

    /// Claim one capacity slot on a NAMED worker, or report that it has none.
    ///
    /// The candidate-list counterpart of [`Self::select_and_reserve`]. Some
    /// dispatch paths do not pick one worker from a pool — the outbox fan-out
    /// walks a candidate list and pushes to the first that accepts — so they
    /// cannot claim at selection. They claim HERE instead, per candidate, at the
    /// moment they are about to push to that specific worker.
    ///
    /// `Ok(None)` means this worker cannot take the dispatch: it has left the
    /// registry, its capacity is not yet announced, or it is already holding
    /// every slot it advertised. The caller moves to the next candidate. That is
    /// what makes the candidate list safe to have been computed earlier — a
    /// worker that filled up between the selection and the push is refused here
    /// rather than pushed past its own admission.
    ///
    /// Like `select_and_reserve`, the returned [`DispatchReservation`] releases
    /// the slot when dropped, so a failed push gives it straight back and the
    /// next candidate is considered against a true count.
    ///
    /// # Errors
    ///
    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
    pub(in crate::worker) fn reserve_worker(
        &self,
        worker_id: WorkerId,
    ) -> Result<Option<DispatchReservation>, ServerError> {
        {
            let mut state = self.state()?;
            // UNKNOWN capacity is not a free slot. Same rule as selection: the
            // server does not invent a number for a worker that has not yet
            // said one. Copied out so the read of `workers` does not overlap the
            // write to `in_flight` below.
            let advertised = match state.workers.get(&worker_id) {
                Some(worker) => match worker.max_concurrency {
                    Some(advertised) => advertised,
                    None => return Ok(None),
                },
                None => return Ok(None),
            };
            let held = state.in_flight.get(&worker_id).copied().unwrap_or(0);
            if held >= advertised.get() {
                // Deliberately does NOT insert on the refusal path: a worker
                // that cannot take this dispatch should not gain a bookkeeping
                // entry for having been asked.
                return Ok(None);
            }
            state.in_flight.insert(worker_id, held.saturating_add(1));
        }
        Ok(Some(DispatchReservation {
            registry: self.clone(),
            worker_id,
            committed: false,
        }))
    }
}

/// One capacity slot on one worker, held from selection until it is dropped.
///
/// Returned by [`ConnectedWorkerRegistry::select_and_reserve`]. Dropping it
/// releases the slot and wakes the selection waits, exactly as a completion
/// does — so a dispatch that dies anywhere between choosing its worker and
/// handing the count to the heartbeat tracker gives the slot back without that
/// path having to know a slot existed. That is the whole point: the release is
/// owed by the type, not by every error arm.
///
/// # Handing the slot over: [`Self::commit`], and why it exists now
///
/// This type used to forbid a `commit`/`forget` escape, on the reasoning that
/// the tracker's own increment would carry the slot afterwards and the two would
/// merely overlap "for the width of one call". That was measured and it was
/// wrong in a way that mattered.
///
/// The overlap is two separate acquisitions of the registry lock — the tracker
/// increments under one, this type's `Drop` decrements under another — and
/// between them the worker's `in_flight` reads ONE HIGHER than the number of
/// dispatches it is actually holding. A concurrent leg that lands in that window
/// asks [`ConnectedWorkerRegistry::reserve_worker`], is told `held >=
/// advertised`, and is refused a slot the worker demonstrably has. Observed
/// directly: `in_flight = 5` on a worker advertising 4, on a fan sized exactly
/// to its pool — which is the normal shape, so the margin is zero and one
/// collision is enough.
///
/// So the handover is now a single act with no intermediate state:
/// [`Self::commit`] disarms this reservation and the tracker skips its own
/// increment, because the slot this reservation already holds IS the slot the
/// tracked dispatch holds. The count goes 1 → 1 rather than 1 → 2 → 1, and
/// never passes through a value that is not true.
///
/// The original worry — a missed handoff silently under-counting a busy worker
/// forever — is answered by making the transfer the ONLY way to commit: the
/// tracker takes the reservation by value, so it either commits it or drops it,
/// and a dropped reservation still releases.
#[must_use = "dropping the reservation immediately releases the slot it was taken to hold"]
#[derive(Debug)]
pub struct DispatchReservation {
    registry: ConnectedWorkerRegistry,
    worker_id: WorkerId,
    /// Set when the slot has been handed to the liveness tracker, which then
    /// owns its release. An uncommitted reservation still releases on `Drop`.
    committed: bool,
}

impl DispatchReservation {
    /// The worker whose slot this reservation holds.
    #[must_use]
    pub fn worker_id(&self) -> WorkerId {
        self.worker_id
    }

    /// Hand this slot to the liveness tracker without ever releasing it.
    ///
    /// Called by
    /// [`HeartbeatTracker::track_task`](crate::worker::heartbeat::HeartbeatTracker::track_task)
    /// and by nothing else — it takes the reservation by value, so the transfer
    /// cannot be half-done. The tracker skips its own increment in exchange:
    /// this reservation's `+1` becomes the tracked dispatch's `+1`, unchanged
    /// and never released in between, so no concurrent selection can observe a
    /// count higher than the work the worker is really holding.
    pub(in crate::worker) fn commit(mut self) {
        self.committed = true;
    }
}

impl Drop for DispatchReservation {
    fn drop(&mut self) {
        // COMMITTED means the liveness tracker took this slot over without the
        // count ever moving. Releasing here would decrement a slot the tracked
        // dispatch is still holding.
        if self.committed {
            return;
        }
        // A poisoned registry lock cannot be propagated out of `Drop`, and it
        // must not be swallowed: an unreleased slot is capacity this worker
        // never gets back, which presents to an operator as a pool that
        // silently shrinks. Report it with the worker named so the leak is
        // attributable.
        if let Err(error) = self.registry.record_dispatch_finished(self.worker_id) {
            tracing::error!(
                worker_id = ?self.worker_id,
                %error,
                "dispatch reservation could not release its capacity slot; this worker's \
                 selection capacity is now under-reported by one until it re-registers"
            );
        }
    }
}

/// The candidates for one dispatch address: eligible, id-ordered, and rotated
/// to begin at the pool's cursor, which this call advances by one.
///
/// This is the ONE selection derivation in the registry. Both
/// [`ConnectedWorkerRegistry::workers_for`] — the gRPC push dispatcher's
/// candidate list — and [`ConnectedWorkerRegistry::select_and_reserve`] — the
/// NIF bridge's wait — read it,
/// so the two cannot drift about who is dispatchable or about whose turn it is.
/// They had drifted: one rotated but never consulted eligibility, the other
/// enforced eligibility but always returned the lowest worker id, so a pool
/// served over anything but the push leg sent all of its work to one worker.
///
/// Reachability is a dispatch PRECONDITION, so it is enforced here rather than
/// discovered at push time. Selecting a worker the server cannot push to
/// produces a dispatch that can only fail, and on the liminal transport it
/// fails by consuming connection capacity — so an unreachable worker chosen
/// anyway makes its own unreachability worse. That reasoning governs both
/// selectors now, so it lives where both read it.
///
/// The id sort matters: `by_activity` holds workers in a `HashMap`, whose
/// iteration order is unspecified. Sorting first makes the rotation the sole,
/// deterministic source of ordering — true round-robin across calls with the
/// same membership, not a wobble layered on hash order.
///
/// An empty eligible set returns empty WITHOUT creating or advancing a cursor.
/// The cursor is keyed on arbitrary caller-supplied strings, and one minted for
/// a pool with nobody to rotate is a leak nothing prunes: the prune in
/// [`ConnectedWorkerRegistry::remove_worker_from_service`] fires only when an
/// activity bucket empties, and a pool that never had a worker has no bucket to
/// empty.
///
/// This does NOT count: `pool_census` and `ineligible_workers_over_tiers` answer
/// "how many" and must never advance the cursor, so they keep their own filters
/// and are deliberately not folded in here.
pub(super) fn eligible_candidates_in_rotation(
    state: &mut RegistryState,
    key: ActivityKey,
    node: Option<&str>,
) -> Vec<WorkerHandle> {
    let mut workers: Vec<WorkerHandle> = state
        .by_activity
        .get(&key)
        .map(|workers| {
            workers
                .values()
                .filter(|worker| worker_matches_node(worker, node))
                .filter(|worker| !state.dispatch_ineligible.contains_key(&worker.id))
                .filter(|worker| !worker_is_at_capacity(state, worker))
                .cloned()
                .collect()
        })
        .unwrap_or_default();
    if workers.is_empty() {
        return workers;
    }
    workers.sort_by_key(WorkerHandle::id);
    let cursor = state.rotation.entry(key).or_insert(0);
    let start = *cursor % workers.len();
    *cursor = cursor.wrapping_add(1);
    let mut rotated = Vec::with_capacity(workers.len());
    rotated.extend_from_slice(&workers[start..]);
    rotated.extend_from_slice(&workers[..start]);
    rotated
}

/// Whether a worker satisfies an optional node filter. `None` (unpinned) matches
/// every worker; `Some(node)` matches only a worker advertising that exact node
/// (NODE affinity = require). A worker with no advertised node never matches a
/// pinned dispatch.
pub(super) fn worker_matches_node(worker: &WorkerHandle, node: Option<&str>) -> bool {
    match node {
        None => true,
        Some(node) => worker.node() == Some(node),
    }
}