This commit is contained in:
Will Temple
2026-05-15 16:28:06 -07:00
parent 67400f1499
commit bfda24ad0a
13 changed files with 782 additions and 363 deletions

View File

@@ -23,11 +23,24 @@ type ShutdownFuture = Pin<Box<dyn Future<Output = io::Result<()>> + 'static>>;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ExecutionPath {
Kqueue,
Offload,
}
pub fn execution_path(_op: &NetOp) -> ExecutionPath {
ExecutionPath::Offload
pub fn execution_path(op: &NetOp) -> ExecutionPath {
match op {
NetOp::Socket { .. }
| NetOp::Connect { .. }
| NetOp::Bind { .. }
| NetOp::Listen { .. }
| NetOp::Accept { .. }
| NetOp::Send { .. }
| NetOp::SendTo { .. }
| NetOp::Recv { .. }
| NetOp::RecvFrom { .. }
| NetOp::Shutdown { .. }
| NetOp::Close { .. } => ExecutionPath::Kqueue,
}
}
pub async fn resolve_addrs<A>(addr: A) -> io::Result<Vec<SocketAddr>>
@@ -59,7 +72,7 @@ pub async fn socket(op: NetOp) -> io::Result<OwnedFd> {
unreachable!("socket backend called with non-socket op");
};
offload(move || socket_sync(domain, socket_type, protocol, flags)).await
socket_sync(domain, socket_type, protocol, flags)
}
pub async fn connect(op: NetOp) -> io::Result<()> {
@@ -67,7 +80,7 @@ pub async fn connect(op: NetOp) -> io::Result<()> {
unreachable!("connect backend called with non-connect op");
};
offload(move || connect_sync(fd, RawSocketAddr::from_socket_addr(addr))).await
connect_async(fd, RawSocketAddr::from_socket_addr(addr)).await
}
pub async fn bind(op: NetOp) -> io::Result<()> {
@@ -75,7 +88,7 @@ pub async fn bind(op: NetOp) -> io::Result<()> {
unreachable!("bind backend called with non-bind op");
};
offload(move || bind_sync(fd, RawSocketAddr::from_socket_addr(addr))).await
bind_sync(fd, RawSocketAddr::from_socket_addr(addr))
}
pub async fn listen(op: NetOp) -> io::Result<()> {
@@ -83,7 +96,7 @@ pub async fn listen(op: NetOp) -> io::Result<()> {
unreachable!("listen backend called with non-listen op");
};
offload(move || listen_sync(fd, backlog)).await
listen_sync(fd, backlog)
}
pub async fn accept(op: NetOp) -> io::Result<AcceptedSocket> {
@@ -91,7 +104,7 @@ pub async fn accept(op: NetOp) -> io::Result<AcceptedSocket> {
unreachable!("accept backend called with non-accept op");
};
offload(move || accept_sync(fd)).await
accept_async(fd).await
}
pub async fn send(op: NetOp) -> io::Result<usize> {
@@ -99,7 +112,7 @@ pub async fn send(op: NetOp) -> io::Result<usize> {
unreachable!("send backend called with non-send op");
};
offload(move || send_sync(fd, data, flags)).await
send_async(fd, data, flags).await
}
pub async fn send_to(op: NetOp) -> io::Result<usize> {
@@ -113,7 +126,7 @@ pub async fn send_to(op: NetOp) -> io::Result<usize> {
unreachable!("send_to backend called with non-send_to op");
};
offload(move || send_to_sync(fd, target, data, flags)).await
send_to_async(fd, target, data, flags).await
}
pub async fn recv(op: NetOp) -> io::Result<Vec<u8>> {
@@ -121,7 +134,7 @@ pub async fn recv(op: NetOp) -> io::Result<Vec<u8>> {
unreachable!("recv backend called with non-recv op");
};
offload(move || recv_sync(fd, len, flags)).await
recv_async(fd, len, flags).await
}
pub async fn recv_from(op: NetOp) -> io::Result<ReceivedDatagram> {
@@ -129,7 +142,7 @@ pub async fn recv_from(op: NetOp) -> io::Result<ReceivedDatagram> {
unreachable!("recv_from backend called with non-recv_from op");
};
offload(move || recv_from_sync(fd, len, flags)).await
recv_from_async(fd, len, flags).await
}
pub async fn shutdown(op: NetOp) -> io::Result<()> {
@@ -137,7 +150,7 @@ pub async fn shutdown(op: NetOp) -> io::Result<()> {
unreachable!("shutdown backend called with non-shutdown op");
};
offload(move || shutdown_sync(fd, how)).await
shutdown_sync(fd, how)
}
pub async fn close(op: NetOp) -> io::Result<()> {
@@ -145,7 +158,7 @@ pub async fn close(op: NetOp) -> io::Result<()> {
unreachable!("close backend called with non-close op");
};
offload(move || close_sync(fd)).await
close_sync(fd)
}
pub async fn connect_stream(addr: SocketAddr) -> io::Result<OwnedFd> {
@@ -234,11 +247,9 @@ async fn bind_datagram_inner(addr: SocketAddr) -> io::Result<OwnedFd> {
}
pub async fn duplicate(fd: RawFd) -> io::Result<OwnedFd> {
offload(move || {
let duplicated = cvt(unsafe { libc::fcntl(fd, libc::F_DUPFD_CLOEXEC, 0) })?;
Ok(unsafe { OwnedFd::from_raw_fd(duplicated) })
})
.await
let duplicated = cvt(unsafe { libc::fcntl(fd, libc::F_DUPFD_CLOEXEC, 0) })?;
set_nonblocking(duplicated)?;
Ok(unsafe { OwnedFd::from_raw_fd(duplicated) })
}
pub async fn recv_timeout(
@@ -247,11 +258,7 @@ pub async fn recv_timeout(
flags: i32,
timeout: Duration,
) -> io::Result<Vec<u8>> {
offload(move || {
wait_for_fd(fd, libc::POLLIN, timeout)?;
recv_sync(fd, len, flags)
})
.await
io_timeout(timeout, recv_async(fd, len, flags)).await
}
pub async fn send_timeout(
@@ -260,11 +267,7 @@ pub async fn send_timeout(
flags: i32,
timeout: Duration,
) -> io::Result<usize> {
offload(move || {
wait_for_fd(fd, libc::POLLOUT, timeout)?;
send_sync(fd, data, flags)
})
.await
io_timeout(timeout, send_async(fd, data, flags)).await
}
pub async fn recv_from_timeout(
@@ -273,11 +276,7 @@ pub async fn recv_from_timeout(
flags: i32,
timeout: Duration,
) -> io::Result<ReceivedDatagram> {
offload(move || {
wait_for_fd(fd, libc::POLLIN, timeout)?;
recv_from_sync(fd, len, flags)
})
.await
io_timeout(timeout, recv_from_async(fd, len, flags)).await
}
pub async fn send_to_timeout(
@@ -287,15 +286,21 @@ pub async fn send_to_timeout(
flags: i32,
timeout: Duration,
) -> io::Result<usize> {
offload(move || {
wait_for_fd(fd, libc::POLLOUT, timeout)?;
send_to_sync(fd, target, data, flags)
})
.await
io_timeout(timeout, send_to_async(fd, target, data, flags)).await
}
pub async fn connect_stream_timeout(addr: SocketAddr, timeout: Duration) -> io::Result<OwnedFd> {
offload(move || connect_stream_timeout_sync(addr, timeout)).await
let fd = socket_sync(socket_domain(addr), libc::SOCK_STREAM, 0, 0)?;
if let Err(error) = io_timeout(
timeout,
connect_async(fd.as_raw_fd(), RawSocketAddr::from_socket_addr(addr)),
)
.await
{
drop(fd);
return Err(error);
}
Ok(fd)
}
pub fn local_addr(fd: RawFd) -> io::Result<SocketAddr> {
@@ -396,6 +401,15 @@ async fn offload<T: Send + 'static>(
future.await
}
async fn io_timeout<T>(
timeout: Duration,
future: impl Future<Output = io::Result<T>>,
) -> io::Result<T> {
crate::time::timeout(timeout, future)
.await
.map_err(|_| io::Error::new(io::ErrorKind::TimedOut, "operation timed out"))?
}
fn socket_domain(addr: SocketAddr) -> i32 {
match addr {
SocketAddr::V4(_) => libc::AF_INET,
@@ -573,11 +587,28 @@ impl RawSocketAddr {
fn socket_sync(domain: i32, socket_type: i32, protocol: i32, _flags: u32) -> io::Result<OwnedFd> {
let fd = cvt(unsafe { libc::socket(domain, socket_type, protocol) })?;
set_cloexec(fd)?;
set_nonblocking(fd)?;
Ok(unsafe { OwnedFd::from_raw_fd(fd) })
}
fn connect_sync(fd: RawFd, addr: RawSocketAddr) -> io::Result<()> {
cvt(unsafe { libc::connect(fd, addr.as_ptr(), addr.len()) }).map(|_| ())
async fn connect_async(fd: RawFd, addr: RawSocketAddr) -> io::Result<()> {
loop {
let result = unsafe { libc::connect(fd, addr.as_ptr(), addr.len()) };
if result == 0 {
return Ok(());
}
let error = io::Error::last_os_error();
match error.raw_os_error() {
Some(libc::EINTR) => continue,
Some(libc::EINPROGRESS) | Some(libc::EALREADY) => {
crate::sys::current::fd::wait_writable(fd).await?;
return socket_error(fd);
}
Some(libc::EISCONN) => return Ok(()),
_ => return Err(error),
}
}
}
fn bind_sync(fd: RawFd, addr: RawSocketAddr) -> io::Result<()> {
@@ -602,12 +633,59 @@ fn accept_sync(fd: RawFd) -> io::Result<AcceptedSocket> {
})
}
fn send_sync(fd: RawFd, data: Vec<u8>, flags: i32) -> io::Result<usize> {
async fn accept_async(fd: RawFd) -> io::Result<AcceptedSocket> {
loop {
match accept_sync(fd) {
Ok(socket) => {
set_nonblocking(socket.fd)?;
return Ok(socket);
}
Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
crate::sys::current::fd::wait_readable(fd).await?;
}
Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
Err(error) => return Err(error),
}
}
}
async fn send_async(fd: RawFd, data: Vec<u8>, flags: i32) -> io::Result<usize> {
loop {
match send_slice_sync(fd, &data, flags) {
Ok(written) => return Ok(written),
Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
crate::sys::current::fd::wait_writable(fd).await?;
}
Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
Err(error) => return Err(error),
}
}
}
fn send_slice_sync(fd: RawFd, data: &[u8], flags: i32) -> io::Result<usize> {
let written = unsafe { libc::send(fd, data.as_ptr().cast::<c_void>(), data.len(), flags) };
cvt_long(written).map(|written| written as usize)
}
fn send_to_sync(fd: RawFd, target: SocketAddr, data: Vec<u8>, flags: i32) -> io::Result<usize> {
async fn send_to_async(
fd: RawFd,
target: SocketAddr,
data: Vec<u8>,
flags: i32,
) -> io::Result<usize> {
loop {
match send_to_slice_sync(fd, target, &data, flags) {
Ok(written) => return Ok(written),
Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
crate::sys::current::fd::wait_writable(fd).await?;
}
Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
Err(error) => return Err(error),
}
}
}
fn send_to_slice_sync(fd: RawFd, target: SocketAddr, data: &[u8], flags: i32) -> io::Result<usize> {
let addr = RawSocketAddr::from_socket_addr(target);
let written = unsafe {
libc::sendto(
@@ -622,6 +700,19 @@ fn send_to_sync(fd: RawFd, target: SocketAddr, data: Vec<u8>, flags: i32) -> io:
cvt_long(written).map(|written| written as usize)
}
async fn recv_async(fd: RawFd, len: usize, flags: i32) -> io::Result<Vec<u8>> {
loop {
match recv_sync(fd, len, flags) {
Ok(data) => return Ok(data),
Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
crate::sys::current::fd::wait_readable(fd).await?;
}
Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
Err(error) => return Err(error),
}
}
}
fn recv_sync(fd: RawFd, len: usize, flags: i32) -> io::Result<Vec<u8>> {
let mut data = vec![0u8; len];
let read = unsafe { libc::recv(fd, data.as_mut_ptr().cast::<c_void>(), len, flags) };
@@ -630,6 +721,19 @@ fn recv_sync(fd: RawFd, len: usize, flags: i32) -> io::Result<Vec<u8>> {
Ok(data)
}
async fn recv_from_async(fd: RawFd, len: usize, flags: i32) -> io::Result<ReceivedDatagram> {
loop {
match recv_from_sync(fd, len, flags) {
Ok(datagram) => return Ok(datagram),
Err(error) if error.kind() == io::ErrorKind::WouldBlock => {
crate::sys::current::fd::wait_readable(fd).await?;
}
Err(error) if error.kind() == io::ErrorKind::Interrupted => {}
Err(error) => return Err(error),
}
}
}
fn recv_from_sync(fd: RawFd, len: usize, flags: i32) -> io::Result<ReceivedDatagram> {
let mut data = vec![0u8; len];
let mut storage = MaybeUninit::<libc::sockaddr_storage>::zeroed();
@@ -659,79 +763,6 @@ fn close_sync(fd: RawFd) -> io::Result<()> {
cvt(unsafe { libc::close(fd) }).map(|_| ())
}
fn connect_stream_timeout_sync(addr: SocketAddr, timeout: Duration) -> io::Result<OwnedFd> {
let fd = cvt(unsafe { libc::socket(socket_domain(addr), libc::SOCK_STREAM, 0) })?;
set_cloexec(fd)?;
set_nonblocking(fd)?;
let raw = RawSocketAddr::from_socket_addr(addr);
let connect_result = unsafe { libc::connect(fd, raw.as_ptr(), raw.len()) };
if connect_result < 0 {
let error = io::Error::last_os_error();
if error.raw_os_error() != Some(libc::EINPROGRESS) {
let _ = unsafe { libc::close(fd) };
return Err(error);
}
}
let wait = wait_for_fd(fd, libc::POLLOUT, timeout);
if let Err(error) = wait {
let _ = unsafe { libc::close(fd) };
return Err(error);
}
let mut so_error: libc::c_int = 0;
let mut len = std::mem::size_of::<libc::c_int>() as libc::socklen_t;
if unsafe {
libc::getsockopt(
fd,
libc::SOL_SOCKET,
libc::SO_ERROR,
&mut so_error as *mut libc::c_int as *mut c_void,
&mut len,
)
} < 0
{
let error = io::Error::last_os_error();
let _ = unsafe { libc::close(fd) };
return Err(error);
}
if so_error != 0 {
let _ = unsafe { libc::close(fd) };
return Err(io::Error::from_raw_os_error(so_error));
}
clear_nonblocking(fd)?;
Ok(unsafe { OwnedFd::from_raw_fd(fd) })
}
fn wait_for_fd(fd: RawFd, events: i16, timeout: Duration) -> io::Result<()> {
let timeout_ms = timeout.as_millis().min(i32::MAX as u128) as i32;
let mut pfd = libc::pollfd {
fd,
events,
revents: 0,
};
loop {
let result = unsafe { libc::poll(&mut pfd, 1, timeout_ms) };
if result > 0 {
return Ok(());
}
if result == 0 {
return Err(io::Error::new(
io::ErrorKind::TimedOut,
"operation timed out",
));
}
let error = io::Error::last_os_error();
if error.kind() == io::ErrorKind::Interrupted {
continue;
}
return Err(error);
}
}
fn set_cloexec(fd: RawFd) -> io::Result<()> {
let flags = cvt(unsafe { libc::fcntl(fd, libc::F_GETFD) })?;
cvt(unsafe { libc::fcntl(fd, libc::F_SETFD, flags | libc::FD_CLOEXEC) })?;
@@ -744,12 +775,6 @@ fn set_nonblocking(fd: RawFd) -> io::Result<()> {
Ok(())
}
fn clear_nonblocking(fd: RawFd) -> io::Result<()> {
let flags = cvt(unsafe { libc::fcntl(fd, libc::F_GETFL) })?;
cvt(unsafe { libc::fcntl(fd, libc::F_SETFL, flags & !libc::O_NONBLOCK) })?;
Ok(())
}
fn should_try_ipv4_loopback(addr: SocketAddr, error: &io::Error) -> bool {
matches!(addr, SocketAddr::V6(v6) if v6.ip().is_loopback())
&& matches!(
@@ -758,6 +783,25 @@ fn should_try_ipv4_loopback(addr: SocketAddr, error: &io::Error) -> bool {
)
}
fn socket_error(fd: RawFd) -> io::Result<()> {
let mut so_error: libc::c_int = 0;
let mut len = std::mem::size_of::<libc::c_int>() as libc::socklen_t;
cvt(unsafe {
libc::getsockopt(
fd,
libc::SOL_SOCKET,
libc::SO_ERROR,
&mut so_error as *mut libc::c_int as *mut c_void,
&mut len,
)
})?;
if so_error == 0 {
Ok(())
} else {
Err(io::Error::from_raw_os_error(so_error))
}
}
fn localhost_v4(addr: SocketAddr) -> SocketAddr {
SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, addr.port()))
}