compio_signal/
unix.rs

1//! Unix-specific types for signal handling.
2
3#[cfg(feature = "lazy_cell")]
4use std::sync::LazyLock;
5use std::{
6    collections::HashMap,
7    io::{self, Read, Write},
8    ops::Deref,
9    sync::Mutex,
10};
11
12use compio_runtime::event::{Event, EventHandle};
13#[cfg(not(feature = "lazy_cell"))]
14use once_cell::sync::Lazy as LazyLock;
15use os_pipe::{PipeReader, PipeWriter};
16use slab::Slab;
17
18static HANDLER: LazyLock<Mutex<HashMap<i32, Slab<EventHandle>>>> =
19    LazyLock::new(|| Mutex::new(HashMap::new()));
20static PIPE: LazyLock<Pipe> = LazyLock::new(|| Pipe::new().unwrap());
21
22struct Pipe {
23    sender: PipeWriter,
24}
25
26impl Pipe {
27    pub fn new() -> io::Result<Self> {
28        let (receiver, sender) = os_pipe::pipe()?;
29
30        std::thread::spawn(move || {
31            real_signal_handler(receiver);
32        });
33
34        Ok(Self { sender })
35    }
36
37    pub fn send(&self, sig: i32) -> io::Result<()> {
38        (&self.sender).write_all(&sig.to_ne_bytes())?;
39        Ok(())
40    }
41}
42
43unsafe extern "C" fn signal_handler(sig: i32) {
44    PIPE.send(sig).unwrap();
45}
46
47fn real_signal_handler(mut receiver: PipeReader) {
48    loop {
49        let mut buffer = [0u8; 4];
50        let res = receiver.read_exact(&mut buffer);
51        if let Ok(()) = res {
52            let sig = i32::from_ne_bytes(buffer);
53            let mut handler = HANDLER.lock().unwrap();
54            if let Some(fds) = handler.get_mut(&sig) {
55                if !fds.is_empty() {
56                    let fds = std::mem::take(fds);
57                    for (_, fd) in fds {
58                        fd.notify();
59                    }
60                }
61            }
62        } else {
63            break;
64        }
65    }
66}
67
68unsafe fn init(sig: i32) -> io::Result<()> {
69    let _ = PIPE.deref();
70    if libc::signal(sig, signal_handler as *const () as usize) == libc::SIG_ERR {
71        Err(io::Error::last_os_error())
72    } else {
73        Ok(())
74    }
75}
76
77unsafe fn uninit(sig: i32) -> io::Result<()> {
78    if libc::signal(sig, libc::SIG_DFL) == libc::SIG_ERR {
79        Err(io::Error::last_os_error())
80    } else {
81        Ok(())
82    }
83}
84
85fn register(sig: i32, fd: &Event) -> io::Result<usize> {
86    unsafe { init(sig)? };
87    let handle = fd.handle();
88    let key = HANDLER
89        .lock()
90        .unwrap()
91        .entry(sig)
92        .or_default()
93        .insert(handle);
94    Ok(key)
95}
96
97fn unregister(sig: i32, key: usize) {
98    let need_uninit = (|| {
99        let mut handler = HANDLER.lock().unwrap();
100        if let Some(fds) = handler.get_mut(&sig) {
101            fds.try_remove(key);
102            if !fds.is_empty() {
103                return false;
104            }
105        }
106        true
107    })();
108    if need_uninit {
109        unsafe { uninit(sig).ok() };
110    }
111}
112
113/// Represents a listener to unix signal event.
114#[derive(Debug)]
115struct SignalFd {
116    sig: i32,
117    key: usize,
118    event: Option<Event>,
119}
120
121impl SignalFd {
122    fn new(sig: i32) -> io::Result<Self> {
123        let event = Event::new();
124        let key = register(sig, &event)?;
125        Ok(Self {
126            sig,
127            key,
128            event: Some(event),
129        })
130    }
131
132    async fn wait(mut self) {
133        self.event
134            .take()
135            .expect("event could not be None")
136            .wait()
137            .await
138    }
139}
140
141impl Drop for SignalFd {
142    fn drop(&mut self) {
143        unregister(self.sig, self.key);
144    }
145}
146
147/// Creates a new listener which will receive notifications when the current
148/// process receives the specified signal.
149///
150/// The first call to this method spawns a thread to execute the signal
151/// handlers.
152pub async fn signal(sig: i32) -> io::Result<()> {
153    let fd = SignalFd::new(sig)?;
154    fd.wait().await;
155    Ok(())
156}