decdn_node/load_shed/
controller.rs

1//! `LoadShedController` ties the swappable policy, the live counters, and the
2//! egress meter into one handle the serve path holds. It mirrors
3//! `ConnectionLimiter`: the handler owns an `Arc<LoadShedController>`, and a
4//! hot-reload swaps the policy in place under an `ArcSwap`.
5
6use std::sync::Arc;
7
8use alloy::primitives::B256;
9use arc_swap::ArcSwap;
10use decdn_common::config::{LoadShedPolicyKind, ResolvedLoadShed};
11
12use super::{
13    AlwaysAdmit, EgressMeter, LoadShedPolicy, PressureSnapshot, RequestClass, ResourcePressure,
14    ResourcePressureParams, ShedDecision, ShedReason, ShedSlot, ShedState,
15};
16
17/// Bytes/sec for a megabit/sec budget: `Mbps * 1_000_000 / 8`.
18#[must_use]
19const fn egress_budget_bps(mbps: u64) -> u64 {
20    mbps.saturating_mul(125_000)
21}
22
23/// Build the policy object a resolved config selects.
24fn build_policy(cfg: &ResolvedLoadShed) -> Arc<dyn LoadShedPolicy> {
25    match cfg.policy {
26        LoadShedPolicyKind::AlwaysAdmit => Arc::new(AlwaysAdmit),
27        LoadShedPolicyKind::ResourcePressure => {
28            Arc::new(ResourcePressure::new(ResourcePressureParams {
29                egress_budget_bps: egress_budget_bps(cfg.egress_budget_mbps),
30                max_serves_high: cfg.max_concurrent_serves_high,
31                max_serves_low: cfg.max_concurrent_serves_low,
32                per_client_cap: cfg.per_client_serve_cap,
33            }))
34        }
35    }
36}
37
38/// The serve path's load-shed handle: a swappable policy over live meters.
39#[allow(missing_debug_implementations)]
40pub struct LoadShedController {
41    // `ArcSwap<Arc<dyn LoadShedPolicy>>` (not `ArcSwap<dyn LoadShedPolicy>`): arc-swap's
42    // `RefCnt` blanket impl `impl<T> RefCnt for Arc<T>` bounds `T: Sized`, so
43    // `Arc<dyn LoadShedPolicy>` does not implement `RefCnt`. Storing the trait object
44    // as the sized `Arc<dyn ...>` pointee is the supported form; `.load()` derefs
45    // through both `Arc`s to the trait object.
46    policy: ArcSwap<Arc<dyn LoadShedPolicy>>,
47    state: Arc<ShedState>,
48    egress: Arc<EgressMeter>,
49}
50
51impl LoadShedController {
52    /// Build a controller with the policy `cfg` selects, a fresh in-flight
53    /// counter, and an unseeded egress meter.
54    #[must_use]
55    pub fn from_config(cfg: &ResolvedLoadShed) -> Arc<Self> {
56        Arc::new(Self {
57            policy: ArcSwap::from_pointee(build_policy(cfg)),
58            state: ShedState::new(),
59            egress: Arc::new(EgressMeter::new()),
60        })
61    }
62
63    /// Decide and, on admit, take a slot. On shed returns the reason; the caller
64    /// maps it to the `NotFound`-shaped refusal.
65    pub fn try_admit(&self, class: RequestClass, client: B256) -> Result<ShedSlot, ShedReason> {
66        let (node, per) = self.state.counts(client);
67        let snap = PressureSnapshot {
68            streams_in_flight: node,
69            client_streams_in_flight: per,
70            egress_bps: self.egress.current_bps(),
71        };
72        let policy = self.policy.load();
73        match policy.decide(class, &snap) {
74            ShedDecision::Admit => Ok(self.state.acquire(client)),
75            ShedDecision::Shed(reason) => Err(reason),
76        }
77    }
78
79    /// Feed delivered bytes into the egress meter (delivery hot path). `bytes`
80    /// is the per-chunk length also added to the billed `delivered` counter —
81    /// the bao-stream content plus interleaved proof, not the full QUIC frame.
82    pub fn record_egress(&self, bytes: u64) {
83        self.egress.record(bytes);
84    }
85
86    /// Fold the interval's bytes into the egress EWMA; returns the new bytes/sec
87    /// (for the gauge).
88    pub fn sample_egress(&self, interval_secs: u64) -> u64 {
89        self.egress.sample(interval_secs)
90    }
91
92    /// Whether the active policy currently considers the node pressured (gauge).
93    #[must_use]
94    pub fn pressure_active(&self) -> bool {
95        self.policy.load().pressure_active()
96    }
97
98    /// Node-wide serves in flight, for the `decdn_load_shed_streams_in_flight`
99    /// gauge. Reads the same live counter `try_admit` feeds into the pressure
100    /// snapshot, so the gauge tracks exactly what the policy decides against.
101    #[must_use]
102    pub fn node_in_flight(&self) -> u32 {
103        self.state.node_in_flight()
104    }
105
106    /// Swap in a policy rebuilt from a new resolved config. Live counters and
107    /// the egress meter are preserved; only the decision policy changes.
108    pub fn reload(&self, cfg: &ResolvedLoadShed) {
109        self.policy.store(Arc::new(build_policy(cfg)));
110    }
111}
112
113#[cfg(test)]
114#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
115mod tests;
decdn_node/load_shed/controller.rs

decdn_node/load_shed/
controller.rs
