rust_expect/util/
timeout.rs

1//! Timeout utilities.
2//!
3//! This module provides utilities for handling timeouts in expect operations.
4
5use std::future::Future;
6use std::time::Duration;
7
8use tokio::time::{Timeout, timeout};
9
10/// Extension trait for adding timeouts to futures.
11pub trait TimeoutExt: Sized {
12    /// Wrap this future with a timeout.
13    fn with_timeout(self, duration: Duration) -> Timeout<Self>;
14
15    /// Wrap this future with a timeout in seconds.
16    fn with_timeout_secs(self, secs: u64) -> Timeout<Self> {
17        self.with_timeout(Duration::from_secs(secs))
18    }
19
20    /// Wrap this future with a timeout in milliseconds.
21    fn with_timeout_ms(self, ms: u64) -> Timeout<Self> {
22        self.with_timeout(Duration::from_millis(ms))
23    }
24}
25
26impl<F: Future> TimeoutExt for F {
27    fn with_timeout(self, duration: Duration) -> Timeout<Self> {
28        timeout(duration, self)
29    }
30}
31
32/// A timeout configuration.
33#[derive(Debug, Clone, Copy)]
34pub struct TimeoutConfig {
35    /// Default timeout for expect operations.
36    pub expect: Duration,
37    /// Timeout for connection operations.
38    pub connect: Duration,
39    /// Timeout for read operations.
40    pub read: Duration,
41    /// Timeout for write operations.
42    pub write: Duration,
43    /// Timeout for close operations.
44    pub close: Duration,
45}
46
47impl Default for TimeoutConfig {
48    fn default() -> Self {
49        Self {
50            expect: Duration::from_secs(30),
51            connect: Duration::from_secs(60),
52            read: Duration::from_secs(10),
53            write: Duration::from_secs(10),
54            close: Duration::from_secs(5),
55        }
56    }
57}
58
59impl TimeoutConfig {
60    /// Create a new timeout configuration with default values.
61    #[must_use]
62    pub fn new() -> Self {
63        Self::default()
64    }
65
66    /// Set the expect timeout.
67    #[must_use]
68    pub const fn expect(mut self, timeout: Duration) -> Self {
69        self.expect = timeout;
70        self
71    }
72
73    /// Set the connect timeout.
74    #[must_use]
75    pub const fn connect(mut self, timeout: Duration) -> Self {
76        self.connect = timeout;
77        self
78    }
79
80    /// Set the read timeout.
81    #[must_use]
82    pub const fn read(mut self, timeout: Duration) -> Self {
83        self.read = timeout;
84        self
85    }
86
87    /// Set the write timeout.
88    #[must_use]
89    pub const fn write(mut self, timeout: Duration) -> Self {
90        self.write = timeout;
91        self
92    }
93
94    /// Set the close timeout.
95    #[must_use]
96    pub const fn close(mut self, timeout: Duration) -> Self {
97        self.close = timeout;
98        self
99    }
100
101    /// Create a configuration with all timeouts set to the same value.
102    #[must_use]
103    pub const fn uniform(timeout: Duration) -> Self {
104        Self {
105            expect: timeout,
106            connect: timeout,
107            read: timeout,
108            write: timeout,
109            close: timeout,
110        }
111    }
112
113    /// Create a configuration with no timeouts (effectively infinite).
114    #[must_use]
115    pub const fn none() -> Self {
116        let max = Duration::from_secs(u64::MAX / 2);
117        Self::uniform(max)
118    }
119}
120
121/// A deadline tracker for operations with multiple steps.
122#[derive(Debug, Clone)]
123pub struct Deadline {
124    /// The deadline instant.
125    deadline: tokio::time::Instant,
126}
127
128impl Deadline {
129    /// Create a new deadline from now.
130    #[must_use]
131    pub fn from_now(duration: Duration) -> Self {
132        Self {
133            deadline: tokio::time::Instant::now() + duration,
134        }
135    }
136
137    /// Check if the deadline has passed.
138    #[must_use]
139    pub fn is_expired(&self) -> bool {
140        tokio::time::Instant::now() >= self.deadline
141    }
142
143    /// Get the remaining time until the deadline.
144    #[must_use]
145    pub fn remaining(&self) -> Duration {
146        self.deadline
147            .saturating_duration_since(tokio::time::Instant::now())
148    }
149
150    /// Check if there is time remaining.
151    #[must_use]
152    pub fn has_time(&self) -> bool {
153        !self.is_expired()
154    }
155
156    /// Sleep until the deadline.
157    pub async fn sleep(&self) {
158        let remaining = self.remaining();
159        if !remaining.is_zero() {
160            tokio::time::sleep(remaining).await;
161        }
162    }
163
164    /// Apply this deadline to a future.
165    pub fn apply<F: Future>(&self, future: F) -> Timeout<F> {
166        timeout(self.remaining(), future)
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    #[test]
175    fn timeout_config_default() {
176        let config = TimeoutConfig::default();
177        assert_eq!(config.expect, Duration::from_secs(30));
178    }
179
180    #[test]
181    fn timeout_config_uniform() {
182        let config = TimeoutConfig::uniform(Duration::from_secs(5));
183        assert_eq!(config.expect, Duration::from_secs(5));
184        assert_eq!(config.connect, Duration::from_secs(5));
185    }
186
187    #[tokio::test]
188    async fn deadline_remaining() {
189        let deadline = Deadline::from_now(Duration::from_secs(10));
190        assert!(deadline.has_time());
191        assert!(deadline.remaining() > Duration::from_secs(9));
192    }
193
194    #[tokio::test]
195    async fn timeout_ext() {
196        let result = async { 42 }.with_timeout(Duration::from_secs(1)).await;
197        assert!(result.is_ok());
198        assert_eq!(result.unwrap(), 42);
199    }
200}
rust_expect/util/timeout.rs

rust_expect/util/
timeout.rs
