From 35f94c1f2eb2127501628fa16b88a1b5866e957b Mon Sep 17 00:00:00 2001
From: wathenjiang <wjjiang19@gmail.com>
Date: Tue, 16 Jul 2024 19:42:07 +0800
Subject: [PATCH] time: update wake_up while holding all the locks of sharded
 time wheels

---
 tokio/src/runtime/time/mod.rs | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/tokio/src/runtime/time/mod.rs b/tokio/src/runtime/time/mod.rs
index 37b04ef0000..f027f9399ce 100644
--- a/tokio/src/runtime/time/mod.rs
+++ b/tokio/src/runtime/time/mod.rs
@@ -190,11 +190,13 @@ 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
@@ -202,7 +204,9 @@ impl Driver {
             .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) => {
                 let now = handle.time_source.now(rt_handle.clock());