pitchfork_cli/proxy/
activity.rs

1//! Per-daemon proxy activity: what the proxy is carrying for each daemon right
2//! now, and when it last carried anything.
3//!
4//! Every unit of work the proxy does for a daemon holds an [`ActivityGuard`]
5//! for as long as it lasts:
6//!
7//! - an HTTP request, until its response body has been sent in full, so a
8//!   streamed or server-sent-events response counts for as long as it streams;
9//! - an upgraded connection (WebSocket), until either side closes it;
10//! - a TLS passthrough connection, until either side closes it;
11//! - an auto-start, for the daemon and every dependency it brings up.
12//!
13//! An idle keep-alive connection holds nothing: it is the requests on it that
14//! count. Name resolution is not activity either, since a browser resolves
15//! names speculatively.
16//!
17//! The same lock that counts activity also records a daemon being stopped for
18//! inactivity, so the two cannot interleave: a stop is only claimed while
19//! nothing is in flight, and nothing new begins while a stop is claimed.
20
21use std::collections::HashMap;
22use std::sync::Mutex;
23use std::time::{Duration, Instant};
24
25use crate::daemon_id::DaemonId;
26
27/// The process-wide tracker the proxy and the supervisor share.
28pub(crate) static ACTIVITY: once_cell::sync::Lazy<ActivityTracker> =
29    once_cell::sync::Lazy::new(ActivityTracker::default);
30
31#[derive(Debug)]
32struct Entry {
33    /// Units of work under way.
34    in_flight: usize,
35    /// When the last unit of work began or ended, or when the daemon was first
36    /// looked at if nothing has happened since.
37    last_activity: Instant,
38    /// A stop for inactivity has been claimed and not yet released.
39    idle_stopping: bool,
40}
41
42impl Entry {
43    fn new(now: Instant) -> Self {
44        Self {
45            in_flight: 0,
46            last_activity: now,
47            idle_stopping: false,
48        }
49    }
50}
51
52/// Activity for every daemon the proxy has done anything for.
53#[derive(Debug, Default)]
54pub(crate) struct ActivityTracker {
55    entries: Mutex<HashMap<DaemonId, Entry>>,
56}
57
58/// How a daemon's activity looks from outside.
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60pub(crate) struct ActivitySnapshot {
61    pub in_flight: usize,
62    pub idle_for: Duration,
63    pub idle_stopping: bool,
64}
65
66impl ActivityTracker {
67    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<DaemonId, Entry>> {
68        // Bookkeeping only; a panic elsewhere must not wedge the proxy.
69        self.entries.lock().unwrap_or_else(|e| e.into_inner())
70    }
71
72    /// Begin a unit of work for `id`, which lasts until the guard is dropped.
73    ///
74    /// Returns `None` while the daemon is being stopped for inactivity: the
75    /// caller should treat it as not running yet and try again once the stop
76    /// has finished, rather than forward to a daemon that is going away.
77    pub(crate) fn begin(&'static self, id: &DaemonId) -> Option<ActivityGuard> {
78        let now = Instant::now();
79        let mut entries = self.lock();
80        let entry = entries.entry(id.clone()).or_insert_with(|| Entry::new(now));
81        if entry.idle_stopping {
82            return None;
83        }
84        entry.in_flight += 1;
85        entry.last_activity = now;
86        Some(ActivityGuard {
87            tracker: self,
88            id: id.clone(),
89        })
90    }
91
92    /// Begin a unit of work for each of `ids`, or for none of them if any is
93    /// being stopped for inactivity.
94    pub(crate) fn begin_all(&'static self, ids: &[DaemonId]) -> Option<Vec<ActivityGuard>> {
95        let mut guards = Vec::with_capacity(ids.len());
96        for id in ids {
97            // Dropping the guards taken so far releases them.
98            guards.push(self.begin(id)?);
99        }
100        Some(guards)
101    }
102
103    fn end(&self, id: &DaemonId) {
104        let mut entries = self.lock();
105        if let Some(entry) = entries.get_mut(id) {
106            entry.in_flight = entry.in_flight.saturating_sub(1);
107            entry.last_activity = Instant::now();
108        }
109    }
110
111    /// Current activity for `id`.
112    ///
113    /// A daemon the proxy has never seen is recorded as of now, so a daemon
114    /// adopted from a previous supervisor gets a full grace period rather than
115    /// being judged idle on the first look.
116    pub(crate) fn snapshot(&self, id: &DaemonId) -> ActivitySnapshot {
117        let now = Instant::now();
118        let mut entries = self.lock();
119        let entry = entries.entry(id.clone()).or_insert_with(|| Entry::new(now));
120        ActivitySnapshot {
121            in_flight: entry.in_flight,
122            idle_for: now.saturating_duration_since(entry.last_activity),
123            idle_stopping: entry.idle_stopping,
124        }
125    }
126
127    /// Claim `id` for a stop, provided nothing is in flight and nothing has
128    /// happened for at least `grace`.
129    ///
130    /// While the claim is held, [`Self::begin`] refuses new work for the
131    /// daemon. Release it with [`Self::release_idle_stop`] whether or not the
132    /// stop went ahead.
133    pub(crate) fn claim_idle_stop(&self, id: &DaemonId, grace: Duration) -> bool {
134        let now = Instant::now();
135        let mut entries = self.lock();
136        let entry = entries.entry(id.clone()).or_insert_with(|| Entry::new(now));
137        if entry.idle_stopping
138            || entry.in_flight > 0
139            || now.saturating_duration_since(entry.last_activity) < grace
140        {
141            return false;
142        }
143        entry.idle_stopping = true;
144        true
145    }
146
147    /// Release a claim taken by [`Self::claim_idle_stop`].
148    ///
149    /// The idle clock restarts, so a stop that was called off does not leave
150    /// the daemon immediately eligible again.
151    pub(crate) fn release_idle_stop(&self, id: &DaemonId) {
152        let mut entries = self.lock();
153        if let Some(entry) = entries.get_mut(id) {
154            entry.idle_stopping = false;
155            entry.last_activity = Instant::now();
156        }
157    }
158
159    /// Whether a stop for inactivity is under way for `id`.
160    pub(crate) fn is_idle_stopping(&self, id: &DaemonId) -> bool {
161        self.lock().get(id).is_some_and(|e| e.idle_stopping)
162    }
163}
164
165/// One unit of proxy work for a daemon. Dropping it ends the work.
166#[must_use = "activity ends as soon as the guard is dropped"]
167pub(crate) struct ActivityGuard {
168    tracker: &'static ActivityTracker,
169    id: DaemonId,
170}
171
172impl std::fmt::Debug for ActivityGuard {
173    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
174        f.debug_struct("ActivityGuard")
175            .field("id", &self.id)
176            .finish()
177    }
178}
179
180impl Drop for ActivityGuard {
181    fn drop(&mut self) {
182        self.tracker.end(&self.id);
183    }
184}
185
186/// A response body that holds an [`ActivityGuard`] until it has been sent in
187/// full or abandoned, so a long download or an event stream counts as
188/// activity for as long as it lasts.
189pub(crate) struct GuardedBody<B> {
190    inner: B,
191    _guard: Option<ActivityGuard>,
192}
193
194impl<B> GuardedBody<B> {
195    pub(crate) fn new(inner: B, guard: Option<ActivityGuard>) -> Self {
196        Self {
197            inner,
198            _guard: guard,
199        }
200    }
201}
202
203impl<B> hyper::body::Body for GuardedBody<B>
204where
205    B: hyper::body::Body + Unpin,
206{
207    type Data = B::Data;
208    type Error = B::Error;
209
210    fn poll_frame(
211        mut self: std::pin::Pin<&mut Self>,
212        cx: &mut std::task::Context<'_>,
213    ) -> std::task::Poll<Option<Result<hyper::body::Frame<Self::Data>, Self::Error>>> {
214        std::pin::Pin::new(&mut self.inner).poll_frame(cx)
215    }
216
217    fn is_end_stream(&self) -> bool {
218        self.inner.is_end_stream()
219    }
220
221    fn size_hint(&self) -> hyper::body::SizeHint {
222        self.inner.size_hint()
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    fn tracker() -> &'static ActivityTracker {
231        Box::leak(Box::default())
232    }
233
234    fn id(name: &str) -> DaemonId {
235        DaemonId::new("proj", name)
236    }
237
238    #[test]
239    fn in_flight_work_blocks_an_idle_claim() {
240        let t = tracker();
241        let guard = t.begin(&id("api")).unwrap();
242        assert!(!t.claim_idle_stop(&id("api"), Duration::ZERO));
243        drop(guard);
244        assert!(t.claim_idle_stop(&id("api"), Duration::ZERO));
245    }
246
247    #[test]
248    fn recent_activity_blocks_an_idle_claim_until_the_grace_passes() {
249        let t = tracker();
250        drop(t.begin(&id("api")).unwrap());
251        assert!(!t.claim_idle_stop(&id("api"), Duration::from_secs(60)));
252        assert!(t.claim_idle_stop(&id("api"), Duration::ZERO));
253    }
254
255    #[test]
256    fn a_claimed_stop_refuses_new_work_until_released() {
257        let t = tracker();
258        assert!(t.claim_idle_stop(&id("api"), Duration::ZERO));
259        assert!(t.is_idle_stopping(&id("api")));
260        assert!(t.begin(&id("api")).is_none());
261        // A second claim for the same daemon is refused too.
262        assert!(!t.claim_idle_stop(&id("api"), Duration::ZERO));
263        t.release_idle_stop(&id("api"));
264        assert!(!t.is_idle_stopping(&id("api")));
265        assert!(t.begin(&id("api")).is_some());
266    }
267
268    #[test]
269    fn begin_all_takes_nothing_when_one_daemon_is_stopping() {
270        let t = tracker();
271        assert!(t.claim_idle_stop(&id("db"), Duration::ZERO));
272        assert!(t.begin_all(&[id("api"), id("db")]).is_none());
273        // The guard briefly taken on `api` was released again.
274        assert_eq!(t.snapshot(&id("api")).in_flight, 0);
275    }
276
277    #[test]
278    fn a_daemon_first_seen_by_a_snapshot_starts_a_fresh_idle_clock() {
279        let t = tracker();
280        let snap = t.snapshot(&id("adopted"));
281        assert_eq!(snap.in_flight, 0);
282        assert!(snap.idle_for < Duration::from_secs(1));
283        assert!(!t.claim_idle_stop(&id("adopted"), Duration::from_secs(60)));
284    }
285
286    #[test]
287    fn releasing_a_claim_restarts_the_idle_clock() {
288        let t = tracker();
289        assert!(t.claim_idle_stop(&id("api"), Duration::ZERO));
290        t.release_idle_stop(&id("api"));
291        assert!(!t.claim_idle_stop(&id("api"), Duration::from_secs(60)));
292    }
293
294    #[test]
295    fn guarded_body_ends_activity_when_dropped() {
296        let t = tracker();
297        let body = GuardedBody::new(
298            http_body_util::Empty::<hyper::body::Bytes>::new(),
299            t.begin(&id("api")),
300        );
301        assert_eq!(t.snapshot(&id("api")).in_flight, 1);
302        drop(body);
303        assert_eq!(t.snapshot(&id("api")).in_flight, 0);
304    }
305}
pitchfork_cli/proxy/activity.rs

pitchfork_cli/proxy/
activity.rs
