From 904dabb23d0b3c5e6f1f53b628ee426d7f9e362d Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Mon, 12 Jun 2023 12:51:23 -0700
Subject: [PATCH] fix more bugs

---
 .../runtime/scheduler/multi_thread/idle.rs    |  5 ++++
 .../runtime/scheduler/multi_thread/worker.rs  | 24 ++++++++++++++++++-
 tokio/src/util/atomic_cell.rs                 |  6 ++++-
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/tokio/src/runtime/scheduler/multi_thread/idle.rs b/tokio/src/runtime/scheduler/multi_thread/idle.rs
index 79e2e184d..07ed08fed 100644
--- a/tokio/src/runtime/scheduler/multi_thread/idle.rs
+++ b/tokio/src/runtime/scheduler/multi_thread/idle.rs
@@ -194,6 +194,11 @@ impl Idle {
         }
 
         self.num_idle.store(0, Release);
+
+        // Wake up any other workers
+        while let Some(index) = synced.idle.sleepers.pop() {
+            shared.condvars[index].notify_one();
+        }
     }
 
     /// The worker releases the given core, making it available to other workers
diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs
index 99b117aba..e46332cd6 100644
--- a/tokio/src/runtime/scheduler/multi_thread/worker.rs
+++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs
@@ -1025,6 +1025,13 @@ impl Worker {
             }
         }
 
+        // If the runtime is shutdown, skip parking
+        self.update_global_flags(cx, &mut synced, &mut core);
+
+        if core.is_shutdown {
+            return Ok((None, core));
+        }
+
         // Core being returned must not be in the searching state
         debug_assert!(!core.is_searching);
 
@@ -1038,10 +1045,15 @@ impl Worker {
             // Wait for driver events
             driver.park(&self.handle.driver);
 
+            synced = cx.shared().synced.lock();
+
             // Put the driver back
             cx.shared().driver.set(driver);
 
-            synced = cx.shared().synced.lock();
+            if cx.shared().inject.is_closed(&mut synced.inject) {
+                self.shutdown_finalize(cx, synced);
+                return Err(());
+            }
 
             // Try to acquire an available core to schedule I/O events
             if let Some(core) = self.try_acquire_available_core(cx, &mut synced) {
@@ -1098,10 +1110,20 @@ impl Worker {
         let mut synced = cx.shared().synced.lock();
         synced.shutdown_cores.push(core);
 
+        self.shutdown_finalize(cx, synced);
+    }
+
+    fn shutdown_finalize(&self, cx: &Context, mut synced: MutexGuard<'_, Synced>) {
+        // Wait for all cores
         if synced.shutdown_cores.len() != cx.shared().remotes.len() {
             return;
         }
 
+        // Wait for driver
+        if cx.shared().driver.is_none() {
+            return;
+        }
+
         debug_assert!(cx.shared().owned.is_empty());
 
         for mut core in synced.shutdown_cores.drain(..) {
diff --git a/tokio/src/util/atomic_cell.rs b/tokio/src/util/atomic_cell.rs
index 07e37303a..e3cc2dd1a 100644
--- a/tokio/src/util/atomic_cell.rs
+++ b/tokio/src/util/atomic_cell.rs
@@ -1,7 +1,7 @@
 use crate::loom::sync::atomic::AtomicPtr;
 
 use std::ptr;
-use std::sync::atomic::Ordering::AcqRel;
+use std::sync::atomic::Ordering::{AcqRel, Acquire};
 
 pub(crate) struct AtomicCell<T> {
     data: AtomicPtr<T>,
@@ -29,6 +29,10 @@ impl<T> AtomicCell<T> {
     pub(crate) fn take(&self) -> Option<Box<T>> {
         self.swap(None)
     }
+
+    pub(crate) fn is_none(&self) -> bool {
+        self.data.load(Acquire).is_null()
+    }
 }
 
 fn to_raw<T>(data: Option<Box<T>>) -> *mut T {