time: advance frozen time in `park_timeout` (#2059)

This patch improves the behavior of frozen time (a testing utility made
available with the `test-util` feature flag). Instead of of requiring
`time::advance` to be called in order to advance the value returned by
`Instant::now`, calls to `time::Driver::park_timeout` will use the
provided duration to advance the time.

This is the desired behavior as the timeout is used to indicate when the
next scheduled delay needs to be fired.

3 files changed, 23 insertions, 7 deletions

diff --git a/tokio/src/time/driver/entry.rs b/tokio/src/time/driver/entry.rs
index d5fab897..079ec7e8 100644
--- a/tokio/src/time/driver/entry.rs
+++ b/tokio/src/time/driver/entry.rs
@@ -104,8 +104,8 @@ const ERROR: u64 = u64::MAX;
 // ===== impl Entry =====
 
 impl Entry {
-    pub(crate) fn new(deadline: Instant, duration: Duration) -> Arc<Entry> {
-        let inner = Handle::current().inner().unwrap();
+    pub(crate) fn new(handle: &Handle, deadline: Instant, duration: Duration) -> Arc<Entry> {
+        let inner = handle.inner().unwrap();
         let entry: Entry;
 
         // Increment the number of active timeouts
diff --git a/tokio/src/time/driver/mod.rs b/tokio/src/time/driver/mod.rs
index 605d9bca..f74a4853 100644
--- a/tokio/src/time/driver/mod.rs
+++ b/tokio/src/time/driver/mod.rs
@@ -4,7 +4,7 @@ mod atomic_stack;
 use self::atomic_stack::AtomicStack;
 
 mod entry;
-use self::entry::Entry;
+pub(super) use self::entry::Entry;
 
 mod handle;
 pub(crate) use self::handle::Handle;
@@ -245,7 +245,14 @@ where
                 let deadline = self.expiration_instant(when);
 
                 if deadline > now {
-                    self.park.park_timeout(deadline - now)?;
+                    let dur = deadline - now;
+
+                    if self.clock.is_frozen() {
+                        self.park.park_timeout(Duration::from_secs(0))?;
+                        self.clock.advance(dur);
+                    } else {
+                        self.park.park_timeout(dur)?;
+                    }
                 } else {
                     self.park.park_timeout(Duration::from_secs(0))?;
                 }
@@ -269,7 +276,14 @@ where
                 let deadline = self.expiration_instant(when);
 
                 if deadline > now {
-                    self.park.park_timeout(cmp::min(deadline - now, duration))?;
+                    let duration = cmp::min(deadline - now, duration);
+
+                    if self.clock.is_frozen() {
+                        self.park.park_timeout(Duration::from_secs(0))?;
+                        self.clock.advance(duration);
+                    } else {
+                        self.park.park_timeout(duration)?;
+                    }
                 } else {
                     self.park.park_timeout(Duration::from_secs(0))?;
                 }
diff --git a/tokio/src/time/driver/registration.rs b/tokio/src/time/driver/registration.rs
index 728d2993..141ca015 100644
--- a/tokio/src/time/driver/registration.rs
+++ b/tokio/src/time/driver/registration.rs
@@ -1,4 +1,4 @@
-use crate::time::driver::Entry;
+use crate::time::driver::{Entry, Handle};
 use crate::time::{Duration, Error, Instant};
 
 use std::sync::Arc;
@@ -15,8 +15,10 @@ pub(crate) struct Registration {
 
 impl Registration {
     pub(crate) fn new(deadline: Instant, duration: Duration) -> Registration {
+        let handle = Handle::current();
+
         Registration {
-            entry: Entry::new(deadline, duration),
+            entry: Entry::new(&handle, deadline, duration),
         }
     }