introduce random make shutdown faster

tokio/src/runtime/scheduler/current_thread/mod.rs

@@ -248,7 +248,7 @@ fn shutdown2(mut core: Box<Core>, handle: &Handle) -> Box<Core> {
     // Drain the OwnedTasks collection. This call also closes the
     // collection, ensuring that no tasks are ever pushed after this
     // call returns.
-    handle.shared.owned.close_and_shutdown_all();
+    handle.shared.owned.close_and_shutdown_all(0);
 
     // Drain local queue
     // We already shut down every task, so we just need to drop the task.
@@ -615,7 +615,7 @@ impl Schedule for Arc<Handle> {
                             // If `None`, the runtime is shutting down, so there is no need to signal shutdown
                             if let Some(core) = core.as_mut() {
                                 core.unhandled_panic = true;
-                                self.shared.owned.close_and_shutdown_all();
+                                self.shared.owned.close_and_shutdown_all(0);
                             }
                         }
                         _ => unreachable!("runtime core not set in CURRENT thread-local"),

tokio/src/runtime/scheduler/multi_thread/worker.rs

@@ -547,7 +547,6 @@ impl Context {
         }
 
         core.pre_shutdown(&self.worker);
-
         // Signal shutdown
         self.worker.handle.shutdown_core(core);
         Err(())
@@ -954,8 +953,16 @@ impl Core {
     /// Signals all tasks to shut down, and waits for them to complete. Must run
     /// before we enter the single-threaded phase of shutdown processing.
     fn pre_shutdown(&mut self, worker: &Worker) {
+        // Start from a random inner list
+        let start = self
+            .rand
+            .fastrand_n(worker.handle.shared.owned.grain as u32);
         // Signal to all tasks to shut down.
-        worker.handle.shared.owned.close_and_shutdown_all();
+        worker
+            .handle
+            .shared
+            .owned
+            .close_and_shutdown_all(start as usize);
 
         self.stats
             .submit(&worker.handle.shared.worker_metrics[worker.index]);

tokio/src/runtime/scheduler/multi_thread_alt/worker.rs

@@ -1460,7 +1460,9 @@ impl Shared {
     }
 
     pub(super) fn shutdown_core(&self, handle: &Handle, mut core: Box<Core>) {
-        self.owned.close_and_shutdown_all();
+        // Start from a random inner list
+        let start = core.rand.fastrand_n(self.owned.grain as u32);
+        self.owned.close_and_shutdown_all(start as usize);
 
         core.stats.submit(&self.worker_metrics[core.index]);
 

tokio/src/runtime/task/list.rs

@@ -38,10 +38,10 @@ fn get_next_id() -> NonZeroU32 {
 }
 
 pub(crate) struct OwnedTasks<S: 'static> {
-    lists: Vec<Mutex<ListInner<S>>>,
+    lists: Box<[Mutex<ListInner<S>>]>,
     pub(crate) id: NonZeroU32,
     closed: AtomicBool,
-    grain: usize,
+    pub(crate) grain: u32,
     count: AtomicUsize,
 }
 
@@ -59,18 +59,18 @@ struct OwnedTasksInner<S: 'static> {
 
 impl<S: 'static> OwnedTasks<S> {
     /// grain must be an integer power of 2
-    pub(crate) fn new(grain: usize) -> Self {
+    pub(crate) fn new(grain: u32) -> Self {
         assert_eq!(
             grain & (grain - 1),
             0,
             "the grain of OwnedTasks must be an integer power of 2"
         );
-        let mut lists = Vec::with_capacity(grain);
+        let mut lists = Vec::with_capacity(grain as usize);
         for _ in 0..grain {
             lists.push(Mutex::new(LinkedList::new()))
         }
         Self {
-            lists,
+            lists: lists.into_boxed_slice(),
             closed: AtomicBool::new(false),
             id: get_next_id(),
             grain,
@@ -124,7 +124,7 @@ impl<S: 'static> OwnedTasks<S> {
     where
         S: Schedule,
     {
-        let mut lock = self.lists[task_id.0 as usize & (self.grain - 1)].lock();
+        let mut lock = self.lists[task_id.0 as usize & (self.grain - 1) as usize].lock();
         // check close flag
         if self.closed.load(Ordering::Acquire) {
             task.shutdown();
@@ -150,17 +150,16 @@ impl<S: 'static> OwnedTasks<S> {
 
     /// Shuts down all tasks in the collection. This call also closes the
     /// collection, preventing new items from being added.
-    pub(crate) fn close_and_shutdown_all(&self)
+    pub(crate) fn close_and_shutdown_all(&self, start: usize)
     where
         S: Schedule,
     {
         // The first iteration of the loop was unrolled so it can set the
         // closed bool.
         self.closed.store(true, Ordering::Release);
-
-        for list in &self.lists {
+        for i in start..self.lists.len() + start {
             loop {
-                let task = match list.lock().pop_back() {
+                let task = match self.lists[i % (self.lists.len())].lock().pop_back() {
                     Some(task) => {
                         self.count.fetch_sub(1, Ordering::Relaxed);
                         task
@@ -190,7 +189,7 @@ impl<S: 'static> OwnedTasks<S> {
 
     #[inline]
     unsafe fn remove_inner(&self, task: &Task<S>) -> Option<Task<S>> {
-        match self.lists[(task.header().task_id.0) as usize & (self.grain - 1)]
+        match self.lists[(task.header().task_id.0) as usize & (self.grain - 1) as usize]
             .lock()
             .remove(task.header_ptr())
         {

tokio/src/runtime/tests/task.rs

@@ -308,7 +308,7 @@ impl Runtime {
     fn shutdown(&self) {
         let mut core = self.0.core.try_lock().unwrap();
 
-        self.0.owned.close_and_shutdown_all();
+        self.0.owned.close_and_shutdown_all(0);
 
         while let Some(task) = core.queue.pop_back() {
             drop(task);