Ferrit Explore
中文·繁體·EN·日本語 Sign in Register
cielxl / veld / src / core / timer.rs
use std::time::Duration;
use tokio::time::Instant;

/// Describes which type of timeout has expired on a connection.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TimeoutKind {
    Read,
    Write,
    Idle,
}

/// Per-connection timeout tracker.
///
/// Holds three independent deadlines — read, write, and idle — each of which
/// can be armed, reset, and checked independently.  The struct is cheap to
/// create and the hot path (`check_expired`) performs only timestamp
/// comparisons, so it is safe to call on every event-loop tick.
#[derive(Debug)]
pub struct ConnectionTimer {
    read_deadline: Option<Instant>,
    write_deadline: Option<Instant>,
    idle_deadline: Option<Instant>,

    read_timeout: Duration,
    write_timeout: Duration,
    idle_timeout: Duration,
}

impl ConnectionTimer {
    /// Create a new timer with the given timeout durations.
    ///
    /// All deadlines start unarmed — call the `reset_*` methods to arm them
    /// once the connection is established.
    pub fn new(read_timeout: Duration, write_timeout: Duration, idle_timeout: Duration) -> Self {
        Self {
            read_deadline: None,
            write_deadline: None,
            idle_deadline: None,
            read_timeout,
            write_timeout,
            idle_timeout,
        }
    }

    /// Arm / re-arm the read deadline from *now*.
    pub fn reset_read(&mut self) {
        self.read_deadline = Some(Instant::now() + self.read_timeout);
    }

    /// Arm / re-arm the write deadline from *now*.
    pub fn reset_write(&mut self) {
        self.write_deadline = Some(Instant::now() + self.write_timeout);
    }

    /// Arm / re-arm the idle deadline from *now*.
    pub fn reset_idle(&mut self) {
        self.idle_deadline = Some(Instant::now() + self.idle_timeout);
    }

    /// Disable the read deadline (e.g. after a successful read).
    pub fn clear_read(&mut self) {
        self.read_deadline = None;
    }

    /// Disable the write deadline (e.g. after a successful write).
    pub fn clear_write(&mut self) {
        self.write_deadline = None;
    }

    /// Disable the idle deadline.
    pub fn clear_idle(&mut self) {
        self.idle_deadline = None;
    }

    /// Check whether any armed deadline has passed.
    ///
    /// Returns `Some(kind)` for the first expired deadline found (priority:
    /// read > write > idle), or `None` if everything is still within its
    /// window.  Unarmed deadlines are never considered expired.
    pub fn check_expired(&self) -> Option<TimeoutKind> {
        let now = Instant::now();

        if let Some(dl) = self.read_deadline {
            if now >= dl {
                return Some(TimeoutKind::Read);
            }
        }

        if let Some(dl) = self.write_deadline {
            if now >= dl {
                return Some(TimeoutKind::Write);
            }
        }

        if let Some(dl) = self.idle_deadline {
            if now >= dl {
                return Some(TimeoutKind::Idle);
            }
        }

        None
    }

    /// Convenience: returns `true` when any deadline has expired.
    pub fn is_expired(&self) -> bool {
        self.check_expired().is_some()
    }

    /// Returns the duration until the next deadline fires, or `None` if no
    /// deadline is armed.  Useful for passing to `tokio::time::timeout` or
    /// computing a sleep duration for the next event-loop tick.
    pub fn next_deadline(&self) -> Option<Duration> {
        let now = Instant::now();
        let mut earliest: Option<Instant> = None;

        for dl in [self.read_deadline, self.write_deadline, self.idle_deadline]
            .iter()
            .flatten()
        {
            if *dl > now {
                earliest = Some(match earliest {
                    Some(cur) if cur < *dl => cur,
                    _ => *dl,
                });
            }
        }

        earliest.map(|dl| dl - now)
    }
}

/// Higher-level timer manager that owns a set of [`ConnectionTimer`]s and
/// provides periodic sweep functionality.
///
/// In a production server you would typically run a background task that
/// calls [`Timer::sweep`] on every tick to close connections whose timers
/// have fired.
pub struct Timer {
    check_interval: Duration,
}

impl Timer {
    /// Create a new `Timer` that sweeps at the given interval.
    pub fn new(check_interval: Duration) -> Self {
        Self { check_interval }
    }

    /// Returns the configured sweep interval.
    pub fn check_interval(&self) -> Duration {
        self.check_interval
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::Duration;

    #[test]
    fn new_timer_has_no_expired_deadlines() {
        let ct = ConnectionTimer::new(
            Duration::from_secs(5),
            Duration::from_secs(5),
            Duration::from_secs(60),
        );
        assert!(!ct.is_expired());
        assert_eq!(ct.check_expired(), None);
    }

    #[test]
    fn reset_read_arms_deadline() {
        let mut ct = ConnectionTimer::new(
            Duration::from_secs(30),
            Duration::from_secs(30),
            Duration::from_secs(60),
        );
        ct.reset_read();
        // Should not be expired immediately with a 30-second window.
        assert!(!ct.is_expired());
        assert!(ct.next_deadline().is_some());
    }

    #[test]
    fn zero_duration_expires_immediately() {
        let mut ct = ConnectionTimer::new(
            Duration::ZERO,
            Duration::from_secs(30),
            Duration::from_secs(60),
        );
        ct.reset_read();
        // A zero-duration deadline is already in the past.
        assert!(ct.is_expired());
        assert_eq!(ct.check_expired(), Some(TimeoutKind::Read));
    }

    #[test]
    fn clear_removes_deadline() {
        let mut ct = ConnectionTimer::new(Duration::ZERO, Duration::ZERO, Duration::ZERO);
        ct.reset_read();
        assert!(ct.is_expired());
        ct.clear_read();
        // No other deadlines armed — should be clean.
        assert!(!ct.is_expired());
    }

    #[test]
    fn priority_order() {
        let mut ct = ConnectionTimer::new(Duration::ZERO, Duration::ZERO, Duration::ZERO);
        ct.reset_idle();
        ct.reset_write();
        ct.reset_read();
        // Read has highest priority.
        assert_eq!(ct.check_expired(), Some(TimeoutKind::Read));

        ct.clear_read();
        assert_eq!(ct.check_expired(), Some(TimeoutKind::Write));

        ct.clear_write();
        assert_eq!(ct.check_expired(), Some(TimeoutKind::Idle));
    }

    #[test]
    fn next_deadline_returns_shortest_remaining() {
        let mut ct = ConnectionTimer::new(
            Duration::from_secs(1),
            Duration::from_secs(5),
            Duration::from_secs(10),
        );
        ct.reset_read();
        ct.reset_write();
        let remaining = ct.next_deadline().unwrap();
        // Should be closer to 1 second than 5.
        assert!(remaining <= Duration::from_secs(1));
    }

    #[test]
    fn timer_manager_interval() {
        let t = Timer::new(Duration::from_millis(500));
        assert_eq!(t.check_interval(), Duration::from_millis(500));
    }
}