rama_net/proxy/
idle.rs

1//! [`IdleGuard`] — a small helper for detecting idleness in a `tokio::select!` loop.
2
3use std::pin::Pin;
4use std::time::Duration;
5
6use tokio::time::{Instant, Sleep};
7
8/// A resettable idle deadline.
9///
10/// Designed to be used as one arm of a `tokio::select!`: poll [`IdleGuard::tick`]
11/// in the select; when another arm fires (i.e. progress was observed) call
12/// [`IdleGuard::reset`] before re-entering the select to extend the idle window.
13///
14/// If the inner [`Sleep`] elapses before [`reset`](IdleGuard::reset) is called,
15/// the idle window has lapsed.
16///
17/// `IdleGuard` is intended for cases where the watched activity does not itself
18/// produce values that can be selected on (e.g. byte progress observed inside
19/// another future). For values that can be selected on, prefer racing the
20/// activity directly.
21pub struct IdleGuard {
22    timeout: Duration,
23    sleep: Pin<Box<Sleep>>,
24}
25
26impl core::fmt::Debug for IdleGuard {
27    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
28        f.debug_struct("IdleGuard")
29            .field("timeout", &self.timeout)
30            .finish_non_exhaustive()
31    }
32}
33
34impl IdleGuard {
35    /// Create a new [`IdleGuard`] that fires after `timeout` of inactivity.
36    #[must_use]
37    pub fn new(timeout: Duration) -> Self {
38        Self {
39            timeout,
40            sleep: Box::pin(tokio::time::sleep(timeout)),
41        }
42    }
43
44    /// The configured idle timeout.
45    #[must_use]
46    pub fn timeout(&self) -> Duration {
47        self.timeout
48    }
49
50    /// Reset the idle deadline to `now + timeout`. Call this whenever progress
51    /// has been observed, before re-entering the select that polls
52    /// [`IdleGuard::tick`].
53    pub fn reset(&mut self) {
54        let deadline = Instant::now() + self.timeout;
55        self.sleep.as_mut().reset(deadline);
56    }
57
58    /// Poll-able future that completes when the idle window has elapsed.
59    ///
60    /// Borrow-and-poll inside `tokio::select!`. Once it completes, the guard
61    /// is considered tripped — call [`IdleGuard::reset`] before re-arming if
62    /// you want to keep waiting.
63    pub fn tick(&mut self) -> &mut Pin<Box<Sleep>> {
64        &mut self.sleep
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use std::time::Duration;
71
72    use super::*;
73
74    #[tokio::test(start_paused = true)]
75    async fn idle_guard_fires_after_timeout() {
76        let mut guard = IdleGuard::new(Duration::from_millis(100));
77        tokio::time::advance(Duration::from_millis(101)).await;
78        // Should be ready immediately after enough virtual time has passed.
79        guard.tick().await;
80    }
81
82    #[tokio::test(start_paused = true)]
83    async fn idle_guard_reset_extends_window() {
84        let mut guard = IdleGuard::new(Duration::from_millis(100));
85
86        // Advance most of the way, then reset.
87        tokio::time::advance(Duration::from_millis(80)).await;
88        guard.reset();
89
90        // Advance another 80ms — would have fired without reset.
91        tokio::time::advance(Duration::from_millis(80)).await;
92
93        // Race a short timeout against the guard; the guard should NOT have
94        // fired yet (we've only spent 80ms since reset out of 100ms window).
95        tokio::select! {
96            biased;
97            _ = guard.tick() => panic!("idle guard fired prematurely"),
98            _ = tokio::time::sleep(Duration::from_millis(0)) => {}
99        }
100
101        // Advance past the reset deadline.
102        tokio::time::advance(Duration::from_millis(30)).await;
103        guard.tick().await;
104    }
105}
rama_net/proxy/idle.rs

rama_net/proxy/
idle.rs
