1 files changed, 27 insertions, 10 deletions

diff --git a/tokio/src/runtime/time.rs b/tokio/src/runtime/time.rs
index b72971f8..1cd58cd6 100644
--- a/tokio/src/runtime/time.rs
+++ b/tokio/src/runtime/time.rs
@@ -3,36 +3,49 @@
 //! shells. This isolates the complexity of dealing with conditional
 //! compilation.
 
-cfg_time! {
+pub(crate) use variant::*;
+
+#[cfg(all(feature = "time", not(loom)))]
+mod variant {
+    use crate::park::Either;
     use crate::runtime::io;
     use crate::time::{self, driver};
 
     pub(crate) type Clock = time::Clock;
-    pub(crate) type Driver = driver::Driver<io::Driver>;
-    pub(crate) type Handle = driver::Handle;
+    pub(crate) type Driver = Either<driver::Driver<io::Driver>, io::Driver>;
+    pub(crate) type Handle = Option<driver::Handle>;
 
     pub(crate) fn create_clock() -> Clock {
         Clock::new()
     }
 
     /// Create a new timer driver / handle pair
-    pub(crate) fn create_driver(io_driver: io::Driver, clock: Clock) -> (Driver, Handle) {
-        let driver = driver::Driver::new(io_driver, clock);
-        let handle = driver.handle();
+    pub(crate) fn create_driver(
+        enable: bool,
+        io_driver: io::Driver,
+        clock: Clock,
+    ) -> (Driver, Handle) {
+        if enable {
+            let driver = driver::Driver::new(io_driver, clock);
+            let handle = driver.handle();
 
-        (driver, handle)
+            (Either::A(driver), Some(handle))
+        } else {
+            (Either::B(io_driver), None)
+        }
     }
 
     pub(crate) fn with_default<F, R>(handle: &Handle, clock: &Clock, f: F) -> R
     where
         F: FnOnce() -> R,
     {
-        let _time = driver::set_default(handle);
+        let _time = handle.as_ref().map(|handle| driver::set_default(handle));
         clock.enter(f)
     }
 }
 
-cfg_not_time! {
+#[cfg(any(not(feature = "time"), loom))]
+mod variant {
     use crate::runtime::io;
 
     pub(crate) type Clock = ();
@@ -44,7 +57,11 @@ cfg_not_time! {
     }
 
     /// Create a new timer driver / handle pair
-    pub(crate) fn create_driver(io_driver: io::Driver, _clock: Clock) -> (Driver, Handle) {
+    pub(crate) fn create_driver(
+        _enable: bool,
+        io_driver: io::Driver,
+        _clock: Clock,
+    ) -> (Driver, Handle) {
         (io_driver, ())
     }