Skip to main content

compio_signal/
windows.rs

1//! Windows-specific types for signal handling.
2
3#[cfg(feature = "once_cell_try")]
4use std::sync::OnceLock;
5use std::{
6    collections::HashMap,
7    io,
8    sync::{LazyLock, Mutex},
9};
10
11#[cfg(not(feature = "once_cell_try"))]
12use once_cell::sync::OnceCell as OnceLock;
13use slab::Slab;
14use synchrony::sync::async_flag::{AsyncFlag as Event, AsyncFlagHandle as EventHandle};
15use windows_sys::{
16    Win32::System::Console::{
17        CTRL_BREAK_EVENT, CTRL_C_EVENT, CTRL_CLOSE_EVENT, CTRL_LOGOFF_EVENT, CTRL_SHUTDOWN_EVENT,
18        SetConsoleCtrlHandler,
19    },
20    core::BOOL,
21};
22
23static HANDLER: LazyLock<Mutex<HashMap<u32, Slab<EventHandle>>>> =
24    LazyLock::new(|| Mutex::new(HashMap::new()));
25
26unsafe extern "system" fn ctrl_event_handler(ctrltype: u32) -> BOOL {
27    let mut handler = HANDLER.lock().unwrap();
28    if let Some(handlers) = handler.get_mut(&ctrltype)
29        && !handlers.is_empty()
30    {
31        let handlers = std::mem::replace(handlers, Slab::new());
32        for (_, handler) in handlers {
33            handler.notify();
34        }
35        return 1;
36    }
37    0
38}
39
40static INIT: OnceLock<()> = OnceLock::new();
41
42fn init() -> io::Result<()> {
43    if unsafe { SetConsoleCtrlHandler(Some(ctrl_event_handler), 1) } != 0 {
44        Ok(())
45    } else {
46        Err(io::Error::last_os_error())
47    }
48}
49
50fn register(ctrltype: u32, e: &Event) -> usize {
51    let mut handler = HANDLER.lock().unwrap();
52    let handle = e.handle();
53    handler.entry(ctrltype).or_default().insert(handle)
54}
55
56fn unregister(ctrltype: u32, key: usize) {
57    let mut handler = HANDLER.lock().unwrap();
58    if let Some(handlers) = handler.get_mut(&ctrltype)
59        && handlers.contains(key)
60    {
61        let _ = handlers.remove(key);
62    }
63}
64
65/// Represents a listener to console CTRL event.
66#[derive(Debug)]
67struct CtrlEvent {
68    ctrltype: u32,
69    event: Option<Event>,
70    handler_key: usize,
71}
72
73impl CtrlEvent {
74    pub(crate) fn new(ctrltype: u32) -> io::Result<Self> {
75        INIT.get_or_try_init(init)?;
76
77        let event = Event::new();
78        let handler_key = register(ctrltype, &event);
79        Ok(Self {
80            ctrltype,
81            event: Some(event),
82            handler_key,
83        })
84    }
85
86    pub async fn wait(mut self) {
87        self.event
88            .take()
89            .expect("event could not be None")
90            .wait()
91            .await
92    }
93}
94
95impl Drop for CtrlEvent {
96    fn drop(&mut self) {
97        unregister(self.ctrltype, self.handler_key);
98    }
99}
100
101async fn ctrl_event(ctrltype: u32) -> io::Result<()> {
102    let event = CtrlEvent::new(ctrltype)?;
103    event.wait().await;
104    Ok(())
105}
106
107/// Creates a new listener which receives "ctrl-break" notifications sent to the
108/// process.
109pub async fn ctrl_break() -> io::Result<()> {
110    ctrl_event(CTRL_BREAK_EVENT).await
111}
112
113/// Creates a new listener which receives "ctrl-close" notifications sent to the
114/// process.
115pub async fn ctrl_close() -> io::Result<()> {
116    ctrl_event(CTRL_CLOSE_EVENT).await
117}
118
119/// Creates a new listener which receives "ctrl-c" notifications sent to the
120/// process.
121pub async fn ctrl_c() -> io::Result<()> {
122    ctrl_event(CTRL_C_EVENT).await
123}
124
125/// Creates a new listener which receives "ctrl-logoff" notifications sent to
126/// the process.
127pub async fn ctrl_logoff() -> io::Result<()> {
128    ctrl_event(CTRL_LOGOFF_EVENT).await
129}
130
131/// Creates a new listener which receives "ctrl-shutdown" notifications sent to
132/// the process.
133pub async fn ctrl_shutdown() -> io::Result<()> {
134    ctrl_event(CTRL_SHUTDOWN_EVENT).await
135}