implement extra reset variants for Interval

Fixes: #5874

Author: victor-timofei

tokio/src/time/interval.rs

@@ -495,6 +495,8 @@ impl Interval {
     ///
     /// This method ignores [`MissedTickBehavior`] strategy.
     ///
+    /// This is equivalent to calling `reset_at(Instant::now() + period)`.
+    ///
     /// # Examples
     ///
     /// ```
@@ -521,6 +523,107 @@ impl Interval {
         self.delay.as_mut().reset(Instant::now() + self.period);
     }
 
+    /// Resets the interval immediately.
+    ///
+    /// This method ignores [`MissedTickBehavior`] strategy.
+    ///
+    /// This is equivalent to calling `reset_at(Instant::now())`.
+    ///
+    /// # 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_immediately();
+    ///
+    ///     interval.tick().await;
+    ///     interval.tick().await;
+    ///
+    ///     // approximately 150ms have elapsed.
+    /// }
+    /// ```
+    pub fn reset_immediately(&mut self) {
+        self.delay.as_mut().reset(Instant::now());
+    }
+
+    /// Resets the interval after the specified [`std::time::Duration`].
+    ///
+    /// This method ignores [`MissedTickBehavior`] strategy.
+    ///
+    /// This is equivalent to calling `reset_at(Instant::now() + after)`.
+    ///
+    /// # 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;
+    ///
+    ///     let after = Duration::from_millis(20);
+    ///     interval.reset_after(after);
+    ///
+    ///     interval.tick().await;
+    ///     interval.tick().await;
+    ///
+    ///     // approximately 170ms have elapsed.
+    /// }
+    /// ```
+    pub fn reset_after(&mut self, after: Duration) {
+        self.delay.as_mut().reset(Instant::now() + after);
+    }
+
+    /// Resets the interval to a [`crate::time::Instant`] deadline.
+    ///
+    /// Sets the next tick to expire at the given instant. If the instant is in
+    /// the past, then the [`MissedTickBehavior`] strategy will be used to
+    /// catch up. If the instant is in the future, then the next tick will
+    /// complete at the given instant, even if that means that it will sleep for
+    /// longer than the duration of this [`Interval`]. If the [`Interval`] had
+    /// any missed ticks before calling this method, then those are discarded.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use tokio::time::{self, Instant};
+    ///
+    /// 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;
+    ///
+    ///     let deadline = Instant::now() + Duration::from_millis(30);
+    ///     interval.reset_at(deadline);
+    ///
+    ///     interval.tick().await;
+    ///     interval.tick().await;
+    ///
+    ///     // approximately 180ms have elapsed.
+    /// }
+    /// ```
+    pub fn reset_at(&mut self, deadline: Instant) {
+        self.delay.as_mut().reset(deadline);
+    }
+
     /// Returns the [`MissedTickBehavior`] strategy currently being used.
     pub fn missed_tick_behavior(&self) -> MissedTickBehavior {
         self.missed_tick_behavior

tokio/tests/time_interval.rs

@@ -204,6 +204,113 @@ async fn reset() {
     check_interval_poll!(i, start, 1001);
 }
 
+#[tokio::test(start_paused = true)]
+async fn reset_immediatelly() {
+    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_immediately();
+
+    // 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, 401);
+
+    time::advance(ms(100)).await;
+    check_interval_poll!(i, start);
+
+    time::advance(ms(200)).await;
+    check_interval_poll!(i, start, 701);
+}
+
+#[tokio::test(start_paused = true)]
+async fn reset_after() {
+    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_after(Duration::from_millis(20));
+
+    // We add one because when using `reset` method, `Interval` adds the
+    // `period` from `Instant::now()`, which will always be off by one
+    time::advance(ms(20)).await;
+    check_interval_poll!(i, start, 421);
+
+    time::advance(ms(100)).await;
+    check_interval_poll!(i, start);
+
+    time::advance(ms(200)).await;
+    check_interval_poll!(i, start, 721);
+}
+
+#[tokio::test(start_paused = true)]
+async fn reset_at() {
+    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_at(Instant::now() + Duration::from_millis(40));
+
+    // We add one because when using `reset` method, `Interval` adds the
+    // `period` from `Instant::now()`, which will always be off by one
+    time::advance(ms(40)).await;
+    check_interval_poll!(i, start, 441);
+
+    time::advance(ms(100)).await;
+    check_interval_poll!(i, start);
+
+    time::advance(ms(200)).await;
+    check_interval_poll!(i, start, 741);
+}
+
 fn poll_next(interval: &mut task::Spawn<time::Interval>) -> Poll<Instant> {
     interval.enter(|cx, mut interval| interval.poll_tick(cx))
 }