aion_server/worker/registry/

census.rs

//! The registry's counting reads: the pool census one dispatch address is
//! judged by, and the excluded-worker headcount a selection refusal quotes.
//!
//! Split out of `registry.rs`: both walk the pool under the SAME lock the
//! selection they explain was taken under, and keeping them beside the
//! registration and routing surfaces pushed that file past the per-file
//! length budget.

use std::collections::HashMap;
use std::time::Duration;

use crate::error::ServerError;
use crate::worker::heartbeat::DispatchExclusion;
use crate::worker::queue_service::PoolCensus;

use super::capacity::worker_is_at_capacity;
use super::reservation::worker_matches_node;
use super::{ActivityKey, ConnectedWorkerRegistry, PoolAddress};

impl ConnectedWorkerRegistry {
    /// How many workers that selection COULD have chosen are currently excluded
    /// from it by the liveness verdict (#197 R3).
    ///
    /// `tiers` is the ordered sequence of node filters the selection actually
    /// walked, and taking it as an argument — rather than re-deriving one here
    /// — is the whole point of the method. Selection over a `Pinned{L}`
    /// namespace walks the required labels and NEVER spills to a `None`
    /// any-node tier; a census taken over the row's own node instead would pass
    /// `None`, match every worker in the pool, and report an unlabelled worker
    /// that was never a candidate as the reason the row found nobody. The
    /// refusal built on that count then tells an operator not to start the
    /// labelled worker that is the only remedy.
    ///
    /// Counted as a UNION over the tiers and over DISTINCT workers: the pool is
    /// walked once and a worker admissible to more than one tier is counted
    /// once, so the number is a headcount rather than a sum of overlapping
    /// matches.
    ///
    /// Read under the SAME lock and with the SAME `worker_matches_node` filter
    /// [`Self::select_and_reserve`] applies, so a refusal quoting this count
    /// describes the fleet selection actually saw.
    ///
    /// # Errors
    ///
    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
    pub fn ineligible_workers_over_tiers(
        &self,
        namespace: &str,
        task_queue: &str,
        activity_type: &str,
        tiers: &[Option<String>],
    ) -> Result<usize, ServerError> {
        let state = self.state()?;
        let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
        Ok(state.by_activity.get(&key).map_or(0, |workers| {
            workers
                .values()
                .filter(|worker| state.dispatch_ineligible.contains_key(&worker.id))
                .filter(|worker| {
                    tiers
                        .iter()
                        .any(|tier| worker_matches_node(worker, tier.as_deref()))
                })
                .count()
        }))
    }

    /// Census the live fleet for one dispatch address (R1).
    ///
    /// Taken under the SAME lock discipline as [`Self::select_and_reserve`] and read
    /// immediately after a selection miss, so the taxonomy verdict describes
    /// the fleet selection actually saw. Three nested counts — the pool, the
    /// activity coverage within it, the node coverage within that — are what
    /// separate `NO_LIVE_POLLERS` from `POLLERS_INCOMPATIBLE`.
    ///
    /// Two further counts split the compatible workers by dispatch eligibility,
    /// which is what separates a pool that is genuinely served from one whose
    /// workers are all excluded, and — within that — an exclusion that clears
    /// itself from one that does not (`POLLERS_UNREACHABLE`). They are taken
    /// here, under this same lock, rather than by a second read: eligibility
    /// can change between two acquisitions, and a verdict assembled from two
    /// readings would describe a fleet that never existed at one instant.
    ///
    /// `last_compatible_poller_age` is zero while a compatible worker is
    /// connected, the elapsed time since the most recent compatible departure
    /// when one has left, and `None` when this server has never had one.
    ///
    /// # Errors
    ///
    /// Returns [`ServerError::LockPoisoned`] if the registry lock is poisoned.
    pub fn pool_census(
        &self,
        namespace: &str,
        task_queue: &str,
        activity_type: &str,
        node: Option<&str>,
    ) -> Result<PoolCensus, ServerError> {
        let state = self.state()?;
        let key = ActivityKey::new(PoolAddress::new(namespace, task_queue), activity_type);
        let workers_in_pool = state
            .workers
            .values()
            .filter(|worker| {
                worker.task_queue == task_queue && worker.namespaces.contains(namespace)
            })
            .count();
        let serving = state.by_activity.get(&key);
        let workers_serving_activity = serving.map_or(0, HashMap::len);
        let compatible: Vec<_> = serving.map_or_else(Vec::new, |workers| {
            workers
                .values()
                .filter(|worker| worker_matches_node(worker, node))
                .collect()
        });
        let compatible_workers = compatible.len();
        // Counted here, under the SAME lock as the selection this census
        // explains, because eligibility can change between two lock
        // acquisitions and a verdict assembled from two readings would describe
        // a fleet that never existed at one instant.
        let eligible_compatible_workers = compatible
            .iter()
            .filter(|worker| !state.dispatch_ineligible.contains_key(&worker.id))
            .filter(|worker| !worker_is_at_capacity(&state, worker))
            .count();
        let compatible_workers_reachability_lost = compatible
            .iter()
            .filter(|worker| {
                matches!(
                    state.dispatch_ineligible.get(&worker.id),
                    Some(&DispatchExclusion::ReachabilityLost)
                )
            })
            .count();
        // Counted under the SAME lock as `eligible_compatible_workers` above and
        // for the same reason: a busy pool and an unreachable one are different
        // conditions with opposite remedies, and a verdict assembled from two
        // readings would describe a fleet that never existed at one instant.
        // DISJOINT counts, deliberately. `worker_is_at_capacity` is true of an
        // unannounced worker as well as a full one — that is what makes
        // selection fail closed — so counting "at capacity" with that predicate
        // alone would fold an idle, mute worker into the busy population and the
        // taxonomy would report a healthy-but-loaded pool for a worker that has
        // simply never said anything. The census is what the operator sentence
        // is derived from, so the split has to happen here, not there.
        let compatible_workers_capacity_unannounced = compatible
            .iter()
            .filter(|worker| worker.max_concurrency.is_none())
            .count();
        let compatible_workers_at_capacity = compatible
            .iter()
            .filter(|worker| worker.max_concurrency.is_some())
            .filter(|worker| worker_is_at_capacity(&state, worker))
            .count();
        let last_compatible_poller_age = if compatible_workers > 0 {
            Some(Duration::ZERO)
        } else {
            state
                .last_departure
                .get(&key)
                .and_then(|by_node| {
                    by_node
                        .iter()
                        .filter(|(departed_node, _)| match node {
                            None => true,
                            Some(node) => departed_node.as_deref() == Some(node),
                        })
                        .map(|(_, departed_at)| *departed_at)
                        .max()
                })
                .map(|departed_at| departed_at.elapsed())
        };
        Ok(PoolCensus {
            workers_in_pool,
            workers_serving_activity,
            compatible_workers,
            eligible_compatible_workers,
            compatible_workers_reachability_lost,
            compatible_workers_at_capacity,
            compatible_workers_capacity_unannounced,
            last_compatible_poller_age,
        })
    }
}