compio_signal/
windows.rs

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

compio_signal/
windows.rs
