1#[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#[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
107pub async fn ctrl_break() -> io::Result<()> {
110 ctrl_event(CTRL_BREAK_EVENT).await
111}
112
113pub async fn ctrl_close() -> io::Result<()> {
116 ctrl_event(CTRL_CLOSE_EVENT).await
117}
118
119pub async fn ctrl_c() -> io::Result<()> {
122 ctrl_event(CTRL_C_EVENT).await
123}
124
125pub async fn ctrl_logoff() -> io::Result<()> {
128 ctrl_event(CTRL_LOGOFF_EVENT).await
129}
130
131pub async fn ctrl_shutdown() -> io::Result<()> {
134 ctrl_event(CTRL_SHUTDOWN_EVENT).await
135}