Skip to main content

aria_primitives/
wakeup.rs

1//! A self-pipe: something that has to interrupt a poll loop writes to it, and
2//! the loop sees an ordinary readable file descriptor.
3//!
4//! Both ends are non-blocking, so neither end can stall the thread it runs on.
5//! What the bytes mean is the caller's business.
6
7use std::io::{self, Read, Write};
8use std::os::fd::{AsFd, BorrowedFd};
9use std::os::unix::net::UnixStream;
10
11const DRAIN_BUFFER_BYTES: usize = 64;
12
13pub struct Wakeup {
14    reader: UnixStream,
15    writer: UnixStream,
16}
17
18impl Wakeup {
19    pub fn new() -> Result<Self, WakeupError> {
20        let (reader, writer) = UnixStream::pair().map_err(WakeupError::Create)?;
21        reader
22            .set_nonblocking(true)
23            .map_err(WakeupError::Configure)?;
24        writer
25            .set_nonblocking(true)
26            .map_err(WakeupError::Configure)?;
27        Ok(Self { reader, writer })
28    }
29
30    pub fn as_fd(&self) -> BorrowedFd<'_> {
31        self.reader.as_fd()
32    }
33
34    /// For a signal handler, and anything else that takes an owned writer.
35    pub fn clone_writer(&self) -> Result<UnixStream, WakeupError> {
36        self.writer.try_clone().map_err(WakeupError::Duplicate)
37    }
38
39    // A full pipe already carries the message, so a blocked write is success.
40    pub fn signal(&self) -> Result<(), WakeupError> {
41        match (&self.writer).write_all(b"X") {
42            Ok(()) => Ok(()),
43            Err(error) if error.kind() == io::ErrorKind::WouldBlock => Ok(()),
44            Err(error) => Err(WakeupError::Signal(error)),
45        }
46    }
47
48    pub fn drain(&self) -> Result<(), WakeupError> {
49        let mut reader = &self.reader;
50        let mut buffer = [0_u8; DRAIN_BUFFER_BYTES];
51        loop {
52            match reader.read(&mut buffer) {
53                Ok(0) => return Ok(()),
54                Ok(_read) => {}
55                Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
56                Err(error) if error.kind() == io::ErrorKind::WouldBlock => return Ok(()),
57                Err(error) => return Err(WakeupError::Drain(error)),
58            }
59        }
60    }
61}
62
63#[derive(Debug, thiserror::Error)]
64pub enum WakeupError {
65    #[error("failed to create the wakeup socket pair")]
66    Create(#[source] io::Error),
67
68    #[error("failed to configure the wakeup socket pair")]
69    Configure(#[source] io::Error),
70
71    #[error("failed to duplicate the wakeup socket")]
72    Duplicate(#[source] io::Error),
73
74    #[error("failed to signal the wakeup socket")]
75    Signal(#[source] io::Error),
76
77    #[error("failed to drain the wakeup socket")]
78    Drain(#[source] io::Error),
79}