async_foundation/net/
timeouts.rs

1use mio::Token;
2use std::time::Instant;
3
4use crate::net::TokenTimeout;
5
6pub struct Timeouts {
7    pub(crate) queue: Vec<TokenTimeout>,
8}
9
10impl Timeouts {
11    pub fn new() -> Timeouts {
12        Timeouts { queue: Vec::new() }
13    }
14
15    pub fn add(&mut self, token: Token, instant: Instant) {
16        let index = self
17            .get_index_by(|item| item.timeout > instant)
18            .map(|index| index + 1)
19            .unwrap_or(0);
20        self.queue.insert(index, TokenTimeout::new(token, instant));
21    }
22
23    pub fn remove(&mut self, token: Token) -> Option<TokenTimeout> {
24        self.get_token_index(token)
25            .map(|index| self.queue.remove(index))
26    }
27
28    fn get_index_by<F>(&self, func: F) -> Option<usize>
29    where
30        F: Fn(&TokenTimeout) -> bool,
31    {
32        for i in (0..self.queue.len()).rev() {
33            let token_timeout = &self.queue[i];
34            if func(token_timeout) {
35                return Some(i);
36            }
37        }
38        None
39    }
40
41    pub fn get_token_index(&self, token: Token) -> Option<usize> {
42        self.get_index_by(|item| item.token == token)
43    }
44
45    pub fn current(&self) -> Option<&TokenTimeout> {
46        self.queue.last()
47    }
48
49    pub(crate) fn pop_timeouts(&mut self, instant: Instant) -> Vec<TokenTimeout> {
50        if self.queue.len() == 0 {
51            return vec![];
52        }
53        let index = match self.get_index_by(|item| item.timeout > instant) {
54            Some(index) => index + 1,
55            None => 0,
56        };
57
58        let mut result = self.queue.split_off(index);
59        result.reverse();
60        result
61    }
62
63    pub fn is_empty(&self) -> bool {
64        self.queue.is_empty()
65    }
66
67    pub fn len(&self) -> usize {
68        self.queue.len()
69    }
70}
71
72// cargo test --features net
73#[cfg(test)]
74mod tests {
75    use std::time::Duration;
76
77    use super::*;
78
79    #[test]
80    fn test_pop_timeouts() {
81        let mut timeouts = Timeouts::new();
82        let instant = Instant::now();
83        timeouts.add(Token(0), instant + Duration::from_secs(1));
84        timeouts.add(Token(1), instant + Duration::from_secs(2));
85        timeouts.add(Token(3), instant + Duration::from_secs(3));
86
87        let result = timeouts.pop_timeouts(instant + Duration::from_secs(2));
88        assert_eq!(result.len(), 2);
89        assert_eq!(result[0].token, Token(0));
90        assert_eq!(result[1].token, Token(1));
91
92        let result = timeouts.pop_timeouts(instant + Duration::from_secs(2));
93        assert_eq!(result.len(), 0);
94
95        let result = timeouts.pop_timeouts(instant + Duration::from_secs(3));
96        assert_eq!(result.len(), 1);
97    }
98
99    #[test]
100    fn test_timeouts_new() {
101        let timeouts = Timeouts::new();
102        assert!(timeouts.is_empty());
103        assert_eq!(timeouts.len(), 0);
104    }
105
106    #[test]
107    fn test_timeouts_add_single() {
108        let mut timeouts = Timeouts::new();
109        let instant = Instant::now();
110        let token = Token(42);
111        
112        timeouts.add(token, instant);
113        
114        assert!(!timeouts.is_empty());
115        assert_eq!(timeouts.len(), 1);
116        
117        let current = timeouts.current();
118        assert!(current.is_some());
119        assert_eq!(current.unwrap().token, token);
120        assert_eq!(current.unwrap().timeout, instant);
121    }
122
123    #[test]
124    fn test_timeouts_add_multiple_ordered() {
125        let mut timeouts = Timeouts::new();
126        let instant = Instant::now();
127        
128        timeouts.add(Token(1), instant + Duration::from_secs(3));
129        timeouts.add(Token(2), instant + Duration::from_secs(1));
130        timeouts.add(Token(3), instant + Duration::from_secs(2));
131        
132        assert_eq!(timeouts.len(), 3);
133        
134        // Should be ordered by timeout (shortest first)
135        let current = timeouts.current();
136        assert!(current.is_some());
137        assert_eq!(current.unwrap().token, Token(2)); // Shortest timeout
138    }
139
140    #[test]
141    fn test_timeouts_remove() {
142        let mut timeouts = Timeouts::new();
143        let instant = Instant::now();
144        
145        timeouts.add(Token(1), instant + Duration::from_secs(1));
146        timeouts.add(Token(2), instant + Duration::from_secs(2));
147        timeouts.add(Token(3), instant + Duration::from_secs(3));
148        
149        assert_eq!(timeouts.len(), 3);
150        
151        // Remove middle token
152        let removed = timeouts.remove(Token(2));
153        assert!(removed.is_some());
154        assert_eq!(removed.unwrap().token, Token(2));
155        assert_eq!(timeouts.len(), 2);
156        
157        // Remove non-existent token
158        let removed = timeouts.remove(Token(99));
159        assert!(removed.is_none());
160        assert_eq!(timeouts.len(), 2);
161    }
162
163    #[test]
164    fn test_timeouts_get_token_index() {
165        let mut timeouts = Timeouts::new();
166        let instant = Instant::now();
167        
168        timeouts.add(Token(1), instant + Duration::from_secs(1));
169        timeouts.add(Token(2), instant + Duration::from_secs(2));
170        timeouts.add(Token(3), instant + Duration::from_secs(3));
171        
172        assert_eq!(timeouts.get_token_index(Token(1)), Some(2));
173        assert_eq!(timeouts.get_token_index(Token(2)), Some(1));
174        assert_eq!(timeouts.get_token_index(Token(3)), Some(0));
175        assert_eq!(timeouts.get_token_index(Token(99)), None);
176    }
177
178    #[test]
179    fn test_timeouts_current() {
180        let mut timeouts = Timeouts::new();
181        let instant = Instant::now();
182        
183        assert!(timeouts.current().is_none());
184        
185        timeouts.add(Token(1), instant + Duration::from_secs(1));
186        let current = timeouts.current();
187        assert!(current.is_some());
188        assert_eq!(current.unwrap().token, Token(1));
189        
190        timeouts.add(Token(2), instant + Duration::from_secs(2));
191        timeouts.add(Token(3), instant + Duration::from_secs(3));
192        
193        let current = timeouts.current();
194        assert!(current.is_some());
195        assert_eq!(current.unwrap().token, Token(1)); // Shortest timeout
196    }
197
198    #[test]
199    fn test_timeouts_pop_timeouts_empty() {
200        let mut timeouts = Timeouts::new();
201        let instant = Instant::now();
202        
203        let result = timeouts.pop_timeouts(instant);
204        assert!(result.is_empty());
205    }
206
207    #[test]
208    fn test_timeouts_pop_timeouts_all() {
209        let mut timeouts = Timeouts::new();
210        let instant = Instant::now();
211        
212        timeouts.add(Token(1), instant + Duration::from_secs(1));
213        timeouts.add(Token(2), instant + Duration::from_secs(2));
214        timeouts.add(Token(3), instant + Duration::from_secs(3));
215        
216        // Pop all timeouts
217        let result = timeouts.pop_timeouts(instant + Duration::from_secs(5));
218        assert_eq!(result.len(), 3);
219        assert_eq!(result[0].token, Token(1));
220        assert_eq!(result[1].token, Token(2));
221        assert_eq!(result[2].token, Token(3));
222        
223        assert!(timeouts.is_empty());
224    }
225
226    #[test]
227    fn test_timeouts_pop_timeouts_partial() {
228        let mut timeouts = Timeouts::new();
229        let instant = Instant::now();
230        
231        timeouts.add(Token(1), instant + Duration::from_secs(1));
232        timeouts.add(Token(2), instant + Duration::from_secs(2));
233        timeouts.add(Token(3), instant + Duration::from_secs(3));
234        timeouts.add(Token(4), instant + Duration::from_secs(4));
235        
236        let result = timeouts.pop_timeouts(instant + Duration::from_secs(2));
237        assert_eq!(result.len(), 2);
238        assert_eq!(result[0].token, Token(1));
239        assert_eq!(result[1].token, Token(2));
240        assert_eq!(timeouts.len(), 2);
241        
242        let result = timeouts.pop_timeouts(instant + Duration::from_secs(5));
243        assert_eq!(result.len(), 2);
244        assert_eq!(result[0].token, Token(3));
245        assert_eq!(result[1].token, Token(4));
246        
247        assert!(timeouts.is_empty());
248    }
249
250    #[test]
251    fn test_timeouts_ordering() {
252        let mut timeouts = Timeouts::new();
253        let instant = Instant::now();
254        
255        timeouts.add(Token(5), instant + Duration::from_secs(5));
256        timeouts.add(Token(3), instant + Duration::from_secs(3));
257        timeouts.add(Token(1), instant + Duration::from_secs(1));
258        timeouts.add(Token(4), instant + Duration::from_secs(4));
259        timeouts.add(Token(2), instant + Duration::from_secs(2));
260        
261        let result = timeouts.pop_timeouts(instant + Duration::from_secs(10));
262        assert_eq!(result.len(), 5);
263        assert_eq!(result[0].token, Token(1));
264        assert_eq!(result[1].token, Token(2));
265        assert_eq!(result[2].token, Token(3));
266        assert_eq!(result[3].token, Token(4));
267        assert_eq!(result[4].token, Token(5));
268    }
269}
async_foundation/net/timeouts.rs

async_foundation/net/
timeouts.rs
