diff --git a/tokio/src/signal/windows/sys.rs b/tokio/src/signal/windows/sys.rs index dda379978..d9c9247f5 100644 --- a/tokio/src/signal/windows/sys.rs +++ b/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 { - new(®istry().ctrl_break) + new(SignalKind::CtrlBreak) } pub(super) fn ctrl_close() -> io::Result { - new(®istry().ctrl_close) + new(SignalKind::CtrlClose) } pub(super) fn ctrl_c() -> io::Result { - new(®istry().ctrl_c) + new(SignalKind::CtrlC) } pub(super) fn ctrl_logoff() -> io::Result { - new(®istry().ctrl_logoff) + new(SignalKind::CtrlLogoff) } pub(super) fn ctrl_shutdown() -> io::Result { - new(®istry().ctrl_shutdown) + new(SignalKind::CtrlShutdown) } -fn new(event_info: &EventInfo) -> io::Result { - global_init()?; - let rx = event_info.subscribe(); +fn new(signal: SignalKind) -> io::Result { + let registry = REGISTRY + .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 { - fn event_info(&self, signum: u32) -> Option<&EventInfo> { - match signum { - console::CTRL_BREAK_EVENT => Some(&self.ctrl_break), - console::CTRL_CLOSE_EVENT => Some(&self.ctrl_close), - console::CTRL_C_EVENT => Some(&self.ctrl_c), - console::CTRL_LOGOFF_EVENT => Some(&self.ctrl_logoff), - console::CTRL_SHUTDOWN_EVENT => Some(&self.ctrl_shutdown), - _ => None, +impl Index for Registry { + type Output = EventInfo; + + fn index(&self, signal: SignalKind) -> &Self::Output { + match signal { + SignalKind::CtrlC => &self.ctrl_c, + SignalKind::CtrlBreak => &self.ctrl_break, + SignalKind::CtrlClose => &self.ctrl_close, + SignalKind::CtrlLogoff => &self.ctrl_logoff, + SignalKind::CtrlShutdown => &self.ctrl_shutdown, } } } -fn registry() -> &'static Registry { - static REGISTRY: OnceLock = OnceLock::new(); - - REGISTRY.get_or_init(Default::default) -} - -fn global_init() -> io::Result<()> { - static INIT: OnceLock>> = 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, - ) - }) -} +static REGISTRY: OnceLock> = OnceLock::new(); unsafe extern "system" fn handler(ty: u32) -> BOOL { - // Ignore unknown control signal types. - let Some(event_info) = registry().event_info(ty) else { + let signal = match ty { + 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(()) { - Ok(_) if event_requires_infinite_sleep_in_handler(ty) => loop { + match registry[signal].send(()) { + 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) { - if event_requires_infinite_sleep_in_handler(signum) { + unsafe fn raise_event(signal: SignalKind) { + 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();