aion_server/worker/registry/

capacity.rs

//! Worker capacity accounting: what a registration advertises, what the
//! server records a worker as announcing later, and how many dispatches it is
//! currently holding.
//!
//! Split out of `registry.rs`: capacity is one ledger with one invariant —
//! UNKNOWN counts as full, and a worker at its advertised number is never
//! selected — and keeping it beside the registration and routing surfaces
//! pushed that file past the per-file length budget.

use std::num::NonZeroU32;

use aion_core::WorkerTransport;
use aion_proto::ProtoRegisterWorker;

use crate::error::ServerError;

use super::{ConnectedWorkerRegistry, RegistryState, WorkerHandle, WorkerId};

impl ConnectedWorkerRegistry {
    /// Record the capacity a worker announced after registering, and WAKE the
    /// selection waits.
    ///
    /// The liminal counterpart of the gRPC `RegisterWorker.max_concurrency`
    /// field. That transport's registration frame is a published wire type with
    /// no capacity field, so a liminal worker states its configured
    /// `max_concurrency` in the announcement it publishes on the reserved
    /// capabilities channel one frame later, and this applies it to the live
    /// handle selection reads.
    ///
    /// Until this lands, the worker's capacity is `None` and selection treats it
    /// as full — so this call is what makes a liminal worker dispatchable at
    /// all, which is why it wakes the parked selections the way a completion
    /// does. A worker parked on could otherwise sit until an unrelated registry
    /// change happened to fire.
    ///
    /// Returns `false` when the worker is no longer registered (a disconnect
    /// racing the announcement — benign).
    ///
    /// # Errors
    ///
    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned,
    /// or [`ServerError::Wire`] if the worker announced a capacity of zero —
    /// refused by name here exactly as the registration funnel refuses it, so
    /// the later channel cannot admit what the earlier one rejects.
    pub fn set_advertised_capacity(
        &self,
        worker_id: WorkerId,
        max_concurrency: u32,
    ) -> Result<bool, ServerError> {
        let advertised = NonZeroU32::new(max_concurrency).ok_or_else(|| ServerError::Wire {
            wire: aion_proto::WireError::backend(
                "worker announced max_concurrency 0; a worker that runs no activities could \
                 never be selected for a dispatch",
            ),
        })?;
        {
            let mut state = self.state()?;
            if !state.workers.contains_key(&worker_id) {
                return Ok(false);
            }
            if let Some(handle) = state.workers.get_mut(&worker_id) {
                handle.max_concurrency = Some(advertised);
            }
            // The selection index holds handle CLONES and is what
            // `worker_is_at_capacity` actually reads. Updating only the primary
            // map would leave the worker permanently unknown-capacity — that is,
            // permanently unselectable — while the console reported it fine.
            for workers in state.by_activity.values_mut() {
                if let Some(handle) = workers.get_mut(&worker_id) {
                    handle.max_concurrency = Some(advertised);
                }
            }
        }
        self.worker_arrived.notify_waiters();
        Ok(true)
    }

    /// Record that `worker_id` has taken one more dispatch.
    ///
    /// Called by [`HeartbeatTracker::track_task`](crate::worker::heartbeat::HeartbeatTracker::track_task)
    /// and by nothing else: the tracker takes this registry as an argument so
    /// the liveness record and the selection projection are written by one
    /// call and cannot be updated apart. A worker that is no longer registered
    /// records nothing — its entry would be unreachable capacity nobody frees.
    ///
    /// # Errors
    ///
    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
    pub(in crate::worker) fn record_dispatch_started(
        &self,
        worker_id: WorkerId,
    ) -> Result<(), ServerError> {
        let mut state = self.state()?;
        if !state.workers.contains_key(&worker_id) {
            return Ok(());
        }
        let held = state.in_flight.entry(worker_id).or_insert(0);
        *held = held.saturating_add(1);
        Ok(())
    }

    /// Record that `worker_id` has released one dispatch, and WAKE the
    /// selection waits.
    ///
    /// The wake is the half that is easy to forget. A dispatch parked because
    /// every compatible worker is at capacity is blocked on exactly this event
    /// and on nothing else — no worker is going to register, and no reachability
    /// verdict is going to change — so a completion that corrected the count
    /// without waking would leave that dispatch asleep until some unrelated
    /// registry change happened to fire. It is unconditional for the same
    /// reason [`Self::set_dispatch_ineligible`]'s is: this method cannot see
    /// which parked dispatch the freed slot serves, and a wake with nobody
    /// parked costs one waiterless `notify_waiters`.
    ///
    /// # Errors
    ///
    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
    pub(in crate::worker) fn record_dispatch_finished(
        &self,
        worker_id: WorkerId,
    ) -> Result<(), ServerError> {
        {
            let mut state = self.state()?;
            if let Some(held) = state.in_flight.get_mut(&worker_id) {
                *held = held.saturating_sub(1);
                if *held == 0 {
                    state.in_flight.remove(&worker_id);
                }
            }
        }
        self.worker_arrived.notify_waiters();
        // A FREED SLOT IS NEWS TO THE OUTBOX LOOP TOO. The selection waits are
        // woken by `worker_arrived`; a row answered "busy" and parked is waiting
        // on exactly this event and has no other way to hear it — its durable
        // re-offer rides the FAILURE backoff, which is the wrong clock for a
        // condition that clears in the time one activity takes.
        //
        // `notify_one` is what makes the non-racing case safe: a pulse with
        // nobody currently waiting STORES a permit, so a dispatcher that is
        // mid-sweep when a slot frees still finds the wake when it next selects.
        // A burst collapses into one permit, and so into one sweep. The racing
        // case — a pulse landing while the busy arm is still writing its durable
        // re-arm — is closed on the dispatcher side, by parking the row before
        // that write.
        if let Some(wake) = &self.capacity_wake {
            wake.notify_one();
        }
        Ok(())
    }

    /// How many dispatches `worker_id` is currently holding, as selection sees
    /// it.
    ///
    /// Exposed so a test can hold this against the heartbeat tracker's own
    /// per-worker count: the two are written by one call, and a divergence
    /// between them is the drift this projection would otherwise be able to
    /// hide.
    ///
    /// # Errors
    ///
    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
    pub fn in_flight_for_worker(&self, worker_id: WorkerId) -> Result<u32, ServerError> {
        Ok(self
            .state()?
            .in_flight
            .get(&worker_id)
            .copied()
            .unwrap_or(0))
    }
}

/// The capacity a registration advertises, or a typed refusal naming the field.
///
/// This half answers the WIRE question — did the worker say anything at all —
/// and it is separate from the zero check in
/// [`ConnectedWorkerRegistry::register_delivery`] because the two are different
/// findings that deserve different words. An ABSENT value is a worker built
/// against a contract it does not implement; a ZERO one is a worker claiming it
/// will run nothing, which every registration path can produce and which is
/// therefore checked at the funnel they share.
///
/// Neither is defaulted. `max_concurrency` decides how much work a process is
/// handed, and a server that supplies its own answer for it is guessing about
/// the one party that knows.
///
/// # Errors
///
/// Returns [`ServerError::Wire`] naming `max_concurrency`.
pub(super) fn advertised_capacity(
    registration: &ProtoRegisterWorker,
    transport: WorkerTransport,
) -> Result<Option<u32>, ServerError> {
    if let Some(advertised) = registration.max_concurrency {
        return Ok(Some(advertised));
    }
    match transport {
        // The gRPC registration frame HAS the field, so its absence is a worker
        // that declined to state its capacity, and the server will not invent
        // one for it.
        WorkerTransport::Grpc => Err(ServerError::Wire {
            wire: aion_proto::WireError::backend(
                "worker registration carries no max_concurrency; every worker advertises the \
                 number of activities it runs at once, and the server will not invent one for it",
            ),
        }),
        // The liminal registration frame has no such field to carry — the
        // worker is not withholding anything, the wire cannot express it. Its
        // capacity arrives on the capabilities channel one frame later; until
        // then the worker is registered with capacity UNKNOWN and selection
        // passes it over.
        //
        // NOT `#[cfg]`-gated: `WorkerTransport` lives in `aion-core` and always
        // has both variants, whatever this crate's features say. Gating the arm
        // would make this match non-exhaustive in a default build.
        WorkerTransport::Liminal => Ok(None),
    }
}

/// Whether a worker is already holding every dispatch it advertised it would
/// run at once.
///
/// Capacity is a dispatch PRECONDITION for the same reason reachability is, and
/// it is enforced at selection for the same reason: a worker chosen beyond its
/// own admission produces a dispatch that cannot start. Before this filter, the
/// only backpressure was the worker's stream channel filling — which the
/// dispatch arm reads as "merely busy" and the liveness arm reads as
/// "unreachable", so the same saturated worker was simultaneously pushed to and
/// judged unreachable, then reaped.
///
/// `>=` rather than `==`: a count that has drifted above the advertised number
/// (a redelivery joining an outstanding generation) must still exclude the
/// worker rather than fall through an equality test.
pub(super) fn worker_is_at_capacity(state: &RegistryState, worker: &WorkerHandle) -> bool {
    // UNKNOWN counts as full. A worker whose capacity has not been advertised
    // yet is one the server cannot dispatch to without inventing a number for
    // it, so it is passed over until it says. This is the safe direction: the
    // condition clears itself one frame later, whereas a guess that runs high
    // pushes a worker past its own admission and a guess that runs low
    // serializes a worker that fans.
    let Some(advertised) = worker.max_concurrency else {
        return true;
    };
    state
        .in_flight
        .get(&worker.id)
        .is_some_and(|held| *held >= advertised.get())
}