From 483e4b9ee7a677014233f6b24fb26bac92bf6899 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Mon, 7 Sep 2026 18:17:16 +0200 Subject: [PATCH] stream: handle overflowing timer durations (#8354) --- tokio-stream/src/stream_ext/throttle.rs | 6 +++--- tokio-stream/src/stream_ext/timeout.rs | 10 ++++------ tokio-stream/tests/stream_timeout.rs | 8 ++++++++ tokio-stream/tests/time_throttle.rs | 7 +++++++ 4 files changed, 22 insertions(+), 9 deletions(-) diff --git a/tokio-stream/src/stream_ext/throttle.rs b/tokio-stream/src/stream_ext/throttle.rs index 4e71debf3..270514571 100644 --- a/tokio-stream/src/stream_ext/throttle.rs +++ b/tokio-stream/src/stream_ext/throttle.rs @@ -1,7 +1,7 @@ //! Slow down a stream by enforcing a delay between items. use crate::Stream; -use tokio::time::{Duration, Instant, Sleep}; +use tokio::time::{sleep, Duration, Sleep}; use std::future::Future; use std::pin::Pin; @@ -14,7 +14,7 @@ where T: Stream, { Throttle { - delay: tokio::time::sleep_until(Instant::now() + duration), + delay: sleep(duration), duration, has_delayed: true, stream, @@ -81,7 +81,7 @@ impl Stream for Throttle { if value.is_some() { if !is_zero(dur) { - me.delay.reset(Instant::now() + dur); + me.delay.set(sleep(dur)); } *me.has_delayed = false; diff --git a/tokio-stream/src/stream_ext/timeout.rs b/tokio-stream/src/stream_ext/timeout.rs index d863af1db..d968ff0c1 100644 --- a/tokio-stream/src/stream_ext/timeout.rs +++ b/tokio-stream/src/stream_ext/timeout.rs @@ -1,6 +1,6 @@ use crate::stream_ext::Fuse; use crate::Stream; -use tokio::time::{Instant, Sleep}; +use tokio::time::{sleep, Sleep}; use core::future::Future; use core::pin::Pin; @@ -29,8 +29,7 @@ pub struct Elapsed(()); impl Timeout { pub(super) fn new(stream: S, duration: Duration) -> Self { - let next = Instant::now() + duration; - let deadline = tokio::time::sleep_until(next); + let deadline = sleep(duration); Timeout { stream: Fuse::new(stream), @@ -45,13 +44,12 @@ impl Stream for Timeout { type Item = Result; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let me = self.project(); + let mut me = self.project(); match me.stream.poll_next(cx) { Poll::Ready(v) => { if v.is_some() { - let next = Instant::now() + *me.duration; - me.deadline.reset(next); + me.deadline.set(sleep(*me.duration)); *me.poll_deadline = true; } return Poll::Ready(v.map(Ok)); diff --git a/tokio-stream/tests/stream_timeout.rs b/tokio-stream/tests/stream_timeout.rs index 19b4c3d53..5e824aa5b 100644 --- a/tokio-stream/tests/stream_timeout.rs +++ b/tokio-stream/tests/stream_timeout.rs @@ -107,3 +107,11 @@ async fn no_timeouts() { assert_ready_eq!(stream.poll_next(), Some(Ok(5))); assert_ready_eq!(stream.poll_next(), None); } + +#[tokio::test] +async fn duration_max_does_not_overflow() { + let stream = stream::iter([1]).timeout(Duration::MAX); + let mut stream = task::spawn(stream); + + assert_ready_eq!(stream.poll_next(), Some(Ok(1))); +} diff --git a/tokio-stream/tests/time_throttle.rs b/tokio-stream/tests/time_throttle.rs index e6c9917be..e118237bd 100644 --- a/tokio-stream/tests/time_throttle.rs +++ b/tokio-stream/tests/time_throttle.rs @@ -26,3 +26,10 @@ async fn usage() { assert_ready!(stream.poll_next()); } + +#[tokio::test] +async fn duration_max_does_not_overflow() { + let mut stream = task::spawn(futures::stream::iter([1]).throttle(Duration::MAX)); + + assert_ready_eq!(stream.poll_next(), Some(1)); +}