From 6c19748f901a1b8cd7fbf84650f1647342b5a66e Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Thu, 27 Oct 2022 13:07:51 -0700
Subject: [PATCH] rt: use signal driver handle via scheduler::Handle (#5135)

The signal driver still uses an `Arc` internally to track if the driver
is still running, however, using the `scheduler::Handle` to access the
signal driver handle lets us delete some code.
---
 tokio/src/runtime/context.rs       | 17 ----------
 tokio/src/runtime/driver.rs        |  9 ++++++
 tokio/src/runtime/scheduler/mod.rs |  7 -----
 tokio/src/runtime/signal/mod.rs    | 50 ++++++------------------------
 tokio/src/signal/unix.rs           |  4 ++-
 5 files changed, 21 insertions(+), 66 deletions(-)

diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs
index 9b20ecbcb50..c816959c17d 100644
--- a/tokio/src/runtime/context.rs
+++ b/tokio/src/runtime/context.rs
@@ -24,23 +24,6 @@ pub(crate) fn current() -> Handle {
     }
 }
 
-cfg_signal_internal! {
-    #[cfg(unix)]
-    pub(crate) fn signal_handle() -> crate::runtime::driver::SignalHandle {
-        match CONTEXT.try_with(|ctx| {
-            let ctx = ctx.borrow();
-            ctx.as_ref()
-                .expect(crate::util::error::CONTEXT_MISSING_ERROR)
-                .inner
-                .signal()
-                .clone()
-        }) {
-            Ok(signal_handle) => signal_handle,
-            Err(_) => panic!("{}", crate::util::error::THREAD_LOCAL_DESTROYED_ERROR),
-        }
-    }
-}
-
 cfg_time! {
     cfg_test_util! {
         pub(crate) fn clock() -> Option<crate::runtime::driver::Clock> {
diff --git a/tokio/src/runtime/driver.rs b/tokio/src/runtime/driver.rs
index 184e736a90a..a72b2569851 100644
--- a/tokio/src/runtime/driver.rs
+++ b/tokio/src/runtime/driver.rs
@@ -90,6 +90,15 @@ impl Handle {
         }
     }
 
+    cfg_signal_internal_and_unix! {
+        #[track_caller]
+        pub(crate) fn signal(&self) -> &crate::runtime::signal::Handle {
+            self.signal
+                .as_ref()
+                .expect("there is no signal driver running, must be called from the context of Tokio runtime")
+        }
+    }
+
     cfg_time! {
         /// Returns a reference to the time driver handle.
         ///
diff --git a/tokio/src/runtime/scheduler/mod.rs b/tokio/src/runtime/scheduler/mod.rs
index 31546673bbe..4805da4720c 100644
--- a/tokio/src/runtime/scheduler/mod.rs
+++ b/tokio/src/runtime/scheduler/mod.rs
@@ -113,13 +113,6 @@ cfg_rt! {
                 Handle::MultiThread(h) => &h.seed_generator,
             }
         }
-
-        #[cfg(unix)]
-        cfg_signal_internal! {
-            pub(crate) fn signal(&self) -> &driver::SignalHandle {
-                &self.driver().signal
-            }
-        }
     }
 
     cfg_metrics! {
diff --git a/tokio/src/runtime/signal/mod.rs b/tokio/src/runtime/signal/mod.rs
index c8ba951c6ed..dba30ba2023 100644
--- a/tokio/src/runtime/signal/mod.rs
+++ b/tokio/src/runtime/signal/mod.rs
@@ -23,18 +23,19 @@ pub(crate) struct Driver {
     /// A pipe for receiving wake events from the signal handler
     receiver: UnixStream,
 
-    /// Shared state
-    inner: Arc<Inner>,
+    /// Shared state. The driver keeps a strong ref and the handle keeps a weak
+    /// ref. The weak ref is used to check if the driver is still active before
+    /// trying to register a signal handler.
+    inner: Arc<()>,
 }
 
-#[derive(Clone, Debug, Default)]
+#[derive(Debug, Default)]
 pub(crate) struct Handle {
-    inner: Weak<Inner>,
+    /// Paired w/ the `Arc` above and is used to check if the driver is still
+    /// around before attempting to register a signal handler.
+    inner: Weak<()>,
 }
 
-#[derive(Debug)]
-pub(super) struct Inner(());
-
 // ===== impl Driver =====
 
 impl Driver {
@@ -75,7 +76,7 @@ impl Driver {
         Ok(Self {
             io,
             receiver,
-            inner: Arc::new(Inner(())),
+            inner: Arc::new(()),
         })
     }
 
@@ -139,36 +140,3 @@ impl Handle {
         }
     }
 }
-
-cfg_rt! {
-    impl Handle {
-        /// Returns a handle to the current driver
-        ///
-        /// # Panics
-        ///
-        /// This function panics if there is no current signal driver set.
-        #[track_caller]
-        pub(crate) fn current() -> Self {
-            crate::runtime::context::signal_handle().expect(
-                "there is no signal driver running, must be called from the context of Tokio runtime",
-            )
-        }
-    }
-}
-
-cfg_not_rt! {
-    impl Handle {
-        /// Returns a handle to the current driver
-        ///
-        /// # Panics
-        ///
-        /// This function panics if there is no current signal driver set.
-        #[track_caller]
-        pub(crate) fn current() -> Self {
-            panic!(
-                "there is no signal driver running, must be called from the context of Tokio runtime or with\
-                `rt` enabled.",
-            )
-        }
-    }
-}
diff --git a/tokio/src/signal/unix.rs b/tokio/src/signal/unix.rs
index f85c09f4dd6..0e1329ed54f 100644
--- a/tokio/src/signal/unix.rs
+++ b/tokio/src/signal/unix.rs
@@ -6,6 +6,7 @@
 #![cfg(unix)]
 #![cfg_attr(docsrs, doc(cfg(all(unix, feature = "signal"))))]
 
+use crate::runtime::scheduler;
 use crate::runtime::signal::Handle;
 use crate::signal::registry::{globals, EventId, EventInfo, Globals, Init, Storage};
 use crate::signal::RxFuture;
@@ -389,7 +390,8 @@ pub struct Signal {
 /// feature flag is not enabled.
 #[track_caller]
 pub fn signal(kind: SignalKind) -> io::Result<Signal> {
-    let rx = signal_with_handle(kind, &Handle::current())?;
+    let handle = scheduler::Handle::current();
+    let rx = signal_with_handle(kind, handle.driver().signal())?;
 
     Ok(Signal {
         inner: RxFuture::new(rx),