mz_ore/netio/
timeout.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! `Async{Read,Write}` wrappers that enforce a configurable timeout on each I/O operation.
17
18use std::io;
19use std::pin::Pin;
20use std::task::{Context, Poll};
21use std::time::Duration;
22
23use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
24use tokio::time::{Instant, Sleep};
25
26#[derive(Debug)]
27struct IoTimeout {
28    duration: Duration,
29    /// Deadline of the blocked operation, if any.
30    deadline: Option<Instant>,
31    /// Timer allocated on the first blocked operation and reused afterwards.
32    /// Its deadline can trail `deadline`, see `poll_expired`.
33    sleep: Option<Pin<Box<Sleep>>>,
34}
35
36impl IoTimeout {
37    fn new(duration: Duration) -> Self {
38        Self {
39            duration,
40            deadline: None,
41            sleep: None,
42        }
43    }
44
45    fn poll<T>(
46        &mut self,
47        cx: &mut Context<'_>,
48        poll_io: impl FnOnce(&mut Context<'_>) -> Poll<io::Result<T>>,
49    ) -> Poll<io::Result<T>> {
50        // Record the deadline before polling I/O, but only touch the timer if
51        // I/O blocks. A deadline past the representable range never expires.
52        let deadline = self
53            .deadline
54            .or_else(|| Instant::now().checked_add(self.duration));
55
56        // Poll I/O before the timer, as `tokio::time::timeout` does. If the
57        // task was not polled until after the deadline, for example because
58        // its runtime stalled, data that arrived in the meantime must win.
59        // Checking the timer first would time out connections to healthy
60        // peers that kept sending.
61        let poll = poll_io(cx);
62        if poll.is_ready() {
63            // Leave the timer registered. Deregistering costs as much as the
64            // re-arm in `poll_expired`, which only happens if it fires.
65            self.deadline = None;
66            return poll;
67        }
68
69        if let Some(deadline) = deadline {
70            if self.poll_expired(cx, deadline) {
71                self.deadline = None;
72                return Poll::Ready(Err(io::ErrorKind::TimedOut.into()));
73            }
74            self.deadline = Some(deadline);
75        }
76        Poll::Pending
77    }
78
79    /// Reports whether `deadline` has passed, registering `cx` for a wakeup
80    /// otherwise.
81    fn poll_expired(&mut self, cx: &mut Context<'_>, deadline: Instant) -> bool {
82        let sleep = self
83            .sleep
84            .get_or_insert_with(|| Box::pin(tokio::time::sleep_until(deadline)));
85        // The timer keeps the deadline of an earlier operation until it fires,
86        // so completing and starting operations does not touch the timer
87        // wheel. Deadlines only grow, so the timer never fires late. When it
88        // fires early, re-arm it to the current deadline.
89        loop {
90            if sleep.as_mut().poll(cx).is_pending() {
91                return false;
92            }
93            if sleep.deadline() >= deadline {
94                return true;
95            }
96            sleep.as_mut().reset(deadline);
97        }
98    }
99}
100
101/// An [`AsyncRead`] wrapper that enforces a timeout on each read.
102#[derive(Debug)]
103pub struct TimedReader<R> {
104    reader: R,
105    timeout: IoTimeout,
106}
107
108impl<R> TimedReader<R> {
109    /// Wrap a reader with a timeout.
110    pub fn new(reader: R, timeout: Duration) -> Self {
111        Self {
112            reader,
113            timeout: IoTimeout::new(timeout),
114        }
115    }
116}
117
118impl<R: AsyncRead + Unpin> AsyncRead for TimedReader<R> {
119    fn poll_read(
120        self: Pin<&mut Self>,
121        cx: &mut Context<'_>,
122        buf: &mut ReadBuf<'_>,
123    ) -> Poll<io::Result<()>> {
124        let this = self.get_mut();
125        this.timeout
126            .poll(cx, |cx| Pin::new(&mut this.reader).poll_read(cx, buf))
127    }
128}
129
130/// An [`AsyncWrite`] wrapper that enforces a timeout on each write.
131#[derive(Debug)]
132pub struct TimedWriter<W> {
133    writer: W,
134    timeout: IoTimeout,
135}
136
137impl<W> TimedWriter<W> {
138    /// Wrap a writer with a timeout.
139    pub fn new(writer: W, timeout: Duration) -> Self {
140        Self {
141            writer,
142            timeout: IoTimeout::new(timeout),
143        }
144    }
145}
146
147impl<W: AsyncWrite + Unpin> AsyncWrite for TimedWriter<W> {
148    fn poll_write(
149        self: Pin<&mut Self>,
150        cx: &mut Context<'_>,
151        buf: &[u8],
152    ) -> Poll<Result<usize, io::Error>> {
153        let this = self.get_mut();
154        this.timeout
155            .poll(cx, |cx| Pin::new(&mut this.writer).poll_write(cx, buf))
156    }
157
158    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
159        let this = self.get_mut();
160        this.timeout
161            .poll(cx, |cx| Pin::new(&mut this.writer).poll_flush(cx))
162    }
163
164    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
165        let this = self.get_mut();
166        this.timeout
167            .poll(cx, |cx| Pin::new(&mut this.writer).poll_shutdown(cx))
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174    use futures::task::noop_waker;
175
176    #[crate::test]
177    fn ready_io_does_not_create_a_timer() {
178        let waker = noop_waker();
179        let mut cx = Context::from_waker(&waker);
180        let mut reader = TimedReader::new(tokio::io::empty(), Duration::from_secs(1));
181        let mut bytes = [0; 8];
182        let mut buf = ReadBuf::new(&mut bytes);
183        assert!(
184            Pin::new(&mut reader)
185                .poll_read(&mut cx, &mut buf)
186                .is_ready()
187        );
188        assert_eq!(buf.filled().len(), 0);
189        assert!(reader.timeout.sleep.is_none());
190        let mut writer = TimedWriter::new(tokio::io::sink(), Duration::from_secs(1));
191        assert!(matches!(
192            Pin::new(&mut writer).poll_write(&mut cx, b"hello"),
193            Poll::Ready(Ok(5))
194        ));
195        assert!(writer.timeout.sleep.is_none());
196        assert!(matches!(
197            Pin::new(&mut writer).poll_flush(&mut cx),
198            Poll::Ready(Ok(()))
199        ));
200        assert!(writer.timeout.sleep.is_none());
201        assert!(matches!(
202            Pin::new(&mut writer).poll_shutdown(&mut cx),
203            Poll::Ready(Ok(()))
204        ));
205        assert!(writer.timeout.sleep.is_none());
206    }
207
208    #[crate::test(tokio::test(start_paused = true))]
209    async fn pending_repoll_keeps_deadline_and_ready_io_wins_over_expired_timer() {
210        let waker = noop_waker();
211        let mut cx = Context::from_waker(&waker);
212        let mut timeout = IoTimeout::new(Duration::from_millis(10));
213        assert!(timeout.poll::<()>(&mut cx, |_| Poll::Pending).is_pending());
214        let deadline = timeout.deadline.unwrap();
215        tokio::time::advance(Duration::from_millis(5)).await;
216        assert!(timeout.poll::<()>(&mut cx, |_| Poll::Pending).is_pending());
217        assert_eq!(timeout.deadline, Some(deadline));
218        tokio::time::advance(Duration::from_millis(6)).await;
219        assert!(matches!(
220            timeout.poll(&mut cx, |_| Poll::Ready(Ok(42))),
221            Poll::Ready(Ok(42))
222        ));
223        assert!(timeout.deadline.is_none());
224    }
225
226    #[crate::test(tokio::test(start_paused = true))]
227    async fn read_returns_data_that_arrived_before_a_late_poll() {
228        let waker = noop_waker();
229        let mut cx = Context::from_waker(&waker);
230        let mut reader = TimedReader::new(ScriptedIo { ready: false }, Duration::from_millis(10));
231        let mut bytes = [0; 8];
232        let mut buf = ReadBuf::new(&mut bytes);
233        assert!(
234            Pin::new(&mut reader)
235                .poll_read(&mut cx, &mut buf)
236                .is_pending()
237        );
238        // The peer sends data, but the task is not polled again until after
239        // the deadline.
240        reader.reader.ready = true;
241        tokio::time::advance(Duration::from_millis(20)).await;
242        assert!(matches!(
243            Pin::new(&mut reader).poll_read(&mut cx, &mut buf),
244            Poll::Ready(Ok(()))
245        ));
246        assert_eq!(buf.filled(), b"x");
247        assert!(reader.timeout.deadline.is_none());
248    }
249
250    #[crate::test(tokio::test(start_paused = true))]
251    async fn completion_and_error_clear_deadline_and_next_operation_reuses_timer() {
252        let waker = noop_waker();
253        let mut cx = Context::from_waker(&waker);
254        let mut timeout = IoTimeout::new(Duration::from_secs(1));
255        assert!(timeout.poll::<()>(&mut cx, |_| Poll::Pending).is_pending());
256        let first_deadline = timeout.deadline.unwrap();
257        let sleep: *const Sleep = &**timeout.sleep.as_ref().unwrap();
258        tokio::time::advance(Duration::from_millis(20)).await;
259        assert!(matches!(
260            timeout.poll(&mut cx, |_| Poll::Ready(Ok(42))),
261            Poll::Ready(Ok(42))
262        ));
263        assert!(timeout.deadline.is_none());
264        assert!(timeout.poll::<()>(&mut cx, |_| Poll::Pending).is_pending());
265        assert!(timeout.deadline.unwrap() > first_deadline);
266        let reused = timeout.sleep.as_ref().unwrap();
267        assert!(std::ptr::eq(&**reused, sleep), "timer must be reused");
268        assert_eq!(
269            reused.deadline(),
270            first_deadline,
271            "unfired timer must not be reset"
272        );
273        let error = timeout.poll::<()>(&mut cx, |_| {
274            Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
275        });
276        assert!(matches!(error, Poll::Ready(Err(err)) if err.kind() == io::ErrorKind::BrokenPipe));
277        assert!(timeout.deadline.is_none());
278        assert!(timeout.poll::<()>(&mut cx, |_| Poll::Pending).is_pending());
279        drop(timeout);
280        tokio::time::advance(Duration::from_secs(2)).await;
281    }
282
283    #[crate::test(tokio::test(start_paused = true))]
284    async fn deadline_starts_before_first_io_poll() {
285        let waker = noop_waker();
286        let mut cx = Context::from_waker(&waker);
287        let mut timeout = IoTimeout::new(Duration::from_secs(1));
288        let start = Instant::now();
289        assert!(timeout.poll::<()>(&mut cx, |_| Poll::Pending).is_pending());
290        assert_eq!(timeout.deadline, Some(start + Duration::from_secs(1)));
291        assert_eq!(
292            timeout.sleep.as_ref().unwrap().deadline(),
293            start + Duration::from_secs(1)
294        );
295    }
296
297    #[crate::test(tokio::test(start_paused = true))]
298    async fn zero_timeout_completes_ready_io_and_times_out_pending_io() {
299        let waker = noop_waker();
300        let mut cx = Context::from_waker(&waker);
301        let mut timeout = IoTimeout::new(Duration::ZERO);
302        assert!(matches!(
303            timeout.poll(&mut cx, |_| Poll::Ready(Ok(()))),
304            Poll::Ready(Ok(()))
305        ));
306        assert!(timeout.deadline.is_none());
307        let result = timeout.poll::<()>(&mut cx, |_| Poll::Pending);
308        assert!(matches!(result, Poll::Ready(Err(err)) if err.kind() == io::ErrorKind::TimedOut));
309        assert!(timeout.deadline.is_none());
310    }
311
312    #[crate::test(tokio::test(start_paused = true))]
313    async fn overflowing_duration_never_times_out() {
314        let waker = noop_waker();
315        let mut cx = Context::from_waker(&waker);
316        let mut timeout = IoTimeout::new(Duration::MAX);
317        assert!(timeout.poll::<()>(&mut cx, |_| Poll::Pending).is_pending());
318        assert!(matches!(
319            timeout.poll(&mut cx, |_| Poll::Ready(Ok(()))),
320            Poll::Ready(Ok(()))
321        ));
322        assert!(timeout.sleep.is_none());
323    }
324
325    struct ScriptedIo {
326        ready: bool,
327    }
328
329    impl AsyncRead for ScriptedIo {
330        fn poll_read(
331            self: Pin<&mut Self>,
332            _: &mut Context<'_>,
333            buf: &mut ReadBuf<'_>,
334        ) -> Poll<io::Result<()>> {
335            if self.ready {
336                buf.put_slice(b"x");
337                Poll::Ready(Ok(()))
338            } else {
339                Poll::Pending
340            }
341        }
342    }
343
344    impl AsyncWrite for ScriptedIo {
345        fn poll_write(
346            self: Pin<&mut Self>,
347            _: &mut Context<'_>,
348            buf: &[u8],
349        ) -> Poll<io::Result<usize>> {
350            if self.ready {
351                Poll::Ready(Ok(buf.len()))
352            } else {
353                Poll::Pending
354            }
355        }
356
357        fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
358            if self.ready {
359                Poll::Ready(Ok(()))
360            } else {
361                Poll::Pending
362            }
363        }
364
365        fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
366            self.poll_flush(cx)
367        }
368    }
369
370    #[crate::test(tokio::test(start_paused = true))]
371    async fn pending_read_completes_before_deadline_and_clears_deadline() {
372        let waker = noop_waker();
373        let mut cx = Context::from_waker(&waker);
374        let mut reader = TimedReader::new(ScriptedIo { ready: false }, Duration::from_millis(10));
375        let mut bytes = [0; 8];
376        let mut buf = ReadBuf::new(&mut bytes);
377        assert!(
378            Pin::new(&mut reader)
379                .poll_read(&mut cx, &mut buf)
380                .is_pending()
381        );
382        tokio::time::advance(Duration::from_millis(5)).await;
383        reader.reader.ready = true;
384        assert!(matches!(
385            Pin::new(&mut reader).poll_read(&mut cx, &mut buf),
386            Poll::Ready(Ok(()))
387        ));
388        assert_eq!(buf.filled(), b"x");
389        assert!(reader.timeout.deadline.is_none());
390    }
391
392    #[crate::test(tokio::test(start_paused = true))]
393    async fn timeout_wakes_pending_io_without_an_io_wakeup() {
394        use tokio::io::AsyncReadExt;
395
396        let mut reader = TimedReader::new(ScriptedIo { ready: false }, Duration::from_millis(10));
397        let start = Instant::now();
398        let mut bytes = [0; 8];
399        let error = reader.read(&mut bytes).await.unwrap_err();
400        assert_eq!(error.kind(), io::ErrorKind::TimedOut);
401        assert!(Instant::now() >= start + Duration::from_millis(10));
402        assert!(reader.timeout.deadline.is_none());
403    }
404
405    #[crate::test]
406    fn immediate_error_does_not_create_a_timer() {
407        let waker = noop_waker();
408        let mut cx = Context::from_waker(&waker);
409        let mut timeout = IoTimeout::new(Duration::from_secs(1));
410        let result = timeout.poll::<()>(&mut cx, |_| {
411            Poll::Ready(Err(io::ErrorKind::BrokenPipe.into()))
412        });
413        assert!(matches!(result, Poll::Ready(Err(err)) if err.kind() == io::ErrorKind::BrokenPipe));
414        assert!(timeout.sleep.is_none());
415    }
416
417    #[crate::test(tokio::test(start_paused = true))]
418    async fn writer_transitions_preserve_pending_deadline_and_reset_on_completion() {
419        let waker = noop_waker();
420        let mut cx = Context::from_waker(&waker);
421        let mut writer = TimedWriter::new(ScriptedIo { ready: false }, Duration::from_millis(10));
422        assert!(
423            Pin::new(&mut writer)
424                .poll_write(&mut cx, b"hello")
425                .is_pending()
426        );
427        let deadline = writer.timeout.deadline.unwrap();
428        tokio::time::advance(Duration::from_millis(5)).await;
429        assert!(Pin::new(&mut writer).poll_flush(&mut cx).is_pending());
430        assert_eq!(writer.timeout.deadline, Some(deadline));
431        writer.writer.ready = true;
432        assert!(matches!(
433            Pin::new(&mut writer).poll_flush(&mut cx),
434            Poll::Ready(Ok(()))
435        ));
436        assert!(writer.timeout.deadline.is_none());
437        writer.writer.ready = false;
438        assert!(Pin::new(&mut writer).poll_shutdown(&mut cx).is_pending());
439        assert!(writer.timeout.deadline.unwrap() > deadline);
440        tokio::time::advance(Duration::from_millis(11)).await;
441        let result = Pin::new(&mut writer).poll_shutdown(&mut cx);
442        assert!(matches!(result, Poll::Ready(Err(err)) if err.kind() == io::ErrorKind::TimedOut));
443        assert!(writer.timeout.deadline.is_none());
444        assert!(Pin::new(&mut writer).poll_write(&mut cx, b"x").is_pending());
445        tokio::time::advance(Duration::from_millis(11)).await;
446        writer.writer.ready = true;
447        assert!(matches!(
448            Pin::new(&mut writer).poll_write(&mut cx, b"x"),
449            Poll::Ready(Ok(1))
450        ));
451        assert!(writer.timeout.deadline.is_none());
452    }
453
454    #[crate::test(tokio::test(start_paused = true))]
455    async fn early_timer_fire_rearms_to_current_deadline() {
456        let waker = noop_waker();
457        let mut cx = Context::from_waker(&waker);
458        let mut timeout = IoTimeout::new(Duration::from_millis(10));
459        let start = Instant::now();
460        assert!(timeout.poll::<()>(&mut cx, |_| Poll::Pending).is_pending());
461        tokio::time::advance(Duration::from_millis(5)).await;
462        assert!(timeout.poll(&mut cx, |_| Poll::Ready(Ok(()))).is_ready());
463        assert!(timeout.poll::<()>(&mut cx, |_| Poll::Pending).is_pending());
464        let deadline = timeout.deadline.unwrap();
465        assert_eq!(deadline, start + Duration::from_millis(15));
466        tokio::time::advance(Duration::from_millis(6)).await;
467        assert!(timeout.poll::<()>(&mut cx, |_| Poll::Pending).is_pending());
468        assert_eq!(timeout.sleep.as_ref().unwrap().deadline(), deadline);
469        tokio::time::advance(Duration::from_millis(4)).await;
470        let result = timeout.poll::<()>(&mut cx, |_| Poll::Pending);
471        assert!(matches!(result, Poll::Ready(Err(err)) if err.kind() == io::ErrorKind::TimedOut));
472    }
473
474    #[crate::test(tokio::test(start_paused = true))]
475    async fn early_timer_fire_wakes_task_until_current_deadline() {
476        let mut timeout = IoTimeout::new(Duration::from_millis(10));
477        let start = Instant::now();
478        let waker = noop_waker();
479        let mut cx = Context::from_waker(&waker);
480        assert!(timeout.poll::<()>(&mut cx, |_| Poll::Pending).is_pending());
481        tokio::time::advance(Duration::from_millis(5)).await;
482        assert!(timeout.poll(&mut cx, |_| Poll::Ready(Ok(()))).is_ready());
483        let error = std::future::poll_fn(|cx| timeout.poll::<()>(cx, |_| Poll::Pending))
484            .await
485            .unwrap_err();
486        assert_eq!(error.kind(), io::ErrorKind::TimedOut);
487        assert!(Instant::now() >= start + Duration::from_millis(15));
488    }
489}
mz_ore/netio/timeout.rs

mz_ore/netio/
timeout.rs
