From 98cd6d484671e213285fbcf76a78c0aabfb84469 Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Thu, 10 Jun 2021 14:58:40 -0700
Subject: [PATCH] time: fix time::advance() with sub-ms durations

Update the advance logic to factor in the timer's ms rounding.

Fixes #3837
---
 tokio/src/time/clock.rs   | 12 +++++++++++-
 tokio/tests/time_pause.rs | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/tokio/src/time/clock.rs b/tokio/src/time/clock.rs
index c5ef86be4e0..865f8e30f19 100644
--- a/tokio/src/time/clock.rs
+++ b/tokio/src/time/clock.rs
@@ -124,7 +124,17 @@ cfg_test_util! {
         let until = clock.now() + duration;
         clock.advance(duration);
 
-        crate::time::sleep_until(until).await;
+        // Fixes tokio-rs/tokio#3837
+        //
+        // "sleep_until" is needed to ensure registered timers are fired after the
+        // advance call, however, it will also round up time to the nearest ms. To
+        // avoid this loss of precision, we sleep until the previous millisecond.
+        // If the advance is less than 1ms, then we just need to yield.
+        if duration >= Duration::from_millis(1) {
+            crate::time::sleep_until(until - Duration::from_millis(1)).await;
+        } else {
+            crate::task::yield_now().await;
+        }
     }
 
     /// Return the current instant, factoring in frozen time.
diff --git a/tokio/tests/time_pause.rs b/tokio/tests/time_pause.rs
index d1834af2157..87310b635f0 100644
--- a/tokio/tests/time_pause.rs
+++ b/tokio/tests/time_pause.rs
@@ -215,6 +215,40 @@ async fn interval() {
     assert_pending!(poll_next(&mut i));
 }
 
+#[tokio::test]
+async fn test_time_advance_sub_ms() {
+    time::pause();
+    let now = Instant::now();
+
+    let dur = Duration::from_micros(51_592);
+    time::advance(dur).await;
+
+    assert_eq!(now.elapsed(), dur);
+
+    let now = Instant::now();
+    let dur = Duration::from_micros(1);
+    time::advance(dur).await;
+
+    assert_eq!(now.elapsed(), dur);
+}
+
+#[tokio::test]
+async fn test_time_advance_3ms_and_change() {
+    time::pause();
+    let now = Instant::now();
+
+    let dur = Duration::from_micros(3_141_592);
+    time::advance(dur).await;
+
+    assert_eq!(now.elapsed(), dur);
+
+    let now = Instant::now();
+    let dur = Duration::from_micros(3_123_456);
+    time::advance(dur).await;
+
+    assert_eq!(now.elapsed(), dur);
+}
+
 fn poll_next(interval: &mut task::Spawn<time::Interval>) -> Poll<Instant> {
     interval.enter(|cx, mut interval| interval.poll_tick(cx))
 }