decdn_node/load_shed/
state.rs

1//! Live in-flight serve counters: one node-wide count and one per client,
2//! generalizing the per-lane atomic-count admission technique to a node-wide
3//! resource gate. A [`ShedSlot`] holds one admitted stream's place and releases
4//! it on every exit — success, error, `?`, disconnect, panic — via `Drop`.
5
6use std::sync::Arc;
7use std::sync::atomic::{AtomicU32, Ordering};
8
9use alloy::primitives::B256;
10use dashmap::DashMap;
11
12/// Node-wide and per-client in-flight serve counts. Held in an `Arc` so a
13/// [`ShedSlot`] can outlive the borrow that created it.
14#[derive(Debug, Default)]
15pub struct ShedState {
16    node_active: AtomicU32,
17    per_client: DashMap<B256, Arc<AtomicU32>>,
18}
19
20impl ShedState {
21    /// A counter with nothing in flight, already wrapped so [`ShedSlot`]s can
22    /// hold it.
23    #[must_use]
24    pub fn new() -> Arc<Self> {
25        Arc::new(Self::default())
26    }
27
28    /// `(node_in_flight, client_in_flight)`. A client with no entry counts 0.
29    #[must_use]
30    pub fn counts(&self, client: B256) -> (u32, u32) {
31        let node = self.node_active.load(Ordering::Relaxed);
32        let per = self
33            .per_client
34            .get(&client)
35            .map_or(0, |c| c.load(Ordering::Relaxed));
36        (node, per)
37    }
38
39    /// Node-wide serves in flight across all clients — the same count
40    /// [`Self::counts`] returns as its first element, without needing a client
41    /// key. Sampled onto the `decdn_load_shed_streams_in_flight` gauge so an
42    /// operator can read live concurrency against the shed high-water mark.
43    #[must_use]
44    pub fn node_in_flight(&self) -> u32 {
45        self.node_active.load(Ordering::Relaxed)
46    }
47
48    /// Increment both counters and hand back the RAII slot.
49    #[must_use]
50    pub fn acquire(self: &Arc<Self>, client: B256) -> ShedSlot {
51        self.node_active.fetch_add(1, Ordering::Relaxed);
52        let counter = self
53            .per_client
54            .entry(client)
55            .or_insert_with(|| Arc::new(AtomicU32::new(0)))
56            .clone();
57        counter.fetch_add(1, Ordering::Relaxed);
58        ShedSlot {
59            state: Arc::clone(self),
60            client,
61            counter,
62        }
63    }
64}
65
66/// RAII place for one admitted serve. `Drop` releases the node-wide and
67/// per-client counts on every exit path, so a finished stream always frees its
68/// slot; the count is owned here, never decremented by hand.
69#[derive(Debug)]
70pub struct ShedSlot {
71    state: Arc<ShedState>,
72    client: B256,
73    counter: Arc<AtomicU32>,
74}
75
76impl Drop for ShedSlot {
77    fn drop(&mut self) {
78        self.state
79            .node_active
80            .update(Ordering::Relaxed, Ordering::Relaxed, |n| {
81                n.saturating_sub(1)
82            });
83        self.counter
84            .update(Ordering::Relaxed, Ordering::Relaxed, |n| {
85                n.saturating_sub(1)
86            });
87        // Best-effort prune: only removes if whatever is currently mapped for
88        // this key is zero; a concurrently-recreated live entry is left intact
89        // (remove_if re-checks under the shard lock).
90        self.state
91            .per_client
92            .remove_if(&self.client, |_, c| c.load(Ordering::Relaxed) == 0);
93    }
94}
95
96#[cfg(test)]
97mod tests;
decdn_node/load_shed/state.rs

decdn_node/load_shed/
state.rs
