rt: remove Handle reference from time driver (#4993)

A step towards the goal of consolidating thread-local variables, this
patch removes the `time::Handle` reference from the time driver struct.
The reference is moved a level up and passed into methods that need it.

The end goal is to have a single `Driver` struct that holds each
sub-driver and a single `Handle` ref-counted struct that holds all the
misc handles. Further work on the time driver is blocked by some other
refactoring that will happen in follow-up patches.

tokio/src/runtime/driver.rs

@@ -147,7 +147,10 @@ cfg_not_process_driver! {
 cfg_time! {
     #[derive(Debug)]
     pub(crate) enum TimeDriver {
-        Enabled(crate::runtime::time::Driver),
+        Enabled {
+            driver: crate::runtime::time::Driver,
+            handle: crate::runtime::time::Handle,
+        },
         Disabled(IoStack),
     }
 
@@ -169,10 +172,9 @@ cfg_time! {
         clock: Clock,
     ) -> (TimeDriver, TimeHandle) {
         if enable {
-            let driver = crate::runtime::time::Driver::new(io_stack, clock);
-            let handle = driver.handle();
+            let (driver, handle) = crate::runtime::time::Driver::new(io_stack, clock);
 
-            (TimeDriver::Enabled(driver), Some(handle))
+            (TimeDriver::Enabled { driver, handle: handle.clone() }, Some(handle))
         } else {
             (TimeDriver::Disabled(io_stack), None)
         }
@@ -181,21 +183,21 @@ cfg_time! {
     impl TimeDriver {
         pub(crate) fn unpark(&self) -> TimerUnpark {
             match self {
-                TimeDriver::Enabled(v) => TimerUnpark::Enabled(v.unpark()),
+                TimeDriver::Enabled { driver, .. } => TimerUnpark::Enabled(driver.unpark()),
                 TimeDriver::Disabled(v) => TimerUnpark::Disabled(v.unpark()),
             }
         }
 
         pub(crate) fn park(&mut self) {
             match self {
-                TimeDriver::Enabled(v) => v.park(),
+                TimeDriver::Enabled { driver, handle } => driver.park(handle),
                 TimeDriver::Disabled(v) => v.park(),
             }
         }
 
         pub(crate) fn park_timeout(&mut self, duration: Duration) {
             match self {
-                TimeDriver::Enabled(v) => v.park_timeout(duration),
+                TimeDriver::Enabled { driver, handle } => driver.park_timeout(handle, duration),
                 TimeDriver::Disabled(v) => v.park_timeout(duration),
             }
         }
@@ -204,13 +206,21 @@ cfg_time! {
         cfg_rt_multi_thread! {
             pub(crate) fn shutdown(&mut self) {
                 match self {
-                    TimeDriver::Enabled(v) => v.shutdown(),
+                    TimeDriver::Enabled { driver, handle } => driver.shutdown(handle),
                     TimeDriver::Disabled(v) => v.shutdown(),
                 }
             }
         }
     }
 
+    impl Drop for TimeDriver {
+        fn drop(&mut self) {
+            if let TimeDriver::Enabled { driver, handle } = self {
+                driver.shutdown(handle);
+            }
+        }
+    }
+
     impl TimerUnpark {
         pub(crate) fn unpark(&self) {
             match self {

tokio/src/runtime/time/mod.rs

@@ -99,9 +99,6 @@ pub(crate) struct Driver {
     /// Timing backend in use.
     time_source: TimeSource,
 
-    /// Shared state.
-    handle: Handle,
-
     /// Parker to delegate to.
     park: IoStack,
 
@@ -149,60 +146,52 @@ impl Driver {
     /// thread and `time_source` to get the current time and convert to ticks.
     ///
     /// Specifying the source of time is useful when testing.
-    pub(crate) fn new(park: IoStack, clock: Clock) -> Driver {
+    pub(crate) fn new(park: IoStack, clock: Clock) -> (Driver, Handle) {
         let time_source = TimeSource::new(clock);
 
         let inner = Inner::new(time_source.clone(), park.unpark());
+        let handle = Handle::new(Arc::new(inner));
 
-        Driver {
+        let driver = Driver {
             time_source,
-            handle: Handle::new(Arc::new(inner)),
             park,
             #[cfg(feature = "test-util")]
             did_wake: Arc::new(AtomicBool::new(false)),
-        }
-    }
+        };
 
-    /// Returns a handle to the timer.
-    ///
-    /// The `Handle` is how `Sleep` instances are created. The `Sleep` instances
-    /// can either be created directly or the `Handle` instance can be passed to
-    /// `with_default`, setting the timer as the default timer for the execution
-    /// context.
-    pub(crate) fn handle(&self) -> Handle {
-        self.handle.clone()
+        (driver, handle)
     }
 
     pub(crate) fn unpark(&self) -> TimerUnpark {
         TimerUnpark::new(self)
     }
 
-    pub(crate) fn park(&mut self) {
-        self.park_internal(None)
+    pub(crate) fn park(&mut self, handle: &Handle) {
+        self.park_internal(handle, None)
     }
 
-    pub(crate) fn park_timeout(&mut self, duration: Duration) {
-        self.park_internal(Some(duration))
+    pub(crate) fn park_timeout(&mut self, handle: &Handle, duration: Duration) {
+        self.park_internal(handle, Some(duration))
     }
 
-    pub(crate) fn shutdown(&mut self) {
-        if self.handle.is_shutdown() {
+    pub(crate) fn shutdown(&mut self, handle: &Handle) {
+        if handle.is_shutdown() {
             return;
         }
 
-        self.handle.get().is_shutdown.store(true, Ordering::SeqCst);
+        handle.get().is_shutdown.store(true, Ordering::SeqCst);
 
         // Advance time forward to the end of time.
 
-        self.handle.process_at_time(u64::MAX);
+        handle.process_at_time(u64::MAX);
 
         self.park.shutdown();
     }
 
-    fn park_internal(&mut self, limit: Option<Duration>) {
-        let mut lock = self.handle.get().state.lock();
+    fn park_internal(&mut self, handle: &Handle, limit: Option<Duration>) {
+        let mut lock = handle.get().state.lock();
 
-        assert!(!self.handle.is_shutdown());
+        assert!(!handle.is_shutdown());
 
         let next_wake = lock.wheel.next_expiration_time();
         lock.next_wake =
@@ -238,7 +227,7 @@ impl Driver {
         }
 
         // Process pending timers after waking up
-        self.handle.process();
+        handle.process();
     }
 
     cfg_test_util! {
@@ -414,12 +403,6 @@ impl Handle {
     }
 }
 
-impl Drop for Driver {
-    fn drop(&mut self) {
-        self.shutdown();
-    }
-}
-
 pub(crate) struct TimerUnpark {
     inner: IoUnpark,