revert some changes

tokio/src/runtime/time/mod.rs

@@ -190,18 +190,22 @@ impl Driver {
         assert!(!handle.is_shutdown());
 
         // Finds out the min expiration time to park.
-        let expiration_time = (0..rt_handle.time().inner.get_shard_size())
-            .filter_map(|id| {
-                let lock = rt_handle.time().inner.lock_sharded_wheel(id);
-                lock.next_expiration_time()
-            })
+        let locks = (0..rt_handle.time().inner.get_shard_size())
+            .map(|id| rt_handle.time().inner.lock_sharded_wheel(id))
+            .collect::<Vec<_>>();
+
+        let expiration_time = locks
+            .iter()
+            .filter_map(|lock| lock.next_expiration_time())
             .min();
 
         rt_handle
             .time()
             .inner
             .next_wake
             .store(next_wake_time(expiration_time));
+        // Safety: After updating the `next_wake`, we drop all the locks.
+        drop(locks);
 
         match expiration_time {
             Some(when) => {

tokio/src/runtime/time/tests/mod.rs

@@ -219,7 +219,6 @@ fn reset_timer_and_drop() {
                 .poll_elapsed(&mut Context::from_waker(futures::task::noop_waker_ref()));
 
             entry.as_mut().reset(start + Duration::from_secs(1), true);
-            drop(entry);
         }
     });
 }

tokio/src/runtime/time/wheel/mod.rs

@@ -190,7 +190,7 @@ impl Wheel {
     }
 
     /// Inserts the `entry` to the `Wheel`.
-    /// Returns the level and slot which `entry` insert into.
+    /// Returns the level and the slot which `entry` insert into.
     ///
     /// Safety: The `cached_when` of this `entry`` must have been updated to `when`.
     unsafe fn inner_insert(