aion-server 0.31.0

//! The `RegisterAck` frame the server guarantees as the first frame on every
//! successful worker response stream, and the registry read-back that fills in
//! the concurrency it echoes.
//!
//! Split out of `worker_grpc.rs`: building the ack and proving the number it
//! carries is one concern with one caller, and keeping it beside the stream
//! handler pushed that file past the per-file length budget.

use aion_proto::generated;
use tonic::Status;

use crate::worker::registry::WorkerId;

use super::status_from_server_error;

/// The concurrency the server RECORDED for this registration, read back off the
/// registered handle.
///
/// Read from the REGISTRY rather than copied from the frame the worker sent, so
/// what the ack tells the worker is what selection will actually act on.
/// Echoing the request back would agree with the worker by construction and
/// could never surface a drift — which is the one thing the echo exists for.
///
/// # Errors
///
/// Returns the registry's own error, or `Internal` when the worker left the
/// registry between being admitted and being acked (a disconnect racing its own
/// registration; there is nothing left to ack).
pub(super) fn recorded_capacity(
    registry: &crate::worker::ConnectedWorkerRegistry,
    worker_id: WorkerId,
) -> Result<std::num::NonZeroU32, Status> {
    registry
        .worker_by_id(worker_id)
        .map_err(|error| status_from_server_error(&error))?
        .ok_or_else(|| {
            Status::internal("worker left the registry before its registration was acked")
        })?
        .max_concurrency()
        // A gRPC registration cannot reach this arm: admission refuses a gRPC
        // frame that carries no `max_concurrency`, so the handle it produced
        // always has one. Unknown capacity exists only on the liminal
        // transport, which has no `RegisterAck` to echo into. Reported rather
        // than unwrapped, because "cannot happen" and "silently acks a wrong
        // number" are one edit apart.
        .ok_or_else(|| {
            Status::internal(
                "gRPC worker was registered with no advertised capacity; the registration \
                 funnel that refuses this was bypassed",
            )
        })
}

/// Build the positive registration acknowledgement frame — the guaranteed
/// first frame on every successful worker response stream.
pub(super) fn register_ack_frame(
    worker_id: WorkerId,
    namespace: &str,
    heartbeat_window: std::time::Duration,
    accepted_max_concurrency: std::num::NonZeroU32,
) -> generated::ServerToWorker {
    generated::ServerToWorker {
        message: Some(generated::server_to_worker::Message::RegisterAck(
            generated::RegisterAck {
                worker_id: worker_id.value(),
                namespace: namespace.to_owned(),
                heartbeat_window_ms: u64::try_from(heartbeat_window.as_millis())
                    .unwrap_or(u64::MAX),
                // Read back off the REGISTERED handle rather than copied from
                // the frame, so what the worker is told is what selection will
                // actually act on. Echoing the request instead would agree with
                // the worker by construction and could never surface a drift.
                accepted_max_concurrency: accepted_max_concurrency.get(),
            },
        )),
    }
}