ursula_runtime/
admission.rs

1//! Orthogonal write-path admission controls.
2//!
3//! Each admission gates accept on one independent "slow downstream":
4//! - cold path full (hot_bytes_per_group, existing [`crate::request::ColdWriteAdmission`])
5//! - raft replication lagging ([`RaftUncommittedAdmission`], uncommitted_bytes_per_group)
6//! - forward queue piling on a remote peer (inflight_forward_bytes_per_peer, lives in `ursula::lib`)
7//! - process memory near OOM (rss vs soft_cap, lives in `ursula::lib`)
8//!
9//! Each admission can be independently configured (`None` = disabled).
10//! Call sites consult the relevant subset; errors surface as HTTP 503.
11
12use std::sync::Arc;
13use std::sync::atomic::AtomicU64;
14use std::sync::atomic::Ordering;
15
16use ursula_shard::RaftGroupId;
17
18/// Per-group admission that rejects new writes when the raft layer has not yet
19/// committed enough previously-submitted bytes. Independently configurable from
20/// the cold-side admission; intended to catch "replication lag" scenarios where
21/// hot bytes have not yet grown because nothing is committing.
22#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
23pub struct RaftUncommittedAdmission {
24    pub max_uncommitted_bytes_per_group: Option<u64>,
25}
26
27impl RaftUncommittedAdmission {
28    pub fn is_enabled(self) -> bool {
29        self.max_uncommitted_bytes_per_group.is_some()
30    }
31}
32
33/// Lock-free per-group counters for in-flight (submitted but not-yet-applied)
34/// raft payload bytes. Shared between the core worker (increment on submit) and
35/// the group actor (decrement on apply completion).
36#[derive(Debug)]
37pub(crate) struct RaftUncommittedBytesTracker {
38    per_group: Vec<AtomicU64>,
39}
40
41impl RaftUncommittedBytesTracker {
42    pub(crate) fn new(group_count: usize) -> Self {
43        Self {
44            per_group: (0..group_count).map(|_| AtomicU64::new(0)).collect(),
45        }
46    }
47
48    pub(crate) fn load(&self, group_id: RaftGroupId) -> u64 {
49        self.slot(group_id).load(Ordering::Relaxed)
50    }
51
52    pub(crate) fn add(&self, group_id: RaftGroupId, bytes: u64) {
53        self.slot(group_id).fetch_add(bytes, Ordering::Relaxed);
54    }
55
56    pub(crate) fn sub(&self, group_id: RaftGroupId, bytes: u64) {
57        // Use saturating subtract to avoid wrap if a request is double-credited
58        // (defense-in-depth; the call sites pair add/sub one-for-one).
59        let slot = self.slot(group_id);
60        let mut current = slot.load(Ordering::Relaxed);
61        loop {
62            let next = current.saturating_sub(bytes);
63            match slot.compare_exchange_weak(current, next, Ordering::Relaxed, Ordering::Relaxed) {
64                Ok(_) => return,
65                Err(observed) => current = observed,
66            }
67        }
68    }
69
70    fn slot(&self, group_id: RaftGroupId) -> &AtomicU64 {
71        let index = usize::try_from(group_id.0).expect("u32 fits usize");
72        &self.per_group[index]
73    }
74}
75
76pub(crate) type SharedRaftUncommittedBytes = Arc<RaftUncommittedBytesTracker>;
77
78/// Guard that decrements the uncommitted bytes counter on drop. Pair an `add`
79/// at submit-time with a guard that lives until the apply (or apply-failure)
80/// completes so we never leak credit on early errors.
81pub(crate) struct UncommittedBytesGuard {
82    tracker: SharedRaftUncommittedBytes,
83    group_id: RaftGroupId,
84    bytes: u64,
85    armed: bool,
86}
87
88impl UncommittedBytesGuard {
89    pub(crate) fn new(
90        tracker: SharedRaftUncommittedBytes,
91        group_id: RaftGroupId,
92        bytes: u64,
93    ) -> Self {
94        tracker.add(group_id, bytes);
95        Self {
96            tracker,
97            group_id,
98            bytes,
99            armed: true,
100        }
101    }
102
103    /// Disarm without releasing; useful when something else has taken over
104    /// the credit (we currently always release on drop, so this is reserved
105    /// for future use).
106    #[allow(dead_code)]
107    pub(crate) fn disarm(mut self) {
108        self.armed = false;
109    }
110}
111
112impl Drop for UncommittedBytesGuard {
113    fn drop(&mut self) {
114        if self.armed {
115            self.tracker.sub(self.group_id, self.bytes);
116        }
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[test]
125    fn disabled_admission_reports_disabled() {
126        let admission = RaftUncommittedAdmission::default();
127        assert!(!admission.is_enabled());
128    }
129
130    #[test]
131    fn enabled_admission_reports_enabled() {
132        let admission = RaftUncommittedAdmission {
133            max_uncommitted_bytes_per_group: Some(1024),
134        };
135        assert!(admission.is_enabled());
136    }
137
138    #[test]
139    fn tracker_add_load_sub_round_trips() {
140        let tracker = RaftUncommittedBytesTracker::new(2);
141        tracker.add(RaftGroupId(0), 32);
142        tracker.add(RaftGroupId(0), 8);
143        tracker.add(RaftGroupId(1), 4);
144        assert_eq!(tracker.load(RaftGroupId(0)), 40);
145        assert_eq!(tracker.load(RaftGroupId(1)), 4);
146        tracker.sub(RaftGroupId(0), 16);
147        assert_eq!(tracker.load(RaftGroupId(0)), 24);
148    }
149
150    #[test]
151    fn tracker_sub_saturates_at_zero() {
152        let tracker = RaftUncommittedBytesTracker::new(1);
153        tracker.add(RaftGroupId(0), 4);
154        tracker.sub(RaftGroupId(0), 10);
155        assert_eq!(tracker.load(RaftGroupId(0)), 0);
156    }
157
158    #[test]
159    fn guard_releases_on_drop() {
160        let tracker = Arc::new(RaftUncommittedBytesTracker::new(1));
161        {
162            let _guard = UncommittedBytesGuard::new(tracker.clone(), RaftGroupId(0), 32);
163            assert_eq!(tracker.load(RaftGroupId(0)), 32);
164        }
165        assert_eq!(tracker.load(RaftGroupId(0)), 0);
166    }
167}
ursula_runtime/admission.rs

ursula_runtime/
admission.rs
