signal: merge windows statics (#8231)

tokio/src/signal/windows/sys.rs

@@ -1,4 +1,6 @@
 use std::io;
+use std::io::Error;
+use std::ops::Index;
 use std::sync::OnceLock;
 
 use crate::signal::RxFuture;
@@ -9,45 +11,66 @@ use windows_sys::Win32::System::Console as console;
 
 type EventInfo = watch::Sender<()>;
 
+#[derive(Clone, Copy)]
+#[repr(u32)]
+enum SignalKind {
+    CtrlC = console::CTRL_C_EVENT,
+    CtrlBreak = console::CTRL_BREAK_EVENT,
+    CtrlClose = console::CTRL_CLOSE_EVENT,
+    CtrlLogoff = console::CTRL_LOGOFF_EVENT,
+    CtrlShutdown = console::CTRL_SHUTDOWN_EVENT,
+}
+
+impl SignalKind {
+    const fn terminates(&self) -> bool {
+        // Returning from the handler function of those events immediately terminates the process.
+        // So for async systems, the easiest solution is to simply never return from
+        // the handler function.
+        //
+        // For more information, see:
+        // https://learn.microsoft.com/en-us/windows/console/handlerroutine#remarks
+        matches!(
+            self,
+            Self::CtrlClose | Self::CtrlLogoff | Self::CtrlShutdown
+        )
+    }
+}
+
 pub(super) fn ctrl_break() -> io::Result<RxFuture> {
-    new(&registry().ctrl_break)
+    new(SignalKind::CtrlBreak)
 }
 
 pub(super) fn ctrl_close() -> io::Result<RxFuture> {
-    new(&registry().ctrl_close)
+    new(SignalKind::CtrlClose)
 }
 
 pub(super) fn ctrl_c() -> io::Result<RxFuture> {
-    new(&registry().ctrl_c)
+    new(SignalKind::CtrlC)
 }
 
 pub(super) fn ctrl_logoff() -> io::Result<RxFuture> {
-    new(&registry().ctrl_logoff)
+    new(SignalKind::CtrlLogoff)
 }
 
 pub(super) fn ctrl_shutdown() -> io::Result<RxFuture> {
-    new(&registry().ctrl_shutdown)
+    new(SignalKind::CtrlShutdown)
 }
 
-fn new(event_info: &EventInfo) -> io::Result<RxFuture> {
+fn new(signal: SignalKind) -> io::Result<RxFuture> {
-    global_init()?;
+    let registry = REGISTRY
-    let rx = event_info.subscribe();
+        .get_or_init(
+            || match unsafe { console::SetConsoleCtrlHandler(Some(handler), 1) } {
+                0 => Err(Error::last_os_error().raw_os_error().expect("unreachable")),
+                _ => Ok(Registry::default()),
+            },
+        )
+        .as_ref()
+        .map_err(|&code| Error::from_raw_os_error(code))?;
+
+    let rx = registry[signal].subscribe();
     Ok(RxFuture::new(rx))
 }
 
-fn event_requires_infinite_sleep_in_handler(signum: u32) -> bool {
-    // Returning from the handler function of those events immediately terminates the process.
-    // So for async systems, the easiest solution is to simply never return from
-    // the handler function.
-    //
-    // For more information, see:
-    // https://learn.microsoft.com/en-us/windows/console/handlerroutine#remarks
-    matches!(
-        signum,
-        console::CTRL_CLOSE_EVENT | console::CTRL_LOGOFF_EVENT | console::CTRL_SHUTDOWN_EVENT
-    )
-}
-
 #[derive(Debug, Default)]
 struct Registry {
     ctrl_break: EventInfo,
@@ -57,47 +80,38 @@ struct Registry {
     ctrl_shutdown: EventInfo,
 }
 
-impl Registry {
+impl Index<SignalKind> for Registry {
-    fn event_info(&self, signum: u32) -> Option<&EventInfo> {
+    type Output = EventInfo;
-        match signum {
+
-            console::CTRL_BREAK_EVENT => Some(&self.ctrl_break),
+    fn index(&self, signal: SignalKind) -> &Self::Output {
-            console::CTRL_CLOSE_EVENT => Some(&self.ctrl_close),
+        match signal {
-            console::CTRL_C_EVENT => Some(&self.ctrl_c),
+            SignalKind::CtrlC => &self.ctrl_c,
-            console::CTRL_LOGOFF_EVENT => Some(&self.ctrl_logoff),
+            SignalKind::CtrlBreak => &self.ctrl_break,
-            console::CTRL_SHUTDOWN_EVENT => Some(&self.ctrl_shutdown),
+            SignalKind::CtrlClose => &self.ctrl_close,
-            _ => None,
+            SignalKind::CtrlLogoff => &self.ctrl_logoff,
+            SignalKind::CtrlShutdown => &self.ctrl_shutdown,
         }
     }
 }
 
-fn registry() -> &'static Registry {
+static REGISTRY: OnceLock<Result<Registry, i32>> = OnceLock::new();
-    static REGISTRY: OnceLock<Registry> = OnceLock::new();
-
-    REGISTRY.get_or_init(Default::default)
-}
-
-fn global_init() -> io::Result<()> {
-    static INIT: OnceLock<Result<(), Option<i32>>> = OnceLock::new();
-
-    INIT.get_or_init(|| {
-        let rc = unsafe { console::SetConsoleCtrlHandler(Some(handler), 1) };
-        if rc == 0 {
-            Err(io::Error::last_os_error().raw_os_error())
-        } else {
-            Ok(())
-        }
-    })
-    .map_err(|e| {
-        e.map_or_else(
-            || io::Error::new(io::ErrorKind::Other, "registering signal handler failed"),
-            io::Error::from_raw_os_error,
-        )
-    })
-}
 
 unsafe extern "system" fn handler(ty: u32) -> BOOL {
-    // Ignore unknown control signal types.
+    let signal = match ty {
-    let Some(event_info) = registry().event_info(ty) else {
+        console::CTRL_C_EVENT => SignalKind::CtrlC,
+        console::CTRL_BREAK_EVENT => SignalKind::CtrlBreak,
+        console::CTRL_CLOSE_EVENT => SignalKind::CtrlClose,
+        console::CTRL_LOGOFF_EVENT => SignalKind::CtrlLogoff,
+        console::CTRL_SHUTDOWN_EVENT => SignalKind::CtrlShutdown,
+        // Ignore unknown signals.
+        _ => return 0,
+    };
+
+    // Note that `OnceLock::get` does not handle the small window between calling
+    // `SetConsoleCtrlHandler` and `REGISTRY` being initialized.
+    let Ok(registry) = REGISTRY.wait().as_ref() else {
+        // Technically unreachable since `handler` is only called if
+        // `SetConsoleCtrlHandler` succeded.
         return 0;
     };
 
@@ -105,8 +119,8 @@ unsafe extern "system" fn handler(ty: u32) -> BOOL {
     // the handler routine is always invoked in a new thread, thus we don't
     // have the same restrictions as in Unix signal handlers, meaning we can
     // go ahead and perform the broadcast here.
-    match event_info.send(()) {
+    match registry[signal].send(()) {
-        Ok(_) if event_requires_infinite_sleep_in_handler(ty) => loop {
+        Ok(_) if signal.terminates() => loop {
             std::thread::park();
         },
         Ok(_) => 1,
@@ -123,13 +137,13 @@ mod tests {
 
     use tokio_test::{assert_ok, assert_pending, assert_ready_ok, task};
 
-    unsafe fn raise_event(signum: u32) {
+    unsafe fn raise_event(signal: SignalKind) {
-        if event_requires_infinite_sleep_in_handler(signum) {
+        if signal.terminates() {
             // Those events will enter an infinite loop in `handler`, so
             // we need to run them on a separate thread
-            std::thread::spawn(move || unsafe { super::handler(signum) });
+            std::thread::spawn(move || unsafe { super::handler(signal as u32) });
         } else {
-            unsafe { super::handler(signum) };
+            unsafe { super::handler(signal as u32) };
         }
     }
 
@@ -146,7 +160,7 @@ mod tests {
         // like sending signals on Unix, so we'll stub out the actual OS
         // integration and test that our handling works.
         unsafe {
-            raise_event(console::CTRL_C_EVENT);
+            raise_event(SignalKind::CtrlC);
         }
 
         assert_ready_ok!(ctrl_c.poll());
@@ -163,7 +177,7 @@ mod tests {
             // like sending signals on Unix, so we'll stub out the actual OS
             // integration and test that our handling works.
             unsafe {
-                raise_event(console::CTRL_BREAK_EVENT);
+                raise_event(SignalKind::CtrlBreak);
             }
 
             ctrl_break.recv().await.unwrap();
@@ -181,7 +195,7 @@ mod tests {
             // like sending signals on Unix, so we'll stub out the actual OS
             // integration and test that our handling works.
             unsafe {
-                raise_event(console::CTRL_CLOSE_EVENT);
+                raise_event(SignalKind::CtrlClose);
             }
 
             ctrl_close.recv().await.unwrap();
@@ -199,7 +213,7 @@ mod tests {
             // like sending signals on Unix, so we'll stub out the actual OS
             // integration and test that our handling works.
             unsafe {
-                raise_event(console::CTRL_SHUTDOWN_EVENT);
+                raise_event(SignalKind::CtrlShutdown);
             }
 
             ctrl_shutdown.recv().await.unwrap();
@@ -217,7 +231,7 @@ mod tests {
             // like sending signals on Unix, so we'll stub out the actual OS
             // integration and test that our handling works.
             unsafe {
-                raise_event(console::CTRL_LOGOFF_EVENT);
+                raise_event(SignalKind::CtrlLogoff);
             }
 
             ctrl_logoff.recv().await.unwrap();