1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

//! Unix-specific types for signal handling.

use std::{
    cell::RefCell,
    collections::HashMap,
    io,
    os::fd::{AsRawFd, RawFd},
};

use compio_runtime::event::{Event, EventHandle};

thread_local! {
    static HANDLER: RefCell<HashMap<i32, HashMap<RawFd, EventHandle>>> =
        RefCell::new(HashMap::new());
}

unsafe extern "C" fn signal_handler(sig: i32) {
    HANDLER.with(|handler| {
        let mut handler = handler.borrow_mut();
        if let Some(fds) = handler.get_mut(&sig) {
            if !fds.is_empty() {
                let fds = std::mem::take(fds);
                for (_, fd) in fds {
                    fd.notify().ok();
                }
            }
        }
    });
}

unsafe fn init(sig: i32) {
    libc::signal(sig, signal_handler as *const () as usize);
}

unsafe fn uninit(sig: i32) {
    libc::signal(sig, libc::SIG_DFL);
}

fn register(sig: i32, fd: &Event) -> io::Result<()> {
    unsafe { init(sig) };
    let raw_fd = fd.as_raw_fd();
    let handle = fd.handle()?;
    HANDLER.with(|handler| {
        handler
            .borrow_mut()
            .entry(sig)
            .or_default()
            .insert(raw_fd, handle)
    });
    Ok(())
}

fn unregister(sig: i32, fd: RawFd) {
    let need_uninit = HANDLER.with(|handler| {
        let mut handler = handler.borrow_mut();
        if let Some(fds) = handler.get_mut(&sig) {
            fds.remove(&fd);
            if !fds.is_empty() {
                return false;
            }
        }
        true
    });
    if need_uninit {
        unsafe { uninit(sig) };
    }
}

/// Represents a listener to unix signal event.
#[derive(Debug)]
struct SignalFd {
    sig: i32,
    fd: Event,
}

impl SignalFd {
    fn new(sig: i32) -> io::Result<Self> {
        let fd = Event::new()?;
        register(sig, &fd)?;
        Ok(Self { sig, fd })
    }

    async fn wait(&self) -> io::Result<()> {
        self.fd.wait().await
    }
}

impl Drop for SignalFd {
    fn drop(&mut self) {
        unregister(self.sig, self.fd.as_raw_fd());
    }
}

/// Creates a new listener which will receive notifications when the current
/// process receives the specified signal.
pub async fn signal(sig: i32) -> io::Result<()> {
    let fd = SignalFd::new(sig)?;
    fd.wait().await
}