itsscb/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2025-09-28 12:10:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Feature: Add more windows signal handlers (#4924)

Browse Source
This commit is contained in:
Jon Kelley 2022-08-30 23:29:31 -07:00 committed by GitHub
parent 5a2bc850af
commit adc774bc4f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 403 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

302
tokio/src/signal/windows.rs
View File

@ -1,9 +1,9 @@
//! Windows-specific types for signal handling.
//!
//! This module is only defined on Windows and allows receiving "ctrl-c"
//! and "ctrl-break" notifications. These events are listened for via the
//! `SetConsoleCtrlHandler` function which receives events of the type
//! `CTRL_C_EVENT` and `CTRL_BREAK_EVENT`.
//! This module is only defined on Windows and allows receiving "ctrl-c",
//! "ctrl-break", "ctrl-logoff", "ctrl-shutdown", and "ctrl-close"
//! notifications. These events are listened for via the `SetConsoleCtrlHandler`
//! function which receives the corresponding winapi event type.
#![cfg(any(windows, docsrs))]
#![cfg_attr(docsrs, doc(cfg(all(windows, feature = "signal"))))]
@ -221,3 +221,297 @@ pub fn ctrl_break() -> io::Result<CtrlBreak> {
inner: self::imp::ctrl_break()?,
})
}
/// Creates a new stream which receives "ctrl-close" notifications sent to the
/// process.
///
/// # Examples
///
/// ```rust,no_run
/// use tokio::signal::windows::ctrl_close;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // An infinite stream of CTRL-CLOSE events.
/// let mut stream = ctrl_close()?;
///
/// // Print whenever a CTRL-CLOSE event is received.
/// for countdown in (0..3).rev() {
/// stream.recv().await;
/// println!("got CTRL-CLOSE. {} more to exit", countdown);
/// }
///
/// Ok(())
/// }
/// ```
pub fn ctrl_close() -> io::Result<CtrlClose> {
Ok(CtrlClose {
inner: self::imp::ctrl_close()?,
})
}
/// Represents a stream which receives "ctrl-close" notitifications sent to the process
/// via 'SetConsoleCtrlHandler'.
///
/// A notification to this process notifies *all* streams listening for
/// this event. Moreover, the notifications **are coalesced** if they aren't processed
/// quickly enough. This means that if two notifications are received back-to-back,
/// then the stream may only receive one item about the two notifications.
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct CtrlClose {
inner: RxFuture,
}
impl CtrlClose {
/// Receives the next signal notification event.
///
/// `None` is returned if no more events can be received by this stream.
///
/// # Examples
///
/// ```rust,no_run
/// use tokio::signal::windows::ctrl_close;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // An infinite stream of CTRL-CLOSE events.
/// let mut stream = ctrl_close()?;
///
/// // Print whenever a CTRL-CLOSE event is received.
/// stream.recv().await;
/// println!("got CTRL-CLOSE. Cleaning up before exiting");
///
/// Ok(())
/// }
/// ```
pub async fn recv(&mut self) -> Option<()> {
self.inner.recv().await
}
/// Polls to receive the next signal notification event, outside of an
/// `async` context.
///
/// `None` is returned if no more events can be received by this stream.
///
/// # Examples
///
/// Polling from a manually implemented future
///
/// ```rust,no_run
/// use std::pin::Pin;
/// use std::future::Future;
/// use std::task::{Context, Poll};
/// use tokio::signal::windows::CtrlClose;
///
/// struct MyFuture {
/// ctrl_close: CtrlClose,
/// }
///
/// impl Future for MyFuture {
/// type Output = Option<()>;
///
/// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
/// println!("polling MyFuture");
/// self.ctrl_close.poll_recv(cx)
/// }
/// }
/// ```
pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
self.inner.poll_recv(cx)
}
}
/// Creates a new stream which receives "ctrl-shutdown" notifications sent to the
/// process.
///
/// # Examples
///
/// ```rust,no_run
/// use tokio::signal::windows::ctrl_shutdown;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // An infinite stream of CTRL-SHUTDOWN events.
/// let mut stream = ctrl_shutdown()?;
///
/// stream.recv().await;
/// println!("got CTRL-SHUTDOWN. Cleaning up before exiting");
///
/// Ok(())
/// }
/// ```
pub fn ctrl_shutdown() -> io::Result<CtrlShutdown> {
Ok(CtrlShutdown {
inner: self::imp::ctrl_shutdown()?,
})
}
/// Represents a stream which receives "ctrl-shutdown" notitifications sent to the process
/// via 'SetConsoleCtrlHandler'.
///
/// A notification to this process notifies *all* streams listening for
/// this event. Moreover, the notifications **are coalesced** if they aren't processed
/// quickly enough. This means that if two notifications are received back-to-back,
/// then the stream may only receive one item about the two notifications.
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct CtrlShutdown {
inner: RxFuture,
}
impl CtrlShutdown {
/// Receives the next signal notification event.
///
/// `None` is returned if no more events can be received by this stream.
///
/// # Examples
///
/// ```rust,no_run
/// use tokio::signal::windows::ctrl_shutdown;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // An infinite stream of CTRL-SHUTDOWN events.
/// let mut stream = ctrl_shutdown()?;
///
/// // Print whenever a CTRL-SHUTDOWN event is received.
/// stream.recv().await;
/// println!("got CTRL-SHUTDOWN. Cleaning up before exiting");
///
/// Ok(())
/// }
/// ```
pub async fn recv(&mut self) -> Option<()> {
self.inner.recv().await
}
/// Polls to receive the next signal notification event, outside of an
/// `async` context.
///
/// `None` is returned if no more events can be received by this stream.
///
/// # Examples
///
/// Polling from a manually implemented future
///
/// ```rust,no_run
/// use std::pin::Pin;
/// use std::future::Future;
/// use std::task::{Context, Poll};
/// use tokio::signal::windows::CtrlShutdown;
///
/// struct MyFuture {
/// ctrl_shutdown: CtrlShutdown,
/// }
///
/// impl Future for MyFuture {
/// type Output = Option<()>;
///
/// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
/// println!("polling MyFuture");
/// self.ctrl_shutdown.poll_recv(cx)
/// }
/// }
/// ```
pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
self.inner.poll_recv(cx)
}
}
/// Creates a new stream which receives "ctrl-logoff" notifications sent to the
/// process.
///
/// # Examples
///
/// ```rust,no_run
/// use tokio::signal::windows::ctrl_logoff;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // An infinite stream of CTRL-LOGOFF events.
/// let mut stream = ctrl_logoff()?;
///
/// stream.recv().await;
/// println!("got CTRL-LOGOFF. Cleaning up before exiting");
///
/// Ok(())
/// }
/// ```
pub fn ctrl_logoff() -> io::Result<CtrlLogoff> {
Ok(CtrlLogoff {
inner: self::imp::ctrl_logoff()?,
})
}
/// Represents a stream which receives "ctrl-logoff" notitifications sent to the process
/// via 'SetConsoleCtrlHandler'.
///
/// A notification to this process notifies *all* streams listening for
/// this event. Moreover, the notifications **are coalesced** if they aren't processed
/// quickly enough. This means that if two notifications are received back-to-back,
/// then the stream may only receive one item about the two notifications.
#[must_use = "streams do nothing unless polled"]
#[derive(Debug)]
pub struct CtrlLogoff {
inner: RxFuture,
}
impl CtrlLogoff {
/// Receives the next signal notification event.
///
/// `None` is returned if no more events can be received by this stream.
///
/// # Examples
///
/// ```rust,no_run
/// use tokio::signal::windows::ctrl_logoff;
///
/// #[tokio::main]
/// async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // An infinite stream of CTRL-LOGOFF events.
/// let mut stream = ctrl_logoff()?;
///
/// // Print whenever a CTRL-LOGOFF event is received.
/// stream.recv().await;
/// println!("got CTRL-LOGOFF. Cleaning up before exiting");
///
/// Ok(())
/// }
/// ```
pub async fn recv(&mut self) -> Option<()> {
self.inner.recv().await
}
/// Polls to receive the next signal notification event, outside of an
/// `async` context.
///
/// `None` is returned if no more events can be received by this stream.
///
/// # Examples
///
/// Polling from a manually implemented future
///
/// ```rust,no_run
/// use std::pin::Pin;
/// use std::future::Future;
/// use std::task::{Context, Poll};
/// use tokio::signal::windows::CtrlLogoff;
///
/// struct MyFuture {
/// ctrl_logoff: CtrlLogoff,
/// }
///
/// impl Future for MyFuture {
/// type Output = Option<()>;
///
/// fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
/// println!("polling MyFuture");
/// self.ctrl_logoff.poll_recv(cx)
/// }
/// }
/// ```
pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<()>> {
self.inner.poll_recv(cx)
}
}

14
tokio/src/signal/windows/stub.rs
View File

@ -4,10 +4,22 @@
use crate::signal::RxFuture;
use std::io;
pub(super) fn ctrl_break() -> io::Result<RxFuture> {
panic!()
}
pub(super) fn ctrl_close() -> io::Result<RxFuture> {
panic!()
}
pub(super) fn ctrl_c() -> io::Result<RxFuture> {
panic!()
}
pub(super) fn ctrl_break() -> io::Result<RxFuture> {
pub(super) fn ctrl_logoff() -> io::Result<RxFuture> {
panic!()
}
pub(super) fn ctrl_shutdown() -> io::Result<RxFuture> {
panic!()
}

106
tokio/src/signal/windows/sys.rs
View File

@ -7,14 +7,26 @@ use crate::signal::RxFuture;
use winapi::shared::minwindef::{BOOL, DWORD, FALSE, TRUE};
use winapi::um::consoleapi::SetConsoleCtrlHandler;
use winapi::um::wincon::{CTRL_BREAK_EVENT, CTRL_C_EVENT};
pub(super) fn ctrl_c() -> io::Result<RxFuture> {
new(CTRL_C_EVENT)
}
use winapi::um::wincon;
pub(super) fn ctrl_break() -> io::Result<RxFuture> {
new(CTRL_BREAK_EVENT)
new(wincon::CTRL_BREAK_EVENT)
}
pub(super) fn ctrl_close() -> io::Result<RxFuture> {
new(wincon::CTRL_CLOSE_EVENT)
}
pub(super) fn ctrl_c() -> io::Result<RxFuture> {
new(wincon::CTRL_C_EVENT)
}
pub(super) fn ctrl_logoff() -> io::Result<RxFuture> {
new(wincon::CTRL_LOGOFF_EVENT)
}
pub(super) fn ctrl_shutdown() -> io::Result<RxFuture> {
new(wincon::CTRL_SHUTDOWN_EVENT)
}
fn new(signum: DWORD) -> io::Result<RxFuture> {
@ -25,15 +37,21 @@ fn new(signum: DWORD) -> io::Result<RxFuture> {
#[derive(Debug)]
pub(crate) struct OsStorage {
ctrl_c: EventInfo,
ctrl_break: EventInfo,
ctrl_close: EventInfo,
ctrl_c: EventInfo,
ctrl_logoff: EventInfo,
ctrl_shutdown: EventInfo,
}
impl Init for OsStorage {
fn init() -> Self {
Self {
ctrl_c: EventInfo::default(),
ctrl_break: EventInfo::default(),
ctrl_break: Default::default(),
ctrl_close: Default::default(),
ctrl_c: Default::default(),
ctrl_logoff: Default::default(),
ctrl_shutdown: Default::default(),
}
}
}
@ -41,8 +59,11 @@ impl Init for OsStorage {
impl Storage for OsStorage {
fn event_info(&self, id: EventId) -> Option<&EventInfo> {
match DWORD::try_from(id) {
Ok(CTRL_C_EVENT) => Some(&self.ctrl_c),
Ok(CTRL_BREAK_EVENT) => Some(&self.ctrl_break),
Ok(wincon::CTRL_BREAK_EVENT) => Some(&self.ctrl_break),
Ok(wincon::CTRL_CLOSE_EVENT) => Some(&self.ctrl_close),
Ok(wincon::CTRL_C_EVENT) => Some(&self.ctrl_c),
Ok(wincon::CTRL_LOGOFF_EVENT) => Some(&self.ctrl_logoff),
Ok(wincon::CTRL_SHUTDOWN_EVENT) => Some(&self.ctrl_shutdown),
_ => None,
}
}
@ -51,8 +72,11 @@ impl Storage for OsStorage {
where
F: FnMut(&'a EventInfo),
{
f(&self.ctrl_c);
f(&self.ctrl_break);
f(&self.ctrl_close);
f(&self.ctrl_c);
f(&self.ctrl_logoff);
f(&self.ctrl_shutdown);
}
}
@ -121,7 +145,7 @@ mod tests {
// like sending signals on Unix, so we'll stub out the actual OS
// integration and test that our handling works.
unsafe {
super::handler(CTRL_C_EVENT);
super::handler(wincon::CTRL_C_EVENT);
}
assert_ready_ok!(ctrl_c.poll());
@ -138,13 +162,67 @@ mod tests {
// like sending signals on Unix, so we'll stub out the actual OS
// integration and test that our handling works.
unsafe {
super::handler(CTRL_BREAK_EVENT);
super::handler(wincon::CTRL_BREAK_EVENT);
}
ctrl_break.recv().await.unwrap();
});
}
#[test]
fn ctrl_close() {
let rt = rt();
rt.block_on(async {
let mut ctrl_close = assert_ok!(crate::signal::windows::ctrl_close());
// Windows doesn't have a good programmatic way of sending events
// like sending signals on Unix, so we'll stub out the actual OS
// integration and test that our handling works.
unsafe {
super::handler(wincon::CTRL_CLOSE_EVENT);
}
ctrl_close.recv().await.unwrap();
});
}
#[test]
fn ctrl_shutdown() {
let rt = rt();
rt.block_on(async {
let mut ctrl_shutdown = assert_ok!(crate::signal::windows::ctrl_shutdown());
// Windows doesn't have a good programmatic way of sending events
// like sending signals on Unix, so we'll stub out the actual OS
// integration and test that our handling works.
unsafe {
super::handler(wincon::CTRL_SHUTDOWN_EVENT);
}
ctrl_shutdown.recv().await.unwrap();
});
}
#[test]
fn ctrl_logoff() {
let rt = rt();
rt.block_on(async {
let mut ctrl_logoff = assert_ok!(crate::signal::windows::ctrl_logoff());
// Windows doesn't have a good programmatic way of sending events
// like sending signals on Unix, so we'll stub out the actual OS
// integration and test that our handling works.
unsafe {
super::handler(wincon::CTRL_LOGOFF_EVENT);
}
ctrl_logoff.recv().await.unwrap();
});
}
fn rt() -> Runtime {
crate::runtime::Builder::new_current_thread()
.build()