use std::os::unix::io::{AsFd, BorrowedFd, OwnedFd};
use std::sync::Arc;
use rustix::event::{eventfd, EventfdFlags};
use rustix::io::{read, write, Errno};
use tracing::warn;
use super::PingError;
use crate::{
generic::Generic, EventSource, Interest, Mode, Poll, PostAction, Readiness, Token, TokenFactory,
};
const INCREMENT_PING: u64 = 0x2;
const INCREMENT_CLOSE: u64 = 0x1;
#[inline]
pub fn make_ping() -> std::io::Result<(Ping, PingSource)> {
let read = eventfd(0, EventfdFlags::CLOEXEC | EventfdFlags::NONBLOCK)?;
let fd = Arc::new(read);
let ping = Ping {
event: Arc::new(FlagOnDrop(Arc::clone(&fd))),
};
let source = PingSource {
event: Generic::new(ArcAsFd(fd), Interest::READ, Mode::Level),
};
Ok((ping, source))
}
#[inline]
fn send_ping(fd: BorrowedFd<'_>, count: u64) -> std::io::Result<()> {
assert!(count > 0);
match write(fd, &count.to_ne_bytes()) {
Ok(_) => Ok(()),
Err(Errno::AGAIN) => Ok(()),
Err(e) => Err(e.into()),
}
}
#[inline]
fn drain_ping(fd: BorrowedFd<'_>) -> std::io::Result<u64> {
const NBYTES: usize = 8;
let mut buf = [0u8; NBYTES];
match read(fd, &mut buf) {
Ok(NBYTES) => Ok(u64::from_ne_bytes(buf)),
Ok(_) => unreachable!(),
Err(e) => Err(e.into()),
}
}
#[derive(Debug)]
struct ArcAsFd(Arc<OwnedFd>);
impl AsFd for ArcAsFd {
fn as_fd(&self) -> BorrowedFd<'_> {
self.0.as_fd()
}
}
#[derive(Debug)]
pub struct PingSource {
event: Generic<ArcAsFd>,
}
impl EventSource for PingSource {
type Event = ();
type Metadata = ();
type Ret = ();
type Error = PingError;
fn process_events<C>(
&mut self,
readiness: Readiness,
token: Token,
mut callback: C,
) -> Result<PostAction, Self::Error>
where
C: FnMut(Self::Event, &mut Self::Metadata) -> Self::Ret,
{
self.event
.process_events(readiness, token, |_, fd| {
let counter = drain_ping(fd.as_fd())?;
let close = (counter & INCREMENT_CLOSE) != 0;
let ping = (counter & (u64::MAX - 1)) != 0;
if ping {
callback((), &mut ());
}
if close {
Ok(PostAction::Remove)
} else {
Ok(PostAction::Continue)
}
})
.map_err(|e| PingError(e.into()))
}
fn register(&mut self, poll: &mut Poll, token_factory: &mut TokenFactory) -> crate::Result<()> {
self.event.register(poll, token_factory)
}
fn reregister(
&mut self,
poll: &mut Poll,
token_factory: &mut TokenFactory,
) -> crate::Result<()> {
self.event.reregister(poll, token_factory)
}
fn unregister(&mut self, poll: &mut Poll) -> crate::Result<()> {
self.event.unregister(poll)
}
}
#[derive(Clone, Debug)]
pub struct Ping {
event: Arc<FlagOnDrop>,
}
impl Ping {
pub fn ping(&self) {
if let Err(e) = send_ping(self.event.0.as_fd(), INCREMENT_PING) {
warn!("Failed to write a ping: {e:?}");
}
}
}
#[derive(Debug)]
struct FlagOnDrop(Arc<OwnedFd>);
impl Drop for FlagOnDrop {
fn drop(&mut self) {
if let Err(e) = send_ping(self.0.as_fd(), INCREMENT_CLOSE) {
warn!("Failed to send close ping: {e:?}");
}
}
}