stream: handle overflowing timer durations (#8354)

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<T: Stream> Stream for Throttle<T> {
 
         if value.is_some() {
             if !is_zero(dur) {
-                me.delay.reset(Instant::now() + dur);
+                me.delay.set(sleep(dur));
             }
 
             *me.has_delayed = false;

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<S: Stream> Timeout<S> {
     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<S: Stream> Stream for Timeout<S> {
     type Item = Result<S::Item, Elapsed>;
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        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));

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)));
+}

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));
+}