1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
//! In-flight run registry. Tracks per-run cancellation tokens and the queue /
//! in-flight counters that drive backpressure (429), the `faucet_serve_runs_*`
//! gauges, `/readyz`, and the shutdown drain. A "queued" run is one that has been
//! spawned but has not yet acquired an execution permit.
use dashmap::DashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::Notify;
use tokio_util::sync::CancellationToken;
pub struct Registry {
tokens: DashMap<String, CancellationToken>,
queued: AtomicUsize,
in_flight: AtomicUsize,
max_queued: usize,
drained: Notify,
}
impl Registry {
pub fn new(max_queued: usize) -> Self {
Self {
tokens: DashMap::new(),
queued: AtomicUsize::new(0),
in_flight: AtomicUsize::new(0),
max_queued: max_queued.max(1),
drained: Notify::new(),
}
}
/// Reserve a queue slot. Returns `true` if a slot was reserved, or `false`
/// if the queue is full. Atomic against concurrent submits via a CAS loop
/// (`AtomicUsize::try_update`, stabilized in Rust 1.95).
pub fn try_reserve(&self) -> bool {
self.queued
.try_update(Ordering::AcqRel, Ordering::Acquire, |cur| {
(cur < self.max_queued).then_some(cur + 1)
})
.is_ok()
}
/// Release a slot reserved by `try_reserve` that will not be spawned
/// (idempotency replay/conflict).
pub fn release_reservation(&self) {
self.dec_queued();
}
pub fn register(&self, run_id: String, token: CancellationToken) {
self.tokens.insert(run_id, token);
}
/// Queued → running: the run acquired its execution permit.
pub fn mark_running(&self) {
self.dec_queued();
self.in_flight.fetch_add(1, Ordering::AcqRel);
}
/// Running transition for a run that never occupied a local queue slot (the
/// cluster claim path: `submit` writes Pending + releases its reservation, so
/// no queued slot exists to consume). Only bumps `in_flight`.
pub fn mark_running_unqueued(&self) {
self.in_flight.fetch_add(1, Ordering::AcqRel);
}
/// Running → terminal: drop the token and wake any shutdown drain waiter.
pub fn mark_finished(&self, run_id: &str) {
self.dec_in_flight();
self.tokens.remove(run_id);
self.drained.notify_waiters();
}
/// Queued → terminal: a run was cancelled (or the server shut down) while
/// still waiting for an execution permit, so it never became in-flight.
/// Release its **queue** slot (not `in_flight`), drop the token, and wake any
/// drain waiter. Distinct from [`Self::mark_finished`], which decrements
/// `in_flight` (#146 R: a cancel on a queued run now takes effect at once).
pub fn mark_queued_cancelled(&self, run_id: &str) {
self.dec_queued();
self.tokens.remove(run_id);
self.drained.notify_waiters();
}
/// Saturating decrement of `queued` (never wraps below 0 — a wrap to
/// usize::MAX would permanently fail `try_reserve` and wedge the server).
fn dec_queued(&self) {
let _ = self
.queued
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |q| {
Some(q.saturating_sub(1))
});
}
/// Saturating decrement of `in_flight`.
fn dec_in_flight(&self) {
let _ = self
.in_flight
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |n| {
Some(n.saturating_sub(1))
});
}
/// Cancel a live run. Returns `true` if a live token existed.
pub fn cancel(&self, run_id: &str) -> bool {
if let Some(t) = self.tokens.get(run_id) {
t.cancel();
true
} else {
false
}
}
pub fn queued(&self) -> usize {
self.queued.load(Ordering::Acquire)
}
pub fn in_flight(&self) -> usize {
self.in_flight.load(Ordering::Acquire)
}
pub fn is_full(&self) -> bool {
self.queued() >= self.max_queued
}
/// Resolve once no run is queued or in flight. Arms the notification *before*
/// re-checking so a transition can't be missed.
pub async fn wait_drained(&self) {
loop {
if self.queued() == 0 && self.in_flight() == 0 {
return;
}
let notified = self.drained.notified();
if self.queued() == 0 && self.in_flight() == 0 {
return;
}
notified.await;
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn reserve_respects_capacity() {
let r = Registry::new(2);
assert!(r.try_reserve());
assert!(r.try_reserve());
assert!(!r.try_reserve());
assert!(r.is_full());
r.release_reservation();
assert!(r.try_reserve());
}
#[test]
fn running_transition_moves_counters() {
let r = Registry::new(4);
r.try_reserve();
assert_eq!(r.queued(), 1);
r.mark_running();
assert_eq!(r.queued(), 0);
assert_eq!(r.in_flight(), 1);
r.mark_finished("x");
assert_eq!(r.in_flight(), 0);
}
#[test]
fn mark_running_unqueued_only_bumps_in_flight() {
let r = Registry::new(4);
// No reservation taken (cluster claim path).
r.mark_running_unqueued();
assert_eq!(
r.queued(),
0,
"queued must NOT be decremented (no slot was held)"
);
assert_eq!(r.in_flight(), 1);
r.mark_finished("x");
assert_eq!(r.in_flight(), 0);
assert_eq!(r.queued(), 0);
}
#[test]
fn queued_decrement_saturates_at_zero() {
let r = Registry::new(4);
// A spurious decrement at 0 must NOT wrap to usize::MAX (that would
// permanently fail try_reserve and wedge backpressure — #228).
r.mark_running(); // dec_queued() at 0 + in_flight++
assert_eq!(r.queued(), 0, "saturating: stays 0, never usize::MAX");
assert!(
r.try_reserve(),
"try_reserve still works (queued not wrapped)"
);
}
#[test]
fn cancel_reports_presence() {
let r = Registry::new(4);
let token = CancellationToken::new();
r.register("run1".into(), token.clone());
assert!(r.cancel("run1"));
assert!(token.is_cancelled());
assert!(!r.cancel("missing"));
}
#[test]
fn is_not_idle_while_queued() {
let r = Registry::new(4);
r.try_reserve();
// queued=1, in_flight=0 — must NOT be considered drained.
assert_eq!(r.queued(), 1);
assert_eq!(r.in_flight(), 0);
}
#[tokio::test]
async fn wait_drained_returns_when_idle() {
let r = Registry::new(4);
// No in-flight work → resolves immediately.
r.wait_drained().await;
}
}