use std::os::fd::{BorrowedFd, RawFd};
use std::time::Duration;
use nix::errno::Errno;
use nix::poll::{poll, PollFd, PollFlags, PollTimeout};
pub(crate) fn wait_readable(fd: RawFd, timeout: Option<Duration>) -> std::io::Result<bool> {
let timeout_ms = poll_timeout(timeout);
loop {
let borrowed = unsafe { BorrowedFd::borrow_raw(fd) };
let mut fds = [PollFd::new(borrowed, PollFlags::POLLIN)];
match poll(&mut fds, timeout_ms) {
Ok(0) => return Ok(false),
Ok(_) => return Ok(true),
Err(Errno::EINTR) => continue,
Err(e) => return Err(std::io::Error::from(e)),
}
}
}
pub(crate) fn wait_any_readable(fds: &[RawFd], timeout: Option<Duration>) -> std::io::Result<bool> {
let timeout_ms = poll_timeout(timeout);
loop {
let borrowed: Vec<BorrowedFd> = fds
.iter()
.map(|&fd| unsafe { BorrowedFd::borrow_raw(fd) })
.collect();
let mut pollfds: Vec<PollFd> = borrowed
.iter()
.map(|&fd| PollFd::new(fd, PollFlags::POLLIN))
.collect();
match poll(&mut pollfds, timeout_ms) {
Ok(0) => return Ok(false),
Ok(_) => return Ok(true),
Err(Errno::EINTR) => continue,
Err(e) => return Err(std::io::Error::from(e)),
}
}
}
pub(crate) const MAX_POLL_FDS: usize = 4;
pub(crate) fn poll_readable_into(
fds: &[RawFd],
ready: &mut [bool],
timeout: Option<Duration>,
) -> std::io::Result<usize> {
assert!(!fds.is_empty(), "poll_readable_into: no fds");
assert!(
fds.len() <= MAX_POLL_FDS,
"poll_readable_into: at most {MAX_POLL_FDS} fds"
);
assert!(
ready.len() >= fds.len(),
"poll_readable_into: ready too short"
);
let timeout_ms = poll_timeout(timeout);
loop {
for r in ready[..fds.len()].iter_mut() {
*r = false;
}
let borrowed: [BorrowedFd; MAX_POLL_FDS] = std::array::from_fn(|i| unsafe {
BorrowedFd::borrow_raw(if i < fds.len() { fds[i] } else { fds[0] })
});
let mut pollfds: [PollFd; MAX_POLL_FDS] =
std::array::from_fn(|i| PollFd::new(borrowed[i], PollFlags::POLLIN));
match poll(&mut pollfds[..fds.len()], timeout_ms) {
Ok(0) => return Ok(0),
Ok(_) => {
let mut n = 0;
for (i, pfd) in pollfds[..fds.len()].iter().enumerate() {
if pfd.revents().is_some_and(|r| !r.is_empty()) {
ready[i] = true;
n += 1;
}
}
return Ok(n);
}
Err(Errno::EINTR) => continue,
Err(e) => return Err(std::io::Error::from(e)),
}
}
}
fn poll_timeout(timeout: Option<Duration>) -> PollTimeout {
match timeout {
None => PollTimeout::NONE,
Some(d) => {
let ms = d.as_millis() + u128::from(d.subsec_nanos() % 1_000_000 != 0);
u16::try_from(ms)
.map(PollTimeout::from)
.unwrap_or(PollTimeout::MAX)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::net::UdpSocket;
use std::os::fd::AsRawFd;
use std::time::Instant;
#[test]
fn sub_millisecond_timeout_rounds_up_and_still_times_out() {
let sock = UdpSocket::bind("127.0.0.1:0").unwrap();
let fd = sock.as_raw_fd();
let start = Instant::now();
let ready = wait_readable(fd, Some(Duration::from_micros(500))).unwrap();
let elapsed = start.elapsed();
assert!(!ready);
assert!(
elapsed >= Duration::from_millis(1),
"expected at least 1ms, got {elapsed:?}"
);
}
#[test]
fn zero_timeout_returns_immediately() {
let sock = UdpSocket::bind("127.0.0.1:0").unwrap();
let fd = sock.as_raw_fd();
let ready = wait_readable(fd, Some(Duration::ZERO)).unwrap();
assert!(!ready);
}
#[test]
fn wait_any_readable_finds_the_ready_fd_among_several() {
let idle = UdpSocket::bind("127.0.0.1:0").unwrap();
let receiver = UdpSocket::bind("127.0.0.1:0").unwrap();
let sender = UdpSocket::bind("127.0.0.1:0").unwrap();
sender
.send_to(b"hi", receiver.local_addr().unwrap())
.unwrap();
let fds = [idle.as_raw_fd(), receiver.as_raw_fd()];
let ready = wait_any_readable(&fds, Some(Duration::from_millis(200))).unwrap();
assert!(ready);
}
#[test]
fn poll_readable_into_marks_only_the_ready_fd() {
let idle = UdpSocket::bind("127.0.0.1:0").unwrap();
let receiver = UdpSocket::bind("127.0.0.1:0").unwrap();
let sender = UdpSocket::bind("127.0.0.1:0").unwrap();
sender
.send_to(b"hi", receiver.local_addr().unwrap())
.unwrap();
let fds = [receiver.as_raw_fd(), idle.as_raw_fd()];
let mut ready = [false; 2];
let n = poll_readable_into(&fds, &mut ready, Some(Duration::from_millis(200))).unwrap();
assert_eq!(n, 1);
assert_eq!(ready, [true, false]);
}
#[test]
fn poll_readable_into_times_out_with_nothing_ready() {
let a = UdpSocket::bind("127.0.0.1:0").unwrap();
let b = UdpSocket::bind("127.0.0.1:0").unwrap();
let fds = [a.as_raw_fd(), b.as_raw_fd()];
let mut ready = [true; 2];
let n = poll_readable_into(&fds, &mut ready, Some(Duration::ZERO)).unwrap();
assert_eq!(n, 0);
assert_eq!(ready, [false, false]);
}
#[test]
fn wait_any_readable_times_out_when_none_ready() {
let a = UdpSocket::bind("127.0.0.1:0").unwrap();
let b = UdpSocket::bind("127.0.0.1:0").unwrap();
let fds = [a.as_raw_fd(), b.as_raw_fd()];
let ready = wait_any_readable(&fds, Some(Duration::ZERO)).unwrap();
assert!(!ready);
}
}