faucet_cli/serve/
registry.rs

1//! In-flight run registry. Tracks per-run cancellation tokens and the queue /
2//! in-flight counters that drive backpressure (429), the `faucet_serve_runs_*`
3//! gauges, `/readyz`, and the shutdown drain. A "queued" run is one that has been
4//! spawned but has not yet acquired an execution permit.
5
6use dashmap::DashMap;
7use std::sync::atomic::{AtomicUsize, Ordering};
8use tokio::sync::Notify;
9use tokio_util::sync::CancellationToken;
10
11pub struct Registry {
12    tokens: DashMap<String, CancellationToken>,
13    queued: AtomicUsize,
14    in_flight: AtomicUsize,
15    max_queued: usize,
16    drained: Notify,
17}
18
19impl Registry {
20    pub fn new(max_queued: usize) -> Self {
21        Self {
22            tokens: DashMap::new(),
23            queued: AtomicUsize::new(0),
24            in_flight: AtomicUsize::new(0),
25            max_queued: max_queued.max(1),
26            drained: Notify::new(),
27        }
28    }
29
30    /// Reserve a queue slot. Returns `true` if a slot was reserved, or `false`
31    /// if the queue is full. Atomic against concurrent submits via a CAS loop.
32    pub fn try_reserve(&self) -> bool {
33        let mut cur = self.queued.load(Ordering::Acquire);
34        loop {
35            if cur >= self.max_queued {
36                return false;
37            }
38            match self.queued.compare_exchange_weak(
39                cur,
40                cur + 1,
41                Ordering::AcqRel,
42                Ordering::Acquire,
43            ) {
44                Ok(_) => return true,
45                Err(actual) => cur = actual,
46            }
47        }
48    }
49
50    /// Release a slot reserved by `try_reserve` that will not be spawned
51    /// (idempotency replay/conflict).
52    pub fn release_reservation(&self) {
53        self.queued.fetch_sub(1, Ordering::AcqRel);
54    }
55
56    pub fn register(&self, run_id: String, token: CancellationToken) {
57        self.tokens.insert(run_id, token);
58    }
59
60    /// Queued → running: the run acquired its execution permit.
61    pub fn mark_running(&self) {
62        self.queued.fetch_sub(1, Ordering::AcqRel);
63        self.in_flight.fetch_add(1, Ordering::AcqRel);
64    }
65
66    /// Running → terminal: drop the token and wake any shutdown drain waiter.
67    pub fn mark_finished(&self, run_id: &str) {
68        self.in_flight.fetch_sub(1, Ordering::AcqRel);
69        self.tokens.remove(run_id);
70        self.drained.notify_waiters();
71    }
72
73    /// Queued → terminal: a run was cancelled (or the server shut down) while
74    /// still waiting for an execution permit, so it never became in-flight.
75    /// Release its **queue** slot (not `in_flight`), drop the token, and wake any
76    /// drain waiter. Distinct from [`Self::mark_finished`], which decrements
77    /// `in_flight` (#146 R: a cancel on a queued run now takes effect at once).
78    pub fn mark_queued_cancelled(&self, run_id: &str) {
79        self.queued.fetch_sub(1, Ordering::AcqRel);
80        self.tokens.remove(run_id);
81        self.drained.notify_waiters();
82    }
83
84    /// Cancel a live run. Returns `true` if a live token existed.
85    pub fn cancel(&self, run_id: &str) -> bool {
86        if let Some(t) = self.tokens.get(run_id) {
87            t.cancel();
88            true
89        } else {
90            false
91        }
92    }
93
94    pub fn queued(&self) -> usize {
95        self.queued.load(Ordering::Acquire)
96    }
97
98    pub fn in_flight(&self) -> usize {
99        self.in_flight.load(Ordering::Acquire)
100    }
101
102    pub fn is_full(&self) -> bool {
103        self.queued() >= self.max_queued
104    }
105
106    /// Resolve once no run is queued or in flight. Arms the notification *before*
107    /// re-checking so a transition can't be missed.
108    pub async fn wait_drained(&self) {
109        loop {
110            if self.queued() == 0 && self.in_flight() == 0 {
111                return;
112            }
113            let notified = self.drained.notified();
114            if self.queued() == 0 && self.in_flight() == 0 {
115                return;
116            }
117            notified.await;
118        }
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn reserve_respects_capacity() {
128        let r = Registry::new(2);
129        assert!(r.try_reserve());
130        assert!(r.try_reserve());
131        assert!(!r.try_reserve());
132        assert!(r.is_full());
133        r.release_reservation();
134        assert!(r.try_reserve());
135    }
136
137    #[test]
138    fn running_transition_moves_counters() {
139        let r = Registry::new(4);
140        r.try_reserve();
141        assert_eq!(r.queued(), 1);
142        r.mark_running();
143        assert_eq!(r.queued(), 0);
144        assert_eq!(r.in_flight(), 1);
145        r.mark_finished("x");
146        assert_eq!(r.in_flight(), 0);
147    }
148
149    #[test]
150    fn cancel_reports_presence() {
151        let r = Registry::new(4);
152        let token = CancellationToken::new();
153        r.register("run1".into(), token.clone());
154        assert!(r.cancel("run1"));
155        assert!(token.is_cancelled());
156        assert!(!r.cancel("missing"));
157    }
158
159    #[test]
160    fn is_not_idle_while_queued() {
161        let r = Registry::new(4);
162        r.try_reserve();
163        // queued=1, in_flight=0 — must NOT be considered drained.
164        assert_eq!(r.queued(), 1);
165        assert_eq!(r.in_flight(), 0);
166    }
167
168    #[tokio::test]
169    async fn wait_drained_returns_when_idle() {
170        let r = Registry::new(4);
171        // No in-flight work → resolves immediately.
172        r.wait_drained().await;
173    }
174}
faucet_cli/serve/registry.rs

faucet_cli/serve/
registry.rs
