rocketmq-transport 1.0.0

// Copyright 2023 The RocketMQ Rust Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::fmt;
use std::net::IpAddr;
use std::sync::Arc;
use std::sync::Mutex;

use crate::contract::TransportContractViolation;
use crate::dispatch::DeferredAdmission;
use crate::dispatch::DeferredWaitLimits;
use crate::dispatch::EmbeddedCaller;
use rocketmq_protocol::code::request_code::RequestCode;
use rocketmq_runtime::BudgetCapacity;
use rocketmq_runtime::BudgetClass;
use rocketmq_runtime::BudgetDimension;
use rocketmq_runtime::BudgetLimit;
pub use rocketmq_runtime::FullPolicy;
use rocketmq_runtime::ResourceBudget;
use rocketmq_runtime::ResourceBudgetTree;
use rocketmq_runtime::ResourcePermit;
use rocketmq_runtime::RuntimeContractViolation;

const ESTABLISHED_CONNECTION_RETAINED_BYTES: usize = 16 * 1024;
const HANDSHAKE_RETAINED_BYTES: usize = 64 * 1024;

#[must_use]
pub(crate) const fn estimated_connection_retained_bytes() -> usize {
    ESTABLISHED_CONNECTION_RETAINED_BYTES
}

#[must_use]
pub(crate) const fn estimated_handshake_retained_bytes() -> usize {
    HANDSHAKE_RETAINED_BYTES
}

/// Simultaneous item and retained-byte limit for one transport resource.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ResourceLimit {
    pub count: usize,
    pub bytes: usize,
}

/// Explicit limits for every resource admitted by the transport.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AdmissionLimits {
    pub connections: ResourceLimit,
    pub handshakes: ResourceLimit,
    pub partial_frames: ResourceLimit,
    pub inflight: ResourceLimit,
    pub queued: ResourceLimit,
    pub processors: ResourceLimit,
    pub per_ip: ResourceLimit,
    pub per_tenant: ResourceLimit,
    pub per_session: ResourceLimit,
    pub control_reserve: ResourceLimit,
    pub max_scope_keys: usize,
}

impl Default for AdmissionLimits {
    fn default() -> Self {
        Self {
            connections: ResourceLimit {
                count: 16_384,
                bytes: 256 * 1024 * 1024,
            },
            handshakes: ResourceLimit {
                count: 1_024,
                bytes: 64 * 1024 * 1024,
            },
            partial_frames: ResourceLimit {
                count: 4_096,
                bytes: 256 * 1024 * 1024,
            },
            inflight: ResourceLimit {
                count: 65_536,
                bytes: 256 * 1024 * 1024,
            },
            queued: ResourceLimit {
                count: 65_536,
                bytes: 256 * 1024 * 1024,
            },
            processors: ResourceLimit {
                count: 4_096,
                bytes: 128 * 1024 * 1024,
            },
            per_ip: ResourceLimit {
                count: 4_096,
                bytes: 64 * 1024 * 1024,
            },
            per_tenant: ResourceLimit {
                count: 16_384,
                bytes: 128 * 1024 * 1024,
            },
            per_session: ResourceLimit {
                count: 4_096,
                bytes: 64 * 1024 * 1024,
            },
            control_reserve: ResourceLimit {
                count: 64,
                bytes: 1024 * 1024,
            },
            max_scope_keys: 16_384,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AdmissionResource {
    Connection,
    Handshake,
    PartialFrame,
    Inflight,
    Queued,
    Processor,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AdmissionClass {
    Data,
    Control,
}

impl AdmissionClass {
    pub fn for_request_code(code: i32) -> Self {
        match RequestCode::from(code) {
            RequestCode::HeartBeat
            | RequestCode::RegisterBroker
            | RequestCode::UnregisterBroker
            | RequestCode::GetRouteinfoByTopic
            | RequestCode::GetBrokerClusterInfo
            | RequestCode::GetBrokerRuntimeInfo => Self::Control,
            _ => Self::Data,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AdmissionScope {
    ip: IpAddr,
    tenant: Option<u64>,
    session: Option<u64>,
}

impl AdmissionScope {
    pub fn new(ip: IpAddr) -> Self {
        Self {
            ip,
            tenant: None,
            session: None,
        }
    }

    pub fn with_tenant(mut self, tenant: u64) -> Self {
        self.tenant = Some(tenant);
        self
    }

    pub fn with_session(mut self, session: u64) -> Self {
        self.session = Some(session);
        self
    }
}

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ResourceSnapshot {
    pub current_count: usize,
    pub current_bytes: usize,
    pub rejected_count: usize,
}

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct AdmissionSnapshot {
    pub connections: ResourceSnapshot,
    pub handshakes: ResourceSnapshot,
    pub partial_frames: ResourceSnapshot,
    pub inflight: ResourceSnapshot,
    pub queued: ResourceSnapshot,
    pub processors: ResourceSnapshot,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AdmissionEventOutcome {
    Acquired,
    Rejected,
    Released,
}

/// Low-cardinality metric event. Scope identities are deliberately omitted.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AdmissionEvent {
    pub resource: AdmissionResource,
    outcome: AdmissionEventOutcome,
    pub bytes: usize,
}

#[cfg(test)]
impl AdmissionEvent {
    pub(crate) const fn is_release(self) -> bool {
        matches!(self.outcome, AdmissionEventOutcome::Released)
    }
}

/// Result of a capacity-controlled Transport admission attempt.
#[must_use]
pub enum AdmissionOutcome<T> {
    /// Capacity was acquired and ownership is returned to the caller.
    Acquired(T),
    /// Capacity was unavailable under the configured policy.
    Rejected(AdmissionRejection),
}

impl<T> AdmissionOutcome<T> {
    pub(crate) fn into_result(self) -> Result<T, AdmissionRejection> {
        match self {
            Self::Acquired(value) => Ok(value),
            Self::Rejected(rejection) => Err(rejection),
        }
    }
}

impl fmt::Display for AdmissionRejection {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("transport admission rejected")
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AdmissionRejection {
    resource: AdmissionResource,
    policy: FullPolicy,
}

impl AdmissionRejection {
    /// Returns the resource whose bounded capacity rejected the request.
    #[must_use]
    pub const fn resource(&self) -> AdmissionResource {
        self.resource
    }

    /// Returns the configured full-capacity disposition.
    #[must_use]
    pub fn policy(&self) -> FullPolicy {
        self.policy
    }
}

impl From<RuntimeContractViolation> for TransportContractViolation {
    fn from(error: RuntimeContractViolation) -> Self {
        Self::AdmissionBudget(error)
    }
}

struct GlobalBudgets {
    connections: ResourceBudget,
    handshakes: ResourceBudget,
    partial_frames: ResourceBudget,
    inflight: ResourceBudget,
    queued: ResourceBudget,
    processors: ResourceBudget,
}

impl GlobalBudgets {
    fn new(limits: AdmissionLimits, process_budget: &ResourceBudget) -> Result<Self, TransportContractViolation> {
        let resources = [
            limits.connections,
            limits.handshakes,
            limits.partial_frames,
            limits.inflight,
            limits.queued,
            limits.processors,
        ];
        let root_count = resources
            .iter()
            .fold(0usize, |total, limit| total.saturating_add(limit.count));
        let requested_root_bytes = resources
            .iter()
            .fold(0usize, |total, limit| total.saturating_add(limit.bytes));
        let process_capacity = process_budget.limit().capacity;
        let root_count = root_count.min(process_capacity.count).max(1);
        let root_bytes = requested_root_bytes.min(process_capacity.bytes).max(1);
        let reservable_resources = [
            limits.connections,
            limits.handshakes,
            limits.inflight,
            limits.queued,
            limits.processors,
        ];
        let reserve_count = reservable_resources
            .iter()
            .fold(0usize, |total, limit| {
                total.saturating_add(effective_reserve(limits.control_reserve.count, limit.count))
            })
            .min(root_count);
        let requested_reserve_bytes = reservable_resources.iter().fold(0usize, |total, limit| {
            total.saturating_add(effective_reserve(limits.control_reserve.bytes, limit.bytes))
        });
        let reserve_bytes = effective_reserve(requested_reserve_bytes, root_bytes);
        let root = process_budget.child(
            "transport",
            BudgetLimit::new(root_count, root_bytes, FullPolicy::Reject)
                .with_control_reserve(BudgetCapacity::new(reserve_count, reserve_bytes)),
        )?;
        Ok(Self {
            connections: global_budget(
                &root,
                "connections",
                limits.connections,
                limits.control_reserve,
                FullPolicy::CloseSlowConsumer,
            )?,
            handshakes: global_budget(
                &root,
                "handshakes",
                limits.handshakes,
                limits.control_reserve,
                FullPolicy::CloseSlowConsumer,
            )?,
            partial_frames: global_budget(
                &root,
                "partial-frames",
                limits.partial_frames,
                ResourceLimit { count: 0, bytes: 0 },
                FullPolicy::CloseSlowConsumer,
            )?,
            inflight: global_budget(
                &root,
                "inflight",
                limits.inflight,
                limits.control_reserve,
                FullPolicy::Reject,
            )?,
            queued: global_budget(
                &root,
                "queued",
                limits.queued,
                limits.control_reserve,
                FullPolicy::Reject,
            )?,
            processors: global_budget(
                &root,
                "processors",
                limits.processors,
                limits.control_reserve,
                FullPolicy::Reject,
            )?,
        })
    }

    fn get(&self, resource: AdmissionResource) -> ResourceBudget {
        match resource {
            AdmissionResource::Connection => self.connections.clone(),
            AdmissionResource::Handshake => self.handshakes.clone(),
            AdmissionResource::PartialFrame => self.partial_frames.clone(),
            AdmissionResource::Inflight => self.inflight.clone(),
            AdmissionResource::Queued => self.queued.clone(),
            AdmissionResource::Processor => self.processors.clone(),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum ScopeKey {
    Ip(AdmissionResource, IpAddr),
    Tenant(AdmissionResource, IpAddr, u64),
    Session(AdmissionResource, IpAddr, Option<u64>, u64),
    EmbeddedCaller(AdmissionResource, EmbeddedCaller),
    EmbeddedSession(AdmissionResource, EmbeddedCaller, u64),
}

/// RAII ownership of global and scoped admission capacity.
pub struct AdmissionPermit {
    _permit: ResourcePermit,
    // Field order returns capacity before publishing Released.
    release: AdmissionRelease,
}

struct AdmissionRelease {
    observer: Option<tokio::sync::mpsc::Sender<AdmissionEvent>>,
    resource: AdmissionResource,
    bytes: usize,
}

/// Capacity retained from the moment a peer announces a frame until that frame enters execution.
#[derive(Debug)]
pub(crate) struct PartialFramePermit {
    permit: AdmissionPermit,
}

impl PartialFramePermit {
    pub(crate) fn new(permit: AdmissionPermit) -> Self {
        Self { permit }
    }

    pub(crate) fn try_rebind(
        mut self,
        handle: &AdmissionScopeHandle,
        resource: AdmissionResource,
        class: AdmissionClass,
    ) -> Result<AdmissionPermit, (Box<Self>, AdmissionRejection)> {
        if class == AdmissionClass::Control {
            self.permit._permit.promote_to_control();
        }
        let budget = handle.budget(resource);
        if self.permit._permit.try_rebind(&budget).is_err() {
            let error = AdmissionRejection {
                resource,
                policy: policy_for(resource),
            };
            if let Some(observer) = &handle.observer {
                let _ = observer.try_send(AdmissionEvent {
                    resource,
                    outcome: AdmissionEventOutcome::Rejected,
                    bytes: self.permit.release.bytes,
                });
            }
            return Err((Box::new(self), error));
        }
        if let Some(observer) = &self.permit.release.observer {
            let _ = observer.try_send(AdmissionEvent {
                resource: self.permit.release.resource,
                outcome: AdmissionEventOutcome::Released,
                bytes: self.permit.release.bytes,
            });
            let _ = observer.try_send(AdmissionEvent {
                resource,
                outcome: AdmissionEventOutcome::Acquired,
                bytes: self.permit.release.bytes,
            });
        }
        self.permit.release.resource = resource;
        Ok(self.permit)
    }
}

#[derive(Debug)]
struct PreparedScopeBudgets {
    connection: ResourceBudget,
    handshake: ResourceBudget,
    partial_frame: ResourceBudget,
    inflight: ResourceBudget,
    queued: ResourceBudget,
    processor: ResourceBudget,
}

/// Prepared per-session admission capability used without registry lookup on the request path.
#[derive(Clone, Debug)]
pub(crate) struct AdmissionScopeHandle {
    budgets: Arc<PreparedScopeBudgets>,
    observer: Option<tokio::sync::mpsc::Sender<AdmissionEvent>>,
    session_id: Option<u64>,
}

impl AdmissionScopeHandle {
    pub(crate) const fn session_id(&self) -> Option<u64> {
        self.session_id
    }

    fn budget(&self, resource: AdmissionResource) -> ResourceBudget {
        match resource {
            AdmissionResource::Connection => self.budgets.connection.clone(),
            AdmissionResource::Handshake => self.budgets.handshake.clone(),
            AdmissionResource::PartialFrame => self.budgets.partial_frame.clone(),
            AdmissionResource::Inflight => self.budgets.inflight.clone(),
            AdmissionResource::Queued => self.budgets.queued.clone(),
            AdmissionResource::Processor => self.budgets.processor.clone(),
        }
    }

    pub(crate) fn try_acquire(
        &self,
        resource: AdmissionResource,
        bytes: usize,
        class: AdmissionClass,
    ) -> Result<AdmissionPermit, AdmissionRejection> {
        let budget = self.budget(resource);
        let budget_class = match class {
            AdmissionClass::Data => BudgetClass::Data,
            AdmissionClass::Control => BudgetClass::Control,
        };
        let permit = budget.try_acquire(bytes, budget_class).map_err(|_| {
            self.observe(resource, AdmissionEventOutcome::Rejected, bytes);
            AdmissionRejection {
                resource,
                policy: policy_for(resource),
            }
        })?;
        self.observe(resource, AdmissionEventOutcome::Acquired, bytes);
        Ok(AdmissionPermit {
            _permit: permit,
            release: AdmissionRelease {
                observer: self.observer.clone(),
                resource,
                bytes,
            },
        })
    }

    pub(crate) fn rebind_permit(
        &self,
        resource: AdmissionResource,
        mut permit: ResourcePermit,
    ) -> Result<AdmissionPermit, AdmissionRejection> {
        let budget = self.budget(resource);
        permit.try_rebind(&budget).map_err(|_| {
            self.observe(resource, AdmissionEventOutcome::Rejected, permit.bytes());
            AdmissionRejection {
                resource,
                policy: policy_for(resource),
            }
        })?;
        let bytes = permit.bytes();
        self.observe(resource, AdmissionEventOutcome::Acquired, bytes);
        Ok(AdmissionPermit {
            _permit: permit,
            release: AdmissionRelease {
                observer: self.observer.clone(),
                resource,
                bytes,
            },
        })
    }

    fn observe(&self, resource: AdmissionResource, outcome: AdmissionEventOutcome, bytes: usize) {
        if let Some(observer) = &self.observer {
            let _ = observer.try_send(AdmissionEvent {
                resource,
                outcome,
                bytes,
            });
        }
    }
}

impl fmt::Debug for AdmissionPermit {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("AdmissionPermit")
            .field("resource", &self.release.resource)
            .field("bytes", &self.release.bytes)
            .finish_non_exhaustive()
    }
}

impl Drop for AdmissionRelease {
    fn drop(&mut self) {
        if let Some(observer) = &self.observer {
            let _ = observer.try_send(AdmissionEvent {
                resource: self.resource,
                outcome: AdmissionEventOutcome::Released,
                bytes: self.bytes,
            });
        }
    }
}

/// Owner of global/per-IP/per-tenant/per-session transport budgets and metrics.
pub struct AdmissionController {
    limits: AdmissionLimits,
    global: GlobalBudgets,
    process_budget: ResourceBudget,
    deferred_admission: Mutex<Option<DeferredAdmission>>,
    scoped: Mutex<HashMap<ScopeKey, ResourceBudget>>,
    observer: Option<tokio::sync::mpsc::Sender<AdmissionEvent>>,
}

impl AdmissionController {
    /// Creates a transport admission controller.
    ///
    /// # Panics
    ///
    /// Panics when a limit is zero or a control reserve exceeds its owning
    /// resource. Production configuration paths should prefer [`Self::try_new`].
    pub fn new(limits: AdmissionLimits) -> Self {
        Self::try_new(limits).unwrap_or_else(|error| panic!("invalid transport admission limits: {error}"))
    }

    pub fn try_new(limits: AdmissionLimits) -> Result<Self, TransportContractViolation> {
        let process_budget = standalone_process_budget(limits)?;
        Self::build(limits, &process_budget, None)
    }

    /// Creates a controller whose transport budgets derive from the injected
    /// process resource root.
    ///
    /// # Errors
    ///
    /// Returns an error when any transport or derived child limit is invalid.
    pub fn try_new_with_budget(
        limits: AdmissionLimits,
        process_budget: &ResourceBudget,
    ) -> Result<Self, TransportContractViolation> {
        Self::build(limits, process_budget, None)
    }

    /// Creates an observed transport admission controller.
    ///
    /// # Panics
    ///
    /// Panics under the same invalid-limit conditions as [`Self::new`].
    pub fn with_observer(limits: AdmissionLimits, observer: tokio::sync::mpsc::Sender<AdmissionEvent>) -> Self {
        Self::try_with_observer(limits, observer)
            .unwrap_or_else(|error| panic!("invalid transport admission limits: {error}"))
    }

    pub fn try_with_observer(
        limits: AdmissionLimits,
        observer: tokio::sync::mpsc::Sender<AdmissionEvent>,
    ) -> Result<Self, TransportContractViolation> {
        let process_budget = standalone_process_budget(limits)?;
        Self::build(limits, &process_budget, Some(observer))
    }

    /// Creates an observed controller under an injected process resource root.
    ///
    /// # Errors
    ///
    /// Returns an error when any transport or derived child limit is invalid.
    pub fn try_with_budget_and_observer(
        limits: AdmissionLimits,
        process_budget: &ResourceBudget,
        observer: tokio::sync::mpsc::Sender<AdmissionEvent>,
    ) -> Result<Self, TransportContractViolation> {
        Self::build(limits, process_budget, Some(observer))
    }

    fn build(
        limits: AdmissionLimits,
        process_budget: &ResourceBudget,
        observer: Option<tokio::sync::mpsc::Sender<AdmissionEvent>>,
    ) -> Result<Self, TransportContractViolation> {
        validate_scope_limits(limits)?;
        Ok(Self {
            limits,
            global: GlobalBudgets::new(limits, process_budget)?,
            process_budget: process_budget.clone(),
            deferred_admission: Mutex::new(None),
            scoped: Mutex::new(HashMap::new()),
            observer,
        })
    }

    pub(crate) fn configure_deferred_admission(
        &self,
        limits: DeferredWaitLimits,
    ) -> Result<DeferredAdmission, TransportContractViolation> {
        let mut configured = self
            .deferred_admission
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if let Some(admission) = configured.as_ref() {
            return if admission.limits() == limits {
                Ok(admission.clone())
            } else {
                Err(TransportContractViolation::DeferredAdmissionConflict)
            };
        }
        let admission = DeferredAdmission::try_new(&self.process_budget, limits)?;
        *configured = Some(admission.clone());
        Ok(admission)
    }

    pub(crate) fn deferred_admission(&self) -> Option<DeferredAdmission> {
        self.deferred_admission
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .clone()
    }

    pub fn try_acquire(
        &self,
        resource: AdmissionResource,
        scope: AdmissionScope,
        bytes: usize,
        class: AdmissionClass,
    ) -> AdmissionOutcome<AdmissionPermit> {
        let policy = policy_for(resource);
        let budget = match self.scoped_budget(resource, scope) {
            Ok(budget) => budget,
            Err(rejection) => return AdmissionOutcome::Rejected(rejection),
        };
        let class = match class {
            AdmissionClass::Data => BudgetClass::Data,
            AdmissionClass::Control => BudgetClass::Control,
        };
        let permit = match budget.try_acquire(bytes, class) {
            Ok(permit) => permit,
            Err(_) => {
                if let Some(observer) = &self.observer {
                    let _ = observer.try_send(AdmissionEvent {
                        resource,
                        outcome: AdmissionEventOutcome::Rejected,
                        bytes,
                    });
                }
                return AdmissionOutcome::Rejected(AdmissionRejection { resource, policy });
            }
        };
        if let Some(observer) = &self.observer {
            let _ = observer.try_send(AdmissionEvent {
                resource,
                outcome: AdmissionEventOutcome::Acquired,
                bytes,
            });
        }
        AdmissionOutcome::Acquired(AdmissionPermit {
            _permit: permit,
            release: AdmissionRelease {
                observer: self.observer.clone(),
                resource,
                bytes,
            },
        })
    }

    pub(crate) fn prepare_scope(&self, scope: AdmissionScope) -> Result<AdmissionScopeHandle, AdmissionRejection> {
        let session_id = scope.session;
        Ok(AdmissionScopeHandle {
            budgets: Arc::new(PreparedScopeBudgets {
                connection: self.scoped_budget(AdmissionResource::Connection, scope)?,
                handshake: self.scoped_budget(AdmissionResource::Handshake, scope)?,
                partial_frame: self.scoped_budget(AdmissionResource::PartialFrame, scope)?,
                inflight: self.scoped_budget(AdmissionResource::Inflight, scope)?,
                queued: self.scoped_budget(AdmissionResource::Queued, scope)?,
                processor: self.scoped_budget(AdmissionResource::Processor, scope)?,
            }),
            observer: self.observer.clone(),
            session_id,
        })
    }

    pub(crate) fn prepare_embedded_scope(
        &self,
        caller: EmbeddedCaller,
        session_id: u64,
    ) -> Result<AdmissionScopeHandle, AdmissionRejection> {
        Ok(AdmissionScopeHandle {
            budgets: Arc::new(PreparedScopeBudgets {
                connection: self.embedded_scoped_budget(AdmissionResource::Connection, caller, session_id)?,
                handshake: self.embedded_scoped_budget(AdmissionResource::Handshake, caller, session_id)?,
                partial_frame: self.embedded_scoped_budget(AdmissionResource::PartialFrame, caller, session_id)?,
                inflight: self.embedded_scoped_budget(AdmissionResource::Inflight, caller, session_id)?,
                queued: self.embedded_scoped_budget(AdmissionResource::Queued, caller, session_id)?,
                processor: self.embedded_scoped_budget(AdmissionResource::Processor, caller, session_id)?,
            }),
            observer: self.observer.clone(),
            session_id: Some(session_id),
        })
    }

    fn embedded_scoped_budget(
        &self,
        resource: AdmissionResource,
        caller: EmbeddedCaller,
        session_id: u64,
    ) -> Result<ResourceBudget, AdmissionRejection> {
        let mut scoped = self.scoped.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
        let global = self.global.get(resource);
        let scoped_reserve = reserve_for(resource, self.limits.control_reserve);
        let caller_key = ScopeKey::EmbeddedCaller(resource, caller);
        let session_key = ScopeKey::EmbeddedSession(resource, caller, session_id);
        let scope_keys = [caller_key, session_key];
        let mut missing_scope_count = scope_keys.into_iter().filter(|key| !scoped.contains_key(key)).count();
        if scoped.len().saturating_add(missing_scope_count) > self.limits.max_scope_keys {
            scoped.retain(|_, budget| budget.snapshot().current_count > 0);
            missing_scope_count = scope_keys.into_iter().filter(|key| !scoped.contains_key(key)).count();
        }
        if scoped.len().saturating_add(missing_scope_count) > self.limits.max_scope_keys {
            return Err(AdmissionRejection {
                resource,
                policy: policy_for(resource),
            });
        }
        let caller_budget = scoped_child(
            &mut scoped,
            self.limits.max_scope_keys,
            caller_key,
            &global,
            format!("embedded-caller-{caller:?}"),
            self.limits.per_ip,
            scoped_reserve,
            resource,
        )?;
        scoped_child(
            &mut scoped,
            self.limits.max_scope_keys,
            session_key,
            &caller_budget,
            format!("embedded-session-{session_id}"),
            self.limits.per_session,
            scoped_reserve,
            resource,
        )
    }

    fn scoped_budget(
        &self,
        resource: AdmissionResource,
        scope: AdmissionScope,
    ) -> Result<ResourceBudget, AdmissionRejection> {
        let mut scoped = self.scoped.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
        let global = self.global.get(resource);
        let scoped_reserve = reserve_for(resource, self.limits.control_reserve);
        let ip_key = ScopeKey::Ip(resource, scope.ip);
        let tenant_key = scope.tenant.map(|tenant| ScopeKey::Tenant(resource, scope.ip, tenant));
        let session_key = scope
            .session
            .map(|session| ScopeKey::Session(resource, scope.ip, scope.tenant, session));
        let scope_keys = [Some(ip_key), tenant_key, session_key];
        let mut missing_scope_count = scope_keys
            .into_iter()
            .flatten()
            .filter(|key| !scoped.contains_key(key))
            .count();
        if scoped.len().saturating_add(missing_scope_count) > self.limits.max_scope_keys {
            scoped.retain(|_, budget| budget.snapshot().current_count > 0);
            missing_scope_count = scope_keys
                .into_iter()
                .flatten()
                .filter(|key| !scoped.contains_key(key))
                .count();
        }
        if scoped.len().saturating_add(missing_scope_count) > self.limits.max_scope_keys {
            return Err(AdmissionRejection {
                resource,
                policy: policy_for(resource),
            });
        }
        let ip_budget = scoped_child(
            &mut scoped,
            self.limits.max_scope_keys,
            ip_key,
            &global,
            format!("ip-{}", scope.ip),
            self.limits.per_ip,
            scoped_reserve,
            resource,
        )?;
        let tenant_budget = match (scope.tenant, tenant_key) {
            (Some(tenant), Some(tenant_key)) => scoped_child(
                &mut scoped,
                self.limits.max_scope_keys,
                tenant_key,
                &ip_budget,
                format!("tenant-{tenant}"),
                self.limits.per_tenant,
                scoped_reserve,
                resource,
            )?,
            (None, None) => ip_budget,
            _ => unreachable!("tenant scope key is derived from the tenant identifier"),
        };
        match (scope.session, session_key) {
            (Some(session), Some(session_key)) => scoped_child(
                &mut scoped,
                self.limits.max_scope_keys,
                session_key,
                &tenant_budget,
                format!("session-{session}"),
                self.limits.per_session,
                scoped_reserve,
                resource,
            ),
            (None, None) => Ok(tenant_budget),
            _ => unreachable!("session scope key is derived from the session identifier"),
        }
    }

    pub fn snapshot(&self) -> AdmissionSnapshot {
        AdmissionSnapshot {
            connections: resource_snapshot(&self.global.connections),
            handshakes: resource_snapshot(&self.global.handshakes),
            partial_frames: resource_snapshot(&self.global.partial_frames),
            inflight: resource_snapshot(&self.global.inflight),
            queued: resource_snapshot(&self.global.queued),
            processors: resource_snapshot(&self.global.processors),
        }
    }
}

fn standalone_process_budget(limits: AdmissionLimits) -> Result<ResourceBudget, TransportContractViolation> {
    let resources = [
        limits.connections,
        limits.handshakes,
        limits.partial_frames,
        limits.inflight,
        limits.queued,
        limits.processors,
    ];
    let count = resources
        .iter()
        .fold(0usize, |total, limit| total.saturating_add(limit.count))
        .max(1);
    let bytes = resources
        .iter()
        .fold(0usize, |total, limit| total.saturating_add(limit.bytes))
        .max(1);
    Ok(ResourceBudgetTree::new(
        "standalone-transport-process",
        BudgetLimit::new(count, bytes, FullPolicy::Reject),
    )?
    .root())
}

fn global_budget(
    root: &ResourceBudget,
    name: &str,
    limit: ResourceLimit,
    reserve: ResourceLimit,
    policy: FullPolicy,
) -> Result<ResourceBudget, RuntimeContractViolation> {
    let root_capacity = root.limit().capacity;
    let limit = ResourceLimit {
        count: limit.count.min(root_capacity.count),
        bytes: limit.bytes.min(root_capacity.bytes),
    };
    let reserve = ResourceLimit {
        count: effective_reserve(reserve.count, limit.count),
        bytes: effective_reserve(reserve.bytes, limit.bytes),
    };
    root.child(
        name,
        BudgetLimit::new(limit.count, limit.bytes, policy)
            .with_control_reserve(BudgetCapacity::new(reserve.count, reserve.bytes)),
    )
}

fn effective_reserve(requested: usize, capacity: usize) -> usize {
    if requested < capacity {
        requested
    } else {
        0
    }
}

fn reserve_for(resource: AdmissionResource, configured: ResourceLimit) -> ResourceLimit {
    if resource == AdmissionResource::PartialFrame {
        ResourceLimit { count: 0, bytes: 0 }
    } else {
        configured
    }
}

fn scoped_child(
    scoped: &mut HashMap<ScopeKey, ResourceBudget>,
    max_scope_keys: usize,
    key: ScopeKey,
    parent: &ResourceBudget,
    name: String,
    requested: ResourceLimit,
    requested_reserve: ResourceLimit,
    resource: AdmissionResource,
) -> Result<ResourceBudget, AdmissionRejection> {
    if let Some(budget) = scoped.get(&key) {
        return Ok(budget.clone());
    }
    if scoped.len() >= max_scope_keys {
        return Err(AdmissionRejection {
            resource,
            policy: policy_for(resource),
        });
    }
    let parent_limit = parent.limit().capacity;
    let limit = ResourceLimit {
        count: requested.count.min(parent_limit.count),
        bytes: requested.bytes.min(parent_limit.bytes),
    };
    let reserve = ResourceLimit {
        count: effective_reserve(requested_reserve.count, limit.count),
        bytes: effective_reserve(requested_reserve.bytes, limit.bytes),
    };
    let budget = parent
        .child(
            name,
            BudgetLimit::new(limit.count, limit.bytes, parent.limit().full_policy)
                .with_control_reserve(BudgetCapacity::new(reserve.count, reserve.bytes)),
        )
        .map_err(|_| AdmissionRejection {
            resource,
            policy: policy_for(resource),
        })?;
    scoped.insert(key, budget.clone());
    Ok(budget)
}

fn validate_scope_limits(limits: AdmissionLimits) -> Result<(), TransportContractViolation> {
    if limits.max_scope_keys == 0 {
        return Err(TransportContractViolation::AdmissionZeroMaxScopeKeys);
    }
    for (scope, limit) in [
        ("per_ip", limits.per_ip),
        ("per_tenant", limits.per_tenant),
        ("per_session", limits.per_session),
    ] {
        if limit.count == 0 {
            return Err(TransportContractViolation::AdmissionZeroScopeCapacity {
                scope,
                dimension: BudgetDimension::Count,
            });
        }
        if limit.bytes == 0 {
            return Err(TransportContractViolation::AdmissionZeroScopeCapacity {
                scope,
                dimension: BudgetDimension::Bytes,
            });
        }
    }
    Ok(())
}

fn policy_for(resource: AdmissionResource) -> FullPolicy {
    match resource {
        AdmissionResource::Connection | AdmissionResource::Handshake | AdmissionResource::PartialFrame => {
            FullPolicy::CloseSlowConsumer
        }
        AdmissionResource::Inflight | AdmissionResource::Queued | AdmissionResource::Processor => FullPolicy::Reject,
    }
}

fn resource_snapshot(budget: &ResourceBudget) -> ResourceSnapshot {
    let snapshot = budget.snapshot();
    ResourceSnapshot {
        current_count: snapshot.current_count,
        current_bytes: snapshot.current_bytes,
        rejected_count: usize::try_from(snapshot.rejected_count).unwrap_or(usize::MAX),
    }
}

#[cfg(test)]
mod tests {
    use std::net::IpAddr;
    use std::net::Ipv4Addr;

    use super::AdmissionClass;
    use super::AdmissionController;
    use super::AdmissionLimits;
    use super::AdmissionResource;
    use super::AdmissionScope;

    #[test]
    fn prepared_scope_reuses_cached_budgets_on_request_path() {
        let controller = AdmissionController::new(AdmissionLimits::default());
        let scope = AdmissionScope::new(IpAddr::V4(Ipv4Addr::LOCALHOST))
            .with_tenant(7)
            .with_session(11);
        let handle = controller
            .prepare_scope(scope)
            .expect("prepare session admission scope");
        let prepared_keys = controller
            .scoped
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .len();

        for _ in 0..1_000 {
            let permit = handle
                .try_acquire(AdmissionResource::Inflight, 128, AdmissionClass::Data)
                .expect("cached session budget remains available");
            drop(permit);
        }

        let final_keys = controller
            .scoped
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .len();
        assert_eq!(final_keys, prepared_keys);
    }
}