Skip to content

Instantly share code, notes, and snippets.

@denezure
Last active March 23, 2022 11:35
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save denezure/1ede65f6033610235931bc925729d753 to your computer and use it in GitHub Desktop.
use nix::sys::{
socket::{MsgFlags, RecvMmsgData, RecvMsg},
uio::IoVec,
};
use tokio::net::UdpSocket;
use thiserror::Error;
type CustomResult<T> = Result<T, Box<dyn std::error::Error + Send + Sync>>;
const DatagramSize: usize = 1472
#[derive(Error, Debug)]
enum RecvmmsgError {
#[error("no datagrams ready to be read")]
WouldBlock,
#[error("unknown")]
Unknown,
}
async fn recvmmsg<'a>(
s: &UdpSocket,
buf: &'a mut [u8],
capacity: usize,
) -> CustomResult<Vec<RecvMsg<'a>>> {
let fd = s.as_raw_fd();
let f = move || {
let mut msgs: Vec<_> = buf
.chunks_exact_mut(DatagramSize)
.map(|buf| [IoVec::from_mut_slice(&mut buf[..])])
.map(|iov| RecvMmsgData {
iov,
cmsg_buffer: None,
})
.collect();
match nix::sys::socket::recvmmsg(fd, &mut msgs, MsgFlags::empty(), None) {
Ok(r) => Ok(r),
// nix gives back raw errno values - map them to what tokio expects (std::io::Error)
Err(nix::errno::Errno::EAGAIN) => {
Err(std::io::Error::from(std::io::ErrorKind::WouldBlock))
}
Err(_) => Err(std::io::Error::new(
std::io::ErrorKind::Other,
"unknown error in recvmmsg",
)),
}
};
s.readable().await?;
match s.try_io(tokio::io::Interest::READABLE, f) {
Ok(messages) => return Ok(messages),
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
Err(Box::new(RecvmmsgError::WouldBlock))
}
Err(e) => return Err(Box::new(RecvmmsgError::Unknown)),
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment