From 347c0cdaba27eea8e54b034db4195fa71652c5a1 Mon Sep 17 00:00:00 2001
From: oblique <psyberbits@gmail.com>
Date: Tue, 23 Nov 2021 14:51:07 +0200
Subject: [PATCH] time: add `Interval::reset` method (#4248)

---
 tokio/src/time/interval.rs   | 30 ++++++++++++++++++++++++++++++
 tokio/tests/time_interval.rs | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/tokio/src/time/interval.rs b/tokio/src/time/interval.rs
index 7e07e5126..6a8994374 100644
--- a/tokio/src/time/interval.rs
+++ b/tokio/src/time/interval.rs
@@ -435,6 +435,36 @@ impl Interval {
         Poll::Ready(timeout)
     }
 
+    /// Resets the interval to complete one period after the current time.
+    ///
+    /// This method ignores [`MissedTickBehavior`] strategy.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio::time;
+    ///
+    /// use std::time::Duration;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let mut interval = time::interval(Duration::from_millis(100));
+    ///
+    ///     interval.tick().await;
+    ///
+    ///     time::sleep(Duration::from_millis(50)).await;
+    ///     interval.reset();
+    ///
+    ///     interval.tick().await;
+    ///     interval.tick().await;
+    ///
+    ///     // approximately 250ms have elapsed.
+    /// }
+    /// ```
+    pub fn reset(&mut self) {
+        self.delay.as_mut().reset(Instant::now() + self.period);
+    }
+
     /// Returns the [`MissedTickBehavior`] strategy currently being used.
     pub fn missed_tick_behavior(&self) -> MissedTickBehavior {
         self.missed_tick_behavior
diff --git a/tokio/tests/time_interval.rs b/tokio/tests/time_interval.rs
index 5f7bf55f2..186582e2e 100644
--- a/tokio/tests/time_interval.rs
+++ b/tokio/tests/time_interval.rs
@@ -166,6 +166,42 @@ async fn skip() {
     check_interval_poll!(i, start, 1800);
 }
 
+#[tokio::test(start_paused = true)]
+async fn reset() {
+    let start = Instant::now();
+
+    // This is necessary because the timer is only so granular, and in order for
+    // all our ticks to resolve, the time needs to be 1ms ahead of what we
+    // expect, so that the runtime will see that it is time to resolve the timer
+    time::advance(ms(1)).await;
+
+    let mut i = task::spawn(time::interval_at(start, ms(300)));
+
+    check_interval_poll!(i, start, 0);
+
+    time::advance(ms(100)).await;
+    check_interval_poll!(i, start);
+
+    time::advance(ms(200)).await;
+    check_interval_poll!(i, start, 300);
+
+    time::advance(ms(100)).await;
+    check_interval_poll!(i, start);
+
+    i.reset();
+
+    time::advance(ms(250)).await;
+    check_interval_poll!(i, start);
+
+    time::advance(ms(50)).await;
+    // We add one because when using `reset` method, `Interval` adds the
+    // `period` from `Instant::now()`, which will always be off by one
+    check_interval_poll!(i, start, 701);
+
+    time::advance(ms(300)).await;
+    check_interval_poll!(i, start, 1001);
+}
+
 fn poll_next(interval: &mut task::Spawn<time::Interval>) -> Poll<Instant> {
     interval.enter(|cx, mut interval| interval.poll_tick(cx))
 }