use crate::{
Buf, Error, IoBufMut, IoBufs,
iouring::{
operation::Operation,
request::{ConnectRequest, PollRequest, RecvRequest, Request, RequestOutput, SendRequest},
sockaddr::SockAddr,
},
};
use std::{
net::{SocketAddr, TcpListener},
os::fd::{AsRawFd, FromRawFd, OwnedFd},
sync::Arc,
time::{Duration, Instant},
};
use tracing::warn;
const DEFAULT_READ_BUFFER_SIZE: usize = 64 * 1024;
#[derive(Clone, Debug)]
pub struct Config {
pub tcp_nodelay: Option<bool>,
pub zero_linger: bool,
pub connect_timeout: Duration,
pub read_write_timeout: Duration,
pub read_buffer_size: usize,
}
impl Default for Config {
fn default() -> Self {
Self {
tcp_nodelay: Some(true),
zero_linger: true,
connect_timeout: Duration::from_secs(10),
read_write_timeout: Duration::from_secs(60),
read_buffer_size: DEFAULT_READ_BUFFER_SIZE,
}
}
}
#[derive(Clone)]
pub struct Network {
tcp_nodelay: Option<bool>,
zero_linger: bool,
connect_timeout: Duration,
read_write_timeout: Duration,
read_buffer_size: usize,
}
impl Network {
pub(crate) const fn new(cfg: Config) -> Self {
Self {
tcp_nodelay: cfg.tcp_nodelay,
zero_linger: cfg.zero_linger,
connect_timeout: cfg.connect_timeout,
read_write_timeout: cfg.read_write_timeout,
read_buffer_size: cfg.read_buffer_size,
}
}
}
fn configure_socket(fd: &OwnedFd, tcp_nodelay: Option<bool>, zero_linger: bool) {
if let Some(enabled) = tcp_nodelay {
let value: libc::c_int = enabled.into();
if unsafe {
libc::setsockopt(
fd.as_raw_fd(),
libc::IPPROTO_TCP,
libc::TCP_NODELAY,
std::ptr::from_ref(&value).cast(),
size_of_val(&value) as libc::socklen_t,
)
} == -1
{
warn!(err = ?std::io::Error::last_os_error(), "failed to set TCP_NODELAY");
}
}
if zero_linger {
let value = libc::linger {
l_onoff: 1,
l_linger: 0,
};
if unsafe {
libc::setsockopt(
fd.as_raw_fd(),
libc::SOL_SOCKET,
libc::SO_LINGER,
std::ptr::from_ref(&value).cast(),
size_of_val(&value) as libc::socklen_t,
)
} == -1
{
warn!(err = ?std::io::Error::last_os_error(), "failed to set SO_LINGER");
}
}
}
impl crate::Network for Network {
type Listener = Listener;
async fn bind(&self, socket: SocketAddr) -> Result<Self::Listener, Error> {
let listener = TcpListener::bind(socket).map_err(|_| Error::BindFailed)?;
listener
.set_nonblocking(true)
.map_err(|_| Error::BindFailed)?;
Ok(Listener {
tcp_nodelay: self.tcp_nodelay,
zero_linger: self.zero_linger,
inner: Arc::new(listener),
read_write_timeout: self.read_write_timeout,
read_buffer_size: self.read_buffer_size,
})
}
async fn dial(
&self,
socket: SocketAddr,
) -> Result<(crate::SinkOf<Self>, crate::StreamOf<Self>), Error> {
let deadline = Instant::now() + self.connect_timeout;
let family = if socket.is_ipv4() {
libc::AF_INET
} else {
libc::AF_INET6
};
let raw = unsafe {
libc::socket(
family,
libc::SOCK_STREAM | libc::SOCK_NONBLOCK | libc::SOCK_CLOEXEC,
0,
)
};
if raw < 0 {
return Err(Error::ConnectionFailed);
}
let fd = Arc::new(unsafe { OwnedFd::from_raw_fd(raw) });
let output = Operation::register(Request::Connect(ConnectRequest {
fd: fd.clone(),
address: Box::new(SockAddr::from(socket)),
deadline: Some(deadline),
}))
.await
.map_err(|_| Error::ConnectionFailed)?;
let RequestOutput::Connect(result) = output else {
unreachable!("connect request returned another output kind");
};
result.map_err(|error| match error {
Error::Timeout => Error::Timeout,
_ => Error::ConnectionFailed,
})?;
configure_socket(&fd, self.tcp_nodelay, self.zero_linger);
Ok((
Sink::new(fd.clone(), self.read_write_timeout),
Stream::new(fd, self.read_write_timeout, self.read_buffer_size),
))
}
}
pub struct Listener {
tcp_nodelay: Option<bool>,
zero_linger: bool,
inner: Arc<TcpListener>,
read_write_timeout: Duration,
read_buffer_size: usize,
}
impl crate::Listener for Listener {
type Stream = Stream;
type Sink = Sink;
async fn accept(&mut self) -> Result<(SocketAddr, Self::Sink, Self::Stream), Error> {
let (stream, remote_addr) = loop {
match self.inner.accept() {
Ok(accepted) => break accepted,
Err(error) if error.kind() == std::io::ErrorKind::Interrupted => continue,
Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {}
Err(_) => return Err(Error::ConnectionFailed),
}
let output = Operation::register(Request::Poll(PollRequest {
fd: self.inner.clone(),
deadline: Some(Instant::now() + self.read_write_timeout),
}))
.await
.map_err(|_| Error::ConnectionFailed)?;
let RequestOutput::Poll(result) = output else {
unreachable!("readiness request returned another output kind");
};
match result {
Ok(()) | Err(Error::Timeout) => {}
Err(_) => return Err(Error::ConnectionFailed),
}
};
stream
.set_nonblocking(true)
.map_err(|_| Error::ConnectionFailed)?;
let fd = Arc::new(OwnedFd::from(stream));
configure_socket(&fd, self.tcp_nodelay, self.zero_linger);
Ok((
remote_addr,
Sink::new(fd.clone(), self.read_write_timeout),
Stream::new(fd, self.read_write_timeout, self.read_buffer_size),
))
}
fn local_addr(&self) -> Result<SocketAddr, std::io::Error> {
self.inner.local_addr()
}
}
pub struct Sink {
fd: Arc<OwnedFd>,
timeout: Duration,
state: SinkState,
}
enum SinkState {
Open,
Sending,
Closed,
}
impl Sink {
const fn new(fd: Arc<OwnedFd>, timeout: Duration) -> Self {
Self {
fd,
timeout,
state: SinkState::Open,
}
}
fn close(&mut self) {
if matches!(self.state, SinkState::Closed) {
return;
}
unsafe {
libc::shutdown(self.fd.as_raw_fd(), libc::SHUT_WR);
}
self.state = SinkState::Closed;
}
}
impl Drop for Sink {
fn drop(&mut self) {
self.close();
}
}
impl crate::Sink for Sink {
async fn send(&mut self, bufs: impl Into<IoBufs> + Send) -> Result<(), Error> {
match self.state {
SinkState::Open => {}
SinkState::Sending => {
self.close();
return Err(Error::Closed);
}
SinkState::Closed => return Err(Error::Closed),
}
let bufs = bufs.into();
if !bufs.has_remaining() {
return Ok(());
}
self.state = SinkState::Sending;
let result = match Operation::register(Request::Send(SendRequest {
fd: self.fd.clone(),
write: bufs.into(),
deadline: Some(Instant::now() + self.timeout),
}))
.await
{
Ok(RequestOutput::Send(result)) => result,
Ok(_) => unreachable!("send request returned another output kind"),
Err(_) => Err(Error::SendFailed),
};
if result.is_err() {
self.close();
return result;
}
self.state = SinkState::Open;
Ok(())
}
}
pub struct Stream {
fd: Arc<OwnedFd>,
timeout: Duration,
poisoned: bool,
buffer: IoBufMut,
buffer_pos: usize,
}
impl Stream {
fn new(fd: Arc<OwnedFd>, timeout: Duration, buffer_capacity: usize) -> Self {
Self {
fd,
timeout,
poisoned: false,
buffer: IoBufMut::with_capacity(buffer_capacity),
buffer_pos: 0,
}
}
async fn submit_recv(
&self,
buffer: IoBufMut,
offset: usize,
len: usize,
exact: bool,
deadline: Instant,
) -> Result<(IoBufMut, usize), Error> {
let output = Operation::register(Request::Recv(RecvRequest {
fd: self.fd.clone(),
buf: buffer,
offset,
len: offset + len,
exact,
deadline: Some(deadline),
}))
.await
.map_err(|_| Error::RecvFailed)?;
let RequestOutput::Recv(result) = output else {
unreachable!("recv request returned another output kind");
};
result
.map(|(buf, total)| (buf, total - offset))
.map_err(|(_, error)| error)
}
async fn fill_buffer(&mut self, deadline: Instant) -> Result<(), Error> {
self.buffer_pos = 0;
let buffer = std::mem::take(&mut self.buffer);
let len = buffer.capacity();
let (buffer, read) = self.submit_recv(buffer, 0, len, false, deadline).await?;
self.buffer = buffer;
unsafe { self.buffer.set_len(read) };
Ok(())
}
}
impl crate::Stream for Stream {
async fn recv(&mut self, len: usize) -> Result<IoBufs, Error> {
if self.poisoned {
return Err(Error::Closed);
}
self.poisoned = true;
let result = async {
let mut owned_buf = IoBufMut::with_capacity(len);
unsafe { owned_buf.set_len(len) };
let mut bytes_received = 0;
let deadline = Instant::now() + self.timeout;
while bytes_received < len {
let buffered = self.buffer.len() - self.buffer_pos;
if buffered > 0 {
let to_copy = std::cmp::min(buffered, len - bytes_received);
owned_buf.as_mut()[bytes_received..bytes_received + to_copy].copy_from_slice(
&self.buffer.as_ref()[self.buffer_pos..self.buffer_pos + to_copy],
);
self.buffer_pos += to_copy;
bytes_received += to_copy;
continue;
}
let remaining = len - bytes_received;
let buffer_capacity = self.buffer.capacity();
if buffer_capacity == 0 || remaining >= buffer_capacity {
let (buf, read) = self
.submit_recv(owned_buf, bytes_received, remaining, true, deadline)
.await?;
owned_buf = buf;
bytes_received += read;
} else {
self.fill_buffer(deadline).await?;
}
}
Ok(IoBufs::from(owned_buf.freeze()))
}
.await;
if result.is_ok() {
self.poisoned = false;
}
result
}
fn peek(&self, max_len: usize) -> &[u8] {
let buffered = self.buffer.len() - self.buffer_pos;
let len = std::cmp::min(buffered, max_len);
&self.buffer.as_ref()[self.buffer_pos..self.buffer_pos + len]
}
}
#[cfg(test)]
mod tests {
use super::{Config, Network, Sink, Stream};
use crate::{
Clock as _, Error, IoBuf, IoBufMut, IoBufs, Listener as _, Network as _, Runner as _,
Sink as _, Spawner as _, Stream as _, Supervisor as _, iouring, network::tests,
};
use commonware_macros::{select, test_group};
use futures::FutureExt as _;
use std::{
io::{Read, Write},
net::TcpStream,
os::{
fd::{AsRawFd, OwnedFd},
unix::net::UnixStream,
},
pin::pin,
sync::{
Arc,
atomic::{AtomicBool, Ordering},
},
task::{Context, Poll, Wake, Waker},
time::{Duration, Instant},
};
#[derive(Default)]
struct Notify(AtomicBool);
impl Wake for Notify {
fn wake(self: Arc<Self>) {
self.0.store(true, Ordering::Relaxed);
}
}
fn socket_options(fd: &OwnedFd) -> (bool, Option<Duration>) {
let stream = TcpStream::from(fd.try_clone().unwrap());
let nodelay = stream.nodelay().unwrap();
let mut linger = libc::linger {
l_onoff: 0,
l_linger: 0,
};
let mut len = size_of_val(&linger) as libc::socklen_t;
let result = unsafe {
libc::getsockopt(
fd.as_raw_fd(),
libc::SOL_SOCKET,
libc::SO_LINGER,
std::ptr::from_mut(&mut linger).cast(),
&mut len,
)
};
assert_eq!(result, 0);
assert_eq!(len as usize, size_of_val(&linger));
let linger =
(linger.l_onoff != 0).then(|| Duration::from_secs(linger.l_linger.try_into().unwrap()));
(nodelay, linger)
}
#[test]
fn test_queued_accept_and_empty_operations_without_worker() {
let network = Network::new(Config::default());
let mut listener =
futures::executor::block_on(network.bind("127.0.0.1:0".parse().unwrap())).unwrap();
let peer = TcpStream::connect(listener.local_addr().unwrap()).unwrap();
let (_, mut sink, mut stream) = futures::executor::block_on(listener.accept()).unwrap();
assert!(stream.peek(1).is_empty());
futures::executor::block_on(sink.send(IoBufs::default())).unwrap();
futures::executor::block_on(sink.send(IoBuf::default())).unwrap();
futures::executor::block_on(sink.send(Vec::<u8>::new())).unwrap();
assert!(
futures::executor::block_on(stream.recv(0))
.unwrap()
.is_empty()
);
drop(peer);
}
#[test]
fn test_cancel_accept_before_service_preserves_next_connection() {
iouring::Runner::new(
iouring::Config::default().with_ring_config(iouring::RingConfig {
size: 1,
..Default::default()
}),
)
.start(|_| async {
let network = Network::new(Config::default());
let mut listener = network.bind("127.0.0.1:0".parse().unwrap()).await.unwrap();
let address = listener.local_addr().unwrap();
assert!(futures::FutureExt::now_or_never(listener.accept()).is_none());
let (mut sender, _receiver) = network.dial(address).await.unwrap();
let (_, _sender, mut receiver) = listener.accept().await.unwrap();
sender.send(b"kept").await.unwrap();
assert_eq!(receiver.recv(4).await.unwrap().coalesce(), b"kept"[..]);
});
}
#[test]
fn test_pending_network_operations_observe_worker_closure() {
let (send, recv, peer) = iouring::Runner::default().start(|_| async {
let (socket, peer) = UnixStream::pair().unwrap();
let fd: Arc<OwnedFd> = Arc::new(socket.into());
let mut sink = Sink::new(fd.clone(), Duration::from_secs(1));
let mut stream = Stream::new(fd, Duration::from_secs(1), 0);
let mut send = Box::pin(async move { sink.send(b"x").await });
let mut recv = Box::pin(async move { stream.recv(1).await });
assert!(futures::poll!(send.as_mut()).is_pending());
assert!(futures::poll!(recv.as_mut()).is_pending());
(send, recv, peer)
});
assert!(matches!(
futures::executor::block_on(send),
Err(Error::SendFailed)
));
assert!(matches!(
futures::executor::block_on(recv),
Err(Error::RecvFailed)
));
drop(peer);
}
#[test]
fn test_trait() {
iouring::Runner::default().start(|context| async move {
tests::test_network_trait(context, || {
Network::new(Config {
read_write_timeout: Duration::from_secs(15),
..Default::default()
})
})
.await;
});
}
#[test]
fn test_connect_timeout() {
iouring::Runner::default().start(|context| async move {
let connect_timeout = Duration::from_millis(100);
let network = Network::new(Config {
connect_timeout,
..Default::default()
});
tests::test_network_connect_timeout(context, network, connect_timeout).await;
});
}
#[test_group("slow")]
#[test]
fn test_stress_trait() {
iouring::Runner::new(
iouring::Config::default().with_ring_config(iouring::RingConfig {
size: 256,
..Default::default()
}),
)
.start(|context| async move {
tests::stress_test_network_trait(context, || Network::new(Config::default())).await;
});
}
#[test]
fn test_read_timeout_with_partial_data() {
iouring::Runner::default().start(|context| async move {
let op_timeout = Duration::from_millis(100);
let network = Network::new(Config {
read_write_timeout: op_timeout,
..Default::default()
});
let mut listener = network.bind("127.0.0.1:0".parse().unwrap()).await.unwrap();
let addr = listener.local_addr().unwrap();
let reader = context.child("reader").spawn(move |_| async move {
let (_addr, _sink, mut stream) = listener.accept().await.unwrap();
let start = Instant::now();
let result = stream.recv(100).await;
let elapsed = start.elapsed();
assert!(stream.peek(100).is_empty());
assert!(matches!(stream.recv(1).await, Err(Error::Closed)));
(result, elapsed)
});
let (mut sink, _stream) = network.dial(addr).await.unwrap();
sink.send([1u8, 2, 3, 4, 5].as_slice()).await.unwrap();
let (result, elapsed) = reader.await.unwrap();
assert!(matches!(result, Err(Error::Timeout)));
assert!(elapsed >= op_timeout);
assert!(elapsed < op_timeout * 3);
});
}
#[test]
fn test_carried_receive_deadline_completes_without_rescheduling() {
iouring::Runner::default().start(|context| async move {
let timeout = Duration::from_secs(1);
let (socket, mut peer) = UnixStream::pair().unwrap();
socket.set_nonblocking(true).unwrap();
let mut observer = socket.try_clone().unwrap();
let mut stream = Stream::new(Arc::new(socket.into()), timeout, 8);
let notified = Arc::new(Notify::default());
let waker = Waker::from(notified.clone());
let mut cx = Context::from_waker(&waker);
peer.write_all(b"x").unwrap();
{
let deadline = Instant::now() + timeout;
let mut receive = pin!(stream.recv(2));
assert!(receive.poll_unpin(&mut cx).is_pending());
while !notified.0.load(Ordering::Relaxed) {
assert!(Instant::now() < deadline, "first refill did not complete");
context.sleep(Duration::from_millis(1)).await;
}
assert!(Instant::now() < deadline);
assert_eq!(
observer.read(&mut [0]).unwrap_err().kind(),
std::io::ErrorKind::WouldBlock
);
context.sleep(timeout).await;
peer.write_all(b"y").unwrap();
assert!(matches!(
receive.poll_unpin(&mut cx),
Poll::Ready(Err(Error::Timeout))
));
}
context.sleep(Duration::from_millis(1)).await;
let mut remaining = [0];
assert_eq!(observer.read(&mut remaining).unwrap(), 1);
assert_eq!(&remaining, b"y");
});
}
#[test]
fn test_unbuffered_mode() {
iouring::Runner::default().start(|context| async move {
let network = Network::new(Config {
read_buffer_size: 0,
..Default::default()
});
let mut listener = network.bind("127.0.0.1:0".parse().unwrap()).await.unwrap();
let addr = listener.local_addr().unwrap();
let reader = context.child("reader").spawn(move |_| async move {
let (_addr, _sink, mut stream) = listener.accept().await.unwrap();
assert!(stream.peek(100).is_empty());
let buf1 = stream.recv(5).await.unwrap();
assert!(stream.peek(100).is_empty());
let buf2 = stream.recv(5).await.unwrap();
assert!(stream.peek(100).is_empty());
(buf1, buf2)
});
let (mut sink, _stream) = network.dial(addr).await.unwrap();
sink.send([1u8, 2, 3, 4, 5].as_slice()).await.unwrap();
sink.send([6u8, 7, 8, 9, 10].as_slice()).await.unwrap();
let (buf1, buf2) = reader.await.unwrap();
assert_eq!(buf1.coalesce(), &[1u8, 2, 3, 4, 5]);
assert_eq!(buf2.coalesce(), &[6u8, 7, 8, 9, 10]);
});
}
#[test]
fn test_cancelled_recv_retains_descriptor_until_completion() {
iouring::Runner::default().start(|context| async move {
let op_timeout = Duration::from_millis(200);
let network = Network::new(Config {
read_write_timeout: op_timeout,
..Default::default()
});
let mut listener = network.bind("127.0.0.1:0".parse().unwrap()).await.unwrap();
let addr = listener.local_addr().unwrap();
let (client_sink, mut client_stream) = network.dial(addr).await.unwrap();
let (_addr, _server_sink, _server_stream) = listener.accept().await.unwrap();
let fd = client_stream.fd.clone();
assert_eq!(Arc::strong_count(&fd), 3);
select! {
_ = client_stream.recv(1) => unreachable!("no data was sent"),
_ = context.sleep(Duration::from_millis(50)) => {},
}
assert_eq!(Arc::strong_count(&fd), 4);
drop(client_sink);
drop(client_stream);
assert_eq!(Arc::strong_count(&fd), 2);
context.sleep(op_timeout).await;
assert_eq!(Arc::strong_count(&fd), 1);
});
}
#[test]
fn test_peek_with_buffered_data() {
let (socket, mut peer) = UnixStream::pair().unwrap();
let mut stream = Stream::new(Arc::new(socket.into()), Duration::from_secs(1), 64);
assert!(stream.peek(100).is_empty());
peer.write_all(b"hello world").unwrap();
let mut stream = iouring::Runner::default().start(|_| async move {
let first = stream.recv(5).await.unwrap();
assert_eq!(first.coalesce(), b"hello");
assert_eq!(stream.peek(100), b" world");
assert_eq!(stream.peek(100), b" world");
assert_eq!(stream.peek(3), b" wo");
assert!(stream.peek(0).is_empty());
peer.write_all(b"xy").unwrap();
assert_eq!(stream.recv(7).await.unwrap().coalesce(), b" worldx");
assert_eq!(stream.peek(100), b"y");
let direct = [b'z'; 64];
peer.write_all(&direct).unwrap();
let received = stream.recv(65).await.unwrap().coalesce();
assert_eq!(&received.as_ref()[..1], b"y");
assert_eq!(&received.as_ref()[1..], &direct);
assert!(stream.peek(100).is_empty());
peer.write_all(b"next").unwrap();
assert_eq!(stream.recv(2).await.unwrap().coalesce(), b"ne");
assert_eq!(stream.peek(100), b"xt");
stream
});
let rest = futures::executor::block_on(stream.recv(2)).unwrap();
assert_eq!(rest.coalesce(), b"xt");
assert!(stream.peek(100).is_empty());
}
#[test]
fn test_submit_recv_returns_bytes_for_this_call() {
iouring::Runner::default().start(|_| async {
let (left, mut right) = UnixStream::pair().unwrap();
let stream = Stream::new(Arc::new(left.into()), Duration::from_secs(1), 0);
right.write_all(b"abc").unwrap();
let buffer = IoBufMut::from(b"xy___");
let result = stream
.submit_recv(buffer, 2, 3, true, Instant::now() + Duration::from_secs(1))
.await;
let (buffer, read) = result.expect("submit_recv should succeed");
assert_eq!(read, 3);
assert_eq!(buffer.as_ref(), b"xyabc");
});
}
#[test]
fn test_large_recv_skips_internal_buffer() {
iouring::Runner::default().start(|context| async move {
let network = Network::new(Config {
read_buffer_size: 8,
..Default::default()
});
let mut listener = network.bind("127.0.0.1:0".parse().unwrap()).await.unwrap();
let addr = listener.local_addr().unwrap();
let expected = b"abcdefgh";
let reader = context.child("reader").spawn(move |_| async move {
let (_addr, _sink, mut stream) = listener.accept().await.unwrap();
let received = stream.recv(expected.len()).await.unwrap();
assert!(stream.peek(1).is_empty());
received
});
let (mut sink, _stream) = network.dial(addr).await.unwrap();
sink.send(expected).await.unwrap();
assert_eq!(reader.await.unwrap().coalesce(), expected);
});
}
#[test]
fn test_socket_options_on_accept_and_dial() {
iouring::Runner::default().start(|_| async {
for (tcp_nodelay, zero_linger) in
[(Some(true), true), (Some(false), false), (None, false)]
{
let network = Network::new(Config {
tcp_nodelay,
zero_linger,
..Default::default()
});
let mut listener = network.bind("127.0.0.1:0".parse().unwrap()).await.unwrap();
let address = listener.local_addr().unwrap();
let (client_sink, _client_stream) = network.dial(address).await.unwrap();
let (_, server_sink, _server_stream) = listener.accept().await.unwrap();
let expected = (
tcp_nodelay.unwrap_or(false),
zero_linger.then_some(Duration::ZERO),
);
for fd in [&client_sink.fd, &server_sink.fd] {
assert_eq!(socket_options(fd), expected);
}
}
});
}
}