use std::collections::HashMap;
use std::net::SocketAddr;
use std::os::unix::io::{AsRawFd, FromRawFd, OwnedFd, RawFd};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc;
use std::thread::JoinHandle;
use std::time::{Duration, Instant};
use io_uring::{IoUring, opcode, types};
use tracing::{debug, error};
use crate::ControlEvent;
use crate::buf_ring::BufRing;
use crate::halt::{HaltGate, RefusalSender};
use crate::policy_control::{AckPolicyPublisher, BootedClaim};
use melin_app::Application;
use melin_app::auth::{ClientRole, RoleId};
use melin_app::decoder::ErasedDecoder;
pub type RequestDecoderArc<A> = Arc<dyn ErasedDecoder<<A as Application>::Event>>;
use melin_app::unix_epoch_nanos;
use melin_pipeline::ring;
use melin_transport_core::pipeline::InputSlot;
use melin_transport_core::tick::TickSchedule;
const BUF_SIZE: usize = 4096;
const BUF_GROUP_ID: u16 = 0;
use crate::client_frames::MAX_FRAME_SIZE;
use crate::connection_limit::{RingSizing, await_startup, ring_setup_error};
const EVENTFD_TOKEN: u64 = u64::MAX;
const PROVIDE_BUFS_TOKEN: u64 = u64::MAX - 1;
const CANCEL_TOKEN: u64 = u64::MAX - 3;
const TICK_TIMEOUT_TOKEN: u64 = u64::MAX - 2;
const IORING_CQE_F_BUFFER: u32 = 1 << 0;
const IORING_CQE_F_MORE: u32 = 1 << 1;
const IORING_CQE_BUFFER_SHIFT: u32 = 16;
use melin_wire_protocol::control::ConnectionId;
pub struct ReaderRegistration<R> {
pub connection_id: ConnectionId,
pub reader: R,
pub addr: SocketAddr,
pub role: ClientRole<RoleId>,
pub key_hash: u64,
}
pub struct UringReaderHandle<R> {
tx: mpsc::Sender<ReaderRegistration<R>>,
event_fd: Arc<OwnedFd>,
join_handle: Option<JoinHandle<()>>,
shutdown: Arc<AtomicBool>,
}
impl<R> UringReaderHandle<R> {
pub fn register(&mut self, registration: ReaderRegistration<R>) {
if self.tx.send(registration).is_ok() {
wake_reader(&self.event_fd);
} else {
error!("reader thread dead, cannot register connection");
self.shutdown.store(true, Ordering::Relaxed);
}
}
pub fn shutdown(&self) {
self.shutdown.store(true, Ordering::Relaxed);
wake_reader(&self.event_fd);
}
pub fn join(mut self) {
if let Some(handle) = self.join_handle.take()
&& let Err(panic) = handle.join()
{
let msg = panic
.downcast_ref::<&str>()
.copied()
.or_else(|| panic.downcast_ref::<String>().map(|s| s.as_str()))
.unwrap_or("<non-string panic>");
error!(message = msg, "reader thread panicked");
}
}
}
pub(crate) fn wake_reader(event_fd: &OwnedFd) {
let val: u64 = 1;
unsafe {
libc::write(
event_fd.as_raw_fd(),
&val as *const u64 as *const libc::c_void,
8,
);
}
}
#[allow(clippy::too_many_arguments)]
pub fn spawn_reader<A: Application, R: AsRawFd + Send + 'static>(
mut producer: ring::Producer<InputSlot<A::Event>>,
decoder: RequestDecoderArc<A>,
halt: HaltGate,
refusals: RefusalSender<A::Report>,
ack_policy: BootedClaim,
control_tx: mpsc::Sender<ControlEvent>,
core: usize,
connection_timeout: Option<Duration>,
tick_cadence: Option<Duration>,
sizing: RingSizing,
shutdown: Arc<AtomicBool>,
) -> std::io::Result<UringReaderHandle<R>>
where
A::Event: Send + Sync + 'static,
A::Report: Send + 'static,
{
let (tx, rx) = mpsc::channel();
let raw_event_fd = unsafe { libc::eventfd(0, libc::EFD_NONBLOCK) };
if raw_event_fd < 0 {
let err = std::io::Error::last_os_error();
return Err(std::io::Error::new(
err.kind(),
format!("uring-reader: cannot create its wakeup eventfd: {err}"),
));
}
let event_fd = Arc::new(unsafe { OwnedFd::from_raw_fd(raw_event_fd) });
let wakeup_fd = Arc::clone(&event_fd);
let shutdown_clone = Arc::clone(&shutdown);
let ack_policy = ack_policy.attach(&mut producer, &halt, Some(Arc::clone(&event_fd)));
let (ready_tx, ready_rx) = mpsc::sync_channel::<std::io::Result<()>>(1);
let handle = melin_app::affinity::spawn_placed("uring-reader", core, move || {
reader_loop::<A, R>(
rx,
wakeup_fd,
producer,
&*decoder,
&halt,
refusals,
ack_policy,
&control_tx,
connection_timeout,
tick_cadence,
sizing,
ready_tx,
&shutdown_clone,
);
})
.map_err(|e| {
let e = std::io::Error::from(e);
std::io::Error::new(
e.kind(),
format!("uring-reader: cannot start its thread: {e}"),
)
})?;
if let Err(e) = await_startup("uring-reader", &ready_rx) {
let _ = handle.join();
return Err(e);
}
Ok(UringReaderHandle {
tx,
event_fd,
join_handle: Some(handle),
shutdown,
})
}
struct ConnectionEntry<R> {
connection_id: u64,
addr: SocketAddr,
role: ClientRole<RoleId>,
key_hash: u64,
_reader: R,
fd: RawFd,
parse_buf: Vec<u8>,
multishot_active: bool,
last_activity: Instant,
dying: bool,
}
struct ConnectionSlab<R> {
entries: Vec<Option<ConnectionEntry<R>>>,
free: Vec<usize>,
}
impl<R> ConnectionSlab<R> {
fn new() -> Self {
Self {
entries: Vec::with_capacity(256),
free: Vec::new(),
}
}
fn insert(&mut self, entry: ConnectionEntry<R>) -> usize {
if let Some(idx) = self.free.pop() {
self.entries[idx] = Some(entry);
idx
} else {
let idx = self.entries.len();
self.entries.push(Some(entry));
idx
}
}
fn get_mut(&mut self, idx: usize) -> Option<&mut ConnectionEntry<R>> {
self.entries.get_mut(idx).and_then(|e| e.as_mut())
}
fn remove(&mut self, idx: usize) -> Option<ConnectionEntry<R>> {
if let Some(slot) = self.entries.get_mut(idx) {
let removed = slot.take();
if removed.is_some() {
self.free.push(idx);
}
removed
} else {
None
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RingEntry {
Skip,
Submit,
SubmitAndWait,
}
#[inline]
fn ring_entry(sq_pending: usize, cq_ready: bool) -> RingEntry {
match (sq_pending > 0, cq_ready) {
(false, true) => RingEntry::Skip,
(true, true) => RingEntry::Submit,
(_, false) => RingEntry::SubmitAndWait,
}
}
#[allow(clippy::too_many_arguments)]
fn reader_loop<A: Application, R: AsRawFd>(
command_rx: mpsc::Receiver<ReaderRegistration<R>>,
wakeup_fd: Arc<OwnedFd>,
mut producer: ring::Producer<InputSlot<A::Event>>,
decoder: &dyn ErasedDecoder<A::Event>,
halt: &HaltGate,
mut refusals: RefusalSender<A::Report>,
mut ack_policy: AckPolicyPublisher,
control_tx: &mpsc::Sender<ControlEvent>,
connection_timeout: Option<Duration>,
tick_cadence: Option<Duration>,
sizing: RingSizing,
ready: mpsc::SyncSender<std::io::Result<()>>,
shutdown: &AtomicBool,
) {
let num_buffers = sizing.provided_buffers();
let ring_entries = sizing.ring_entries();
let mut eventfd_buf: Box<[u8; 8]> = Box::new([0u8; 8]);
let mut buffer_pool = vec![0u8; num_buffers as usize * BUF_SIZE].into_boxed_slice();
let mut buf_ring = BufRing::new(num_buffers, buffer_pool.as_mut_ptr(), BUF_SIZE);
let mut ring: IoUring = match IoUring::builder().setup_single_issuer().build(ring_entries) {
Ok(ring) => ring,
Err(e) => {
let _ = ready.send(Err(ring_setup_error("uring-reader", ring_entries, e)));
return;
}
};
let use_buf_ring = match buf_ring.register(&ring, BUF_GROUP_ID) {
Ok(()) => true,
Err(e) => {
tracing::warn!(
error = %e,
"buf_ring registration rejected — falling back to legacy \
ProvideBuffers recycling (kernel < 5.19, or a hypervisor \
filtering io_uring register opcodes)"
);
if let Err(e) = register_buffer_pool(&mut ring, buffer_pool.as_mut_ptr(), num_buffers) {
let _ = ready.send(Err(std::io::Error::new(
e.kind(),
format!("uring-reader: cannot provide its receive buffers: {e}"),
)));
return;
}
false
}
};
if ready.send(Ok(())).is_err() {
return;
}
let mut slab = ConnectionSlab::<R>::new();
let mut fd_to_slab: HashMap<RawFd, usize> = HashMap::with_capacity(256);
let mut cqes: Vec<(u64, i32, u32)> = Vec::with_capacity(ring_entries as usize * 2);
push_eventfd_read(&mut ring, wakeup_fd.as_raw_fd(), eventfd_buf.as_mut_ptr());
let mut eventfd_armed = true;
#[cfg(feature = "latency-trace")]
let mut publish_rec =
melin_transport_core::trace::register_stage("reader: publish (decode → disruptor publish)");
#[cfg(feature = "tick-to-trade")]
let mut ingest_rec =
melin_transport_core::trace::register_stage("reader: ingest (recv_ts → publish complete)");
#[cfg(feature = "latency-trace")]
let mut last_stats_flush = Instant::now();
let mut last_timeout_scan = Instant::now();
let mut stale: Vec<usize> = Vec::new();
let tick_enabled = tick_cadence.is_some();
let cadence = tick_cadence.unwrap_or(Duration::ZERO);
let mut tick_schedule = TickSchedule::new(cadence, Instant::now());
let mut tick_armed = false;
let mut tick_ts = types::Timespec::new()
.sec(cadence.as_secs())
.nsec(cadence.subsec_nanos());
if tick_enabled {
let sqe = opcode::Timeout::new(&tick_ts)
.build()
.user_data(TICK_TIMEOUT_TOKEN);
unsafe {
ring.submission()
.push(&sqe)
.expect("io_uring SQ full while arming initial tick timeout");
}
tick_armed = true;
tracing::info!(
cadence_ms = cadence.as_millis() as u64,
"tick generator integrated into reader thread"
);
}
loop {
if shutdown.load(Ordering::Relaxed) {
break;
}
ack_policy.poll(&mut producer, halt);
if tick_enabled {
if tick_schedule.publish_if_due(Instant::now(), unix_epoch_nanos, &mut producer) {
tick_armed = false;
}
if !tick_armed {
let remaining = tick_schedule
.next_deadline()
.saturating_duration_since(Instant::now());
tick_ts = types::Timespec::new()
.sec(remaining.as_secs())
.nsec(remaining.subsec_nanos());
let sqe = opcode::Timeout::new(&tick_ts)
.build()
.user_data(TICK_TIMEOUT_TOKEN);
unsafe { push_sqe(&mut ring, &sqe, "the tick TIMEOUT") };
tick_armed = true;
}
}
let sq_pending = ring.submission().len();
let cq_ready = !ring.completion().is_empty();
let entered = match ring_entry(sq_pending, cq_ready) {
RingEntry::Skip => Ok(0),
RingEntry::Submit => ring.submit(),
RingEntry::SubmitAndWait => ring.submit_and_wait(1),
};
match entered {
Ok(_) => {}
Err(ref e) if e.raw_os_error() == Some(libc::EINTR) => continue,
Err(e) => {
error!(error = %e, "io_uring submit/wait error");
break;
}
}
cqes.clear();
cqes.extend(
ring.completion()
.map(|cqe| (cqe.user_data(), cqe.result(), cqe.flags())),
);
let batch_now = Instant::now();
let batch_wall_ns = unix_epoch_nanos();
for &(token, result, flags) in &cqes {
if token == TICK_TIMEOUT_TOKEN {
tick_armed = false;
continue;
}
if token == PROVIDE_BUFS_TOKEN {
if result < 0 {
error!(error = result, "ProvideBuffers failed");
}
continue;
}
if token == CANCEL_TOKEN {
continue;
}
if token == EVENTFD_TOKEN {
if result >= 0 {
while let Ok(reg) = command_rx.try_recv() {
let fd = reg.reader.as_raw_fd();
let entry = ConnectionEntry {
connection_id: reg.connection_id.0,
addr: reg.addr,
role: reg.role,
key_hash: reg.key_hash,
fd,
_reader: reg.reader,
parse_buf: Vec::with_capacity(MAX_FRAME_SIZE + 4),
multishot_active: false,
last_activity: Instant::now(),
dying: false,
};
let idx = slab.insert(entry);
fd_to_slab.insert(fd, idx);
push_recv_multi(&mut ring, &mut slab, idx);
}
} else {
error!(error = result, "eventfd read error");
}
push_eventfd_read(&mut ring, wakeup_fd.as_raw_fd(), eventfd_buf.as_mut_ptr());
continue;
}
let slab_idx = token as usize;
let has_more = (flags & IORING_CQE_F_MORE) != 0;
if !has_more && let Some(entry) = slab.get_mut(slab_idx) {
entry.multishot_active = false;
}
let buf_id = if (flags & IORING_CQE_F_BUFFER) != 0 {
Some((flags >> IORING_CQE_BUFFER_SHIFT) as usize)
} else {
None
};
if slab.get_mut(slab_idx).is_some_and(|e| e.dying) {
if let Some(bid) = buf_id {
recycle_buffer(&mut ring, &mut buf_ring, use_buf_ring, &buffer_pool, bid);
}
if !has_more && let Some(dead) = slab.remove(slab_idx) {
debug!(
connection_id = dead.connection_id,
"teardown complete, slab index released"
);
}
continue;
}
if result == -libc::ENOBUFS {
if let Some(bid) = buf_id {
recycle_buffer(&mut ring, &mut buf_ring, use_buf_ring, &buffer_pool, bid);
}
if has_more {
continue;
}
push_recv_multi(&mut ring, &mut slab, slab_idx);
continue;
}
if result <= 0 {
if let Some(bid) = buf_id {
recycle_buffer(&mut ring, &mut buf_ring, use_buf_ring, &buffer_pool, bid);
}
if let Some(removed) = slab.remove(slab_idx) {
if result == 0 {
debug!(
connection_id = removed.connection_id,
addr = %removed.addr,
"client disconnected"
);
} else {
debug!(
connection_id = removed.connection_id,
addr = %removed.addr,
error = result,
"recv error"
);
}
fd_to_slab.remove(&removed.fd);
let _ = control_tx.send(ControlEvent::Disconnected {
connection_id: removed.connection_id,
});
}
continue;
}
let n = result as usize;
#[allow(clippy::let_unit_value)] let recv_ts = melin_transport_core::trace::mono_trace_ns();
let Some(buf_id) = buf_id else {
debug!(slab_idx, "recv CQE without buffer flag");
if !has_more {
push_recv_multi(&mut ring, &mut slab, slab_idx);
}
continue;
};
let action = if let Some(entry) = slab.get_mut(slab_idx) {
entry.last_activity = batch_now;
let buf_start = buf_id * BUF_SIZE;
entry
.parse_buf
.extend_from_slice(&buffer_pool[buf_start..buf_start + n]);
let drop_conn = process_frames::<A, R>(
entry,
&mut producer,
decoder,
halt,
&mut refusals,
control_tx,
batch_wall_ns,
recv_ts,
#[cfg(feature = "latency-trace")]
&mut publish_rec,
#[cfg(feature = "tick-to-trade")]
&mut ingest_rec,
);
if drop_conn {
Action::Remove
} else if !has_more {
Action::Resubmit
} else {
Action::None
}
} else {
Action::None
};
recycle_buffer(&mut ring, &mut buf_ring, use_buf_ring, &buffer_pool, buf_id);
match action {
Action::Remove => {
begin_teardown(&mut ring, &mut slab, &mut fd_to_slab, control_tx, slab_idx);
}
Action::Resubmit => {
push_recv_multi(&mut ring, &mut slab, slab_idx);
}
Action::None => {}
}
}
if let Some(timeout) = connection_timeout {
let now = Instant::now();
if now.duration_since(last_timeout_scan) >= Duration::from_secs(1) {
last_timeout_scan = now;
stale.clear();
for (idx, slot) in slab.entries.iter().enumerate() {
if let Some(entry) = slot
&& !entry.dying
&& now.duration_since(entry.last_activity) > timeout
{
debug!(
connection_id = entry.connection_id,
addr = %entry.addr,
"connection timed out"
);
stale.push(idx);
}
}
for &idx in &stale {
begin_teardown(&mut ring, &mut slab, &mut fd_to_slab, control_tx, idx);
}
}
}
#[cfg(feature = "latency-trace")]
if batch_now.duration_since(last_stats_flush)
>= melin_transport_core::trace::IDLE_FLUSH_INTERVAL
{
last_stats_flush = batch_now;
publish_rec.flush();
#[cfg(feature = "tick-to-trade")]
ingest_rec.flush();
}
}
let proven = {
let cancel_all = opcode::AsyncCancel2::new(types::CancelBuilder::any())
.build()
.user_data(CANCEL_TOKEN);
unsafe {
let _ = ring.submission().push(&cancel_all);
}
for entry in slab.entries.iter().flatten() {
crate::uring_teardown::wake_pending_ops(entry.fd);
}
let wake: u64 = 1;
unsafe {
libc::write(
wakeup_fd.as_raw_fd(),
&wake as *const u64 as *const libc::c_void,
8,
);
}
loop {
match ring.submit() {
Ok(_) => break,
Err(ref e) if e.raw_os_error() == Some(libc::EINTR) => continue,
Err(_) => break,
}
}
let mut backoff = crate::uring_teardown::DrainBackoff::new();
loop {
for cqe in ring.completion() {
match cqe.user_data() {
EVENTFD_TOKEN => eventfd_armed = false,
PROVIDE_BUFS_TOKEN | CANCEL_TOKEN | TICK_TIMEOUT_TOKEN => {}
idx => {
if (cqe.flags() & IORING_CQE_F_MORE) == 0
&& let Some(entry) = slab.get_mut(idx as usize)
{
entry.multishot_active = false;
}
}
}
}
let owed = eventfd_armed || slab.entries.iter().flatten().any(|e| e.multishot_active);
if !owed {
break true;
}
if !backoff.wait() {
break false;
}
}
};
if !proven {
tracing::warn!(
eventfd_armed,
armed_connections = slab
.entries
.iter()
.flatten()
.filter(|e| e.multishot_active)
.count(),
"io_uring teardown drain did not complete; leaking reader buffers"
);
std::mem::forget(buffer_pool);
std::mem::forget(buf_ring);
std::mem::forget(eventfd_buf);
}
}
enum Action {
Resubmit,
Remove,
None,
}
unsafe fn push_sqe(ring: &mut IoUring, sqe: &io_uring::squeue::Entry, what: &str) {
if unsafe { ring.submission().push(sqe) }.is_ok() {
return;
}
loop {
match ring.submit() {
Ok(_) => break,
Err(ref e) if e.raw_os_error() == Some(libc::EINTR) => continue,
Err(e) => panic!("io_uring submit failed while making room for {what}: {e}"),
}
}
unsafe { ring.submission().push(sqe) }.unwrap_or_else(|_| {
panic!("io_uring SQ still full for {what} after submitting every queued SQE")
});
}
fn register_buffer_pool(
ring: &mut IoUring,
pool_ptr: *mut u8,
num_buffers: u16,
) -> std::io::Result<()> {
let sqe = opcode::ProvideBuffers::new(pool_ptr, BUF_SIZE as i32, num_buffers, BUF_GROUP_ID, 0)
.build()
.user_data(PROVIDE_BUFS_TOKEN);
unsafe {
ring.submission().push(&sqe).map_err(|_| {
std::io::Error::other("io_uring SQ full during buffer pool registration")
})?;
}
loop {
match ring.submit_and_wait(1) {
Ok(_) => break,
Err(ref e) if e.raw_os_error() == Some(libc::EINTR) => continue,
Err(e) => return Err(e),
}
}
let cqe = ring
.completion()
.next()
.ok_or_else(|| std::io::Error::other("no CQE after ProvideBuffers"))?;
if cqe.result() < 0 {
return Err(std::io::Error::from_raw_os_error(-cqe.result()));
}
Ok(())
}
fn re_provide_buffer(ring: &mut IoUring, pool_ptr: *mut u8, buf_id: usize) {
let buf_ptr = unsafe { pool_ptr.add(buf_id * BUF_SIZE) };
let sqe = opcode::ProvideBuffers::new(buf_ptr, BUF_SIZE as i32, 1, BUF_GROUP_ID, buf_id as u16)
.build()
.user_data(PROVIDE_BUFS_TOKEN);
unsafe { push_sqe(ring, &sqe, "a legacy ProvideBuffers recycle") };
}
fn recycle_buffer(
ring: &mut IoUring,
buf_ring: &mut BufRing,
use_buf_ring: bool,
buffer_pool: &[u8],
buf_id: usize,
) {
assert!(
buf_id < buffer_pool.len() / BUF_SIZE,
"kernel returned out-of-pool buffer id {buf_id}"
);
if use_buf_ring {
buf_ring.push(buf_id as u16);
} else {
re_provide_buffer(ring, buffer_pool.as_ptr() as *mut u8, buf_id);
}
}
fn begin_teardown<R>(
ring: &mut IoUring,
slab: &mut ConnectionSlab<R>,
fd_to_slab: &mut HashMap<RawFd, usize>,
control_tx: &mpsc::Sender<ControlEvent>,
idx: usize,
) {
let Some(entry) = slab.get_mut(idx) else {
return;
};
if entry.dying {
return;
}
let connection_id = entry.connection_id;
let fd = entry.fd;
unsafe {
libc::shutdown(fd, libc::SHUT_RDWR);
}
fd_to_slab.remove(&fd);
let _ = control_tx.send(ControlEvent::Disconnected { connection_id });
if entry.multishot_active {
entry.dying = true;
let sqe = opcode::AsyncCancel::new(idx as u64)
.build()
.user_data(CANCEL_TOKEN);
unsafe { push_sqe(ring, &sqe, "a connection teardown's AsyncCancel") };
} else {
slab.remove(idx);
}
}
fn push_recv_multi<R>(ring: &mut IoUring, slab: &mut ConnectionSlab<R>, idx: usize) {
let entry = match slab.get_mut(idx) {
Some(e) => e,
None => return,
};
if entry.multishot_active {
return;
}
let sqe = opcode::RecvMulti::new(types::Fd(entry.fd), BUF_GROUP_ID)
.build()
.user_data(idx as u64);
unsafe { push_sqe(ring, &sqe, "a multishot RECV") };
entry.multishot_active = true;
}
fn push_eventfd_read(ring: &mut IoUring, wakeup_fd: RawFd, buf: *mut u8) {
let sqe = opcode::Read::new(types::Fd(wakeup_fd), buf, 8)
.build()
.user_data(EVENTFD_TOKEN);
unsafe { push_sqe(ring, &sqe, "the eventfd READ") };
}
#[allow(clippy::too_many_arguments)]
fn process_frames<A: Application, R>(
conn: &mut ConnectionEntry<R>,
producer: &mut ring::Producer<InputSlot<A::Event>>,
decoder: &dyn ErasedDecoder<A::Event>,
halt: &HaltGate,
refusals: &mut RefusalSender<A::Report>,
control_tx: &mpsc::Sender<ControlEvent>,
batch_wall_ns: u64,
recv_ts: melin_transport_core::trace::MonoTraceInstant,
#[cfg(feature = "latency-trace")] publish_rec: &mut melin_transport_core::trace::StageRecorder,
#[cfg(feature = "tick-to-trade")] ingest_rec: &mut melin_transport_core::trace::StageRecorder,
) -> bool {
use crate::client_frames::{FrameAction, process_client_frames};
let action = process_client_frames::<A>(
&mut conn.parse_buf,
conn.connection_id,
conn.key_hash,
conn.role,
producer,
decoder,
halt,
refusals,
batch_wall_ns,
recv_ts,
#[cfg(feature = "latency-trace")]
publish_rec,
#[cfg(feature = "tick-to-trade")]
ingest_rec,
);
match action {
FrameAction::Continue => false,
FrameAction::Disconnect => true,
FrameAction::PipelineFull => {
debug!(
connection_id = conn.connection_id,
"pipeline full, routing ServerBusy via response stage"
);
let _ = control_tx.send(ControlEvent::PipelineBusy {
connection_id: conn.connection_id,
});
false
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::client_frames::MAX_REQUEST_BODY;
use melin_app::auth::NoRoles;
use melin_app::decoder::{Decoded, RequestDecoder};
use melin_app::{AppEvent, Application, ApplyCtx, CodecError, QueryCtx, RejectReason};
use melin_journal::JournalEvent;
use melin_pipeline::ring::DisruptorBuilder;
use melin_wire_protocol::control_codec::{
TAG_APP, TAG_CHALLENGE_RESPONSE, TAG_LEN, TAG_RESPONSE_HEARTBEAT,
};
use std::io::{ErrorKind, Read};
use std::os::unix::net::UnixStream;
use std::sync::atomic::AtomicU64;
use std::time::Duration;
mod ring_entry {
use super::super::{RingEntry, ring_entry};
#[test]
fn visible_completions_and_nothing_to_submit_need_no_syscall() {
assert_eq!(ring_entry(0, true), RingEntry::Skip);
}
#[test]
fn queued_sqes_are_always_handed_over() {
assert_eq!(ring_entry(1, true), RingEntry::Submit);
assert_eq!(ring_entry(8, false), RingEntry::SubmitAndWait);
}
#[test]
fn an_empty_cq_is_the_only_thing_that_parks_the_thread() {
assert_eq!(ring_entry(0, false), RingEntry::SubmitAndWait);
for sq in [0, 1, 64] {
assert_ne!(ring_entry(sq, true), RingEntry::SubmitAndWait);
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TestEvent {
Cmd(u8),
Query,
}
impl AppEvent for TestEvent {
const MAX_ENCODED_SIZE: usize = 2;
fn encoded_size(&self) -> usize {
2
}
fn encode(&self, _buf: &mut [u8]) -> usize {
unreachable!("process_frames does not encode app events")
}
fn decode(_buf: &[u8]) -> Result<Self, CodecError> {
unreachable!("process_frames does not decode app events directly")
}
fn is_query(&self) -> bool {
matches!(self, TestEvent::Query)
}
}
#[derive(Default)]
struct TestApp;
type TestReport = (TestEvent, RejectReason);
impl Application for TestApp {
type Event = TestEvent;
type Report = TestReport;
type QueryResponse = ();
type Sizing = ();
const APP_VERSION: u16 = 0;
fn apply(&mut self, _event: TestEvent, _ctx: &ApplyCtx, _out: &mut Vec<TestReport>) {
unreachable!()
}
fn query(&self, _event: TestEvent, _ctx: &QueryCtx) -> Option<()> {
unreachable!()
}
fn tick(&mut self, _now_ns: u64, _out: &mut Vec<TestReport>) {
unreachable!()
}
fn build_reject(event: &TestEvent, reason: RejectReason) -> TestReport {
(*event, reason)
}
fn snapshot<W: std::io::Write>(&self, _w: &mut W) -> std::io::Result<()> {
unreachable!()
}
fn restore<R: Read>(_r: &mut R) -> std::io::Result<Self> {
unreachable!()
}
}
struct ByteDecoder;
impl RequestDecoder for ByteDecoder {
type Event = TestEvent;
type Role = NoRoles;
fn decode(&self, body: &[u8], _role: ClientRole<NoRoles>) -> Decoded<TestEvent> {
match body.first().copied() {
None => Decoded::DecodeError("empty body"),
Some(0xFC) => Decoded::Filter,
Some(0xFD) => Decoded::PermissionDenied("denied"),
Some(0xFE) => Decoded::DecodeError("bad"),
Some(0xFF) => Decoded::Permitted(TestEvent::Query),
Some(b) => Decoded::Permitted(TestEvent::Cmd(b)),
}
}
}
fn request_frame(tag: u8, body: &[u8]) -> Vec<u8> {
let len = (TAG_LEN + body.len()) as u32;
let mut f = Vec::with_capacity(4 + len as usize);
f.extend_from_slice(&len.to_le_bytes());
f.push(tag);
f.extend_from_slice(body);
f
}
fn frame(byte: u8) -> Vec<u8> {
request_frame(TAG_APP, &[byte])
}
fn oversize_prefix() -> [u8; 4] {
((MAX_FRAME_SIZE as u32) + 1).to_le_bytes()
}
struct Fixture {
conn: ConnectionEntry<UnixStream>,
producer: ring::Producer<InputSlot<TestEvent>>,
consumer: ring::Consumer<InputSlot<TestEvent>>,
peer: UnixStream,
}
fn make_fixture(ring_capacity: usize) -> Fixture {
let (server_side, peer) = UnixStream::pair().expect("UnixStream::pair");
peer.set_read_timeout(Some(Duration::from_millis(50)))
.expect("set_read_timeout");
let entry = ConnectionEntry::<UnixStream> {
connection_id: 7,
addr: "127.0.0.1:1".parse().expect("addr parses"),
role: ClientRole::Operator,
key_hash: 0xC0FFEE_u64,
fd: server_side.as_raw_fd(),
_reader: server_side,
parse_buf: Vec::with_capacity(64),
multishot_active: false,
last_activity: Instant::now(),
dying: false,
};
let (producer, mut consumers) =
DisruptorBuilder::<InputSlot<TestEvent>>::new(ring_capacity)
.add_consumer()
.build(melin_pipeline::wait::WaitStrategy::SpinThenYield);
let consumer = consumers.pop().expect("consumer present");
Fixture {
conn: entry,
producer,
consumer,
peer,
}
}
fn run_process_frames(
conn: &mut ConnectionEntry<UnixStream>,
producer: &mut ring::Producer<InputSlot<TestEvent>>,
) -> (bool, mpsc::Receiver<ControlEvent>) {
let (mut refusals, _queue) = refusal_channel();
run_process_frames_with(conn, producer, &gate(None, false), &mut refusals)
}
fn gate(replicas: Option<u32>, fenced: bool) -> HaltGate {
gate_counting_into(replicas, fenced, Arc::new(AtomicU64::new(0)))
}
fn gate_counting_into(
replicas: Option<u32>,
fenced: bool,
refused: Arc<AtomicU64>,
) -> HaltGate {
let fence = Arc::new(melin_transport_core::fence::FenceState::new(0));
if fenced {
fence.fence();
}
HaltGate::new(
replicas.map(crate::halt::test_support::presence),
fence,
refused,
)
}
fn refused_count(refused: &AtomicU64) -> u64 {
refused.load(std::sync::atomic::Ordering::Relaxed)
}
fn policy_claim() -> crate::policy_control::BootedClaim {
crate::policy_control::test_support::booted(
&crate::policy_control::test_support::replicated(crate::ack_policy::AckPolicy::Disk),
)
}
#[test]
fn a_parked_reader_journals_an_ack_policy_request() {
use crate::ack_policy::AckPolicy;
use crate::halt::Verdict;
use crate::halt::test_support::{gate_over, presence};
use crate::policy_control::RequestOutcome;
let replicas = presence(0);
let halt = || gate_over(Some(Arc::clone(&replicas)));
assert_eq!(
halt().verdict(),
Verdict::Refuse(RejectReason::ReplicaDisconnected)
);
let (handle, mut consumer, control) = spawn_parked_reader(halt(), None);
assert_eq!(control.request(AckPolicy::Disk), RequestOutcome::Published);
let (_, slot) = consumer.try_consume().expect("the entry is on the ring");
assert_eq!(
slot.event,
JournalEvent::AckPolicyChange(melin_journal::AckPolicyChange {
policy: AckPolicy::Disk.to_byte(),
unwitnessed: true,
})
);
assert_eq!(halt().verdict(), Verdict::Take);
handle.shutdown();
handle.join();
assert_eq!(
control.request(AckPolicy::Ram),
RequestOutcome::Staged,
"a stopped reader is detached"
);
}
#[test]
fn a_reader_journals_a_policy_staged_during_bring_up_first() {
use crate::ack_policy::AckPolicy;
use crate::halt::test_support::{gate_over, presence};
use crate::policy_control::RequestOutcome;
let (handle, mut consumer, control) =
spawn_parked_reader(gate_over(Some(presence(0))), Some(AckPolicy::Disk));
assert_eq!(control.request(AckPolicy::Disk), RequestOutcome::Published);
handle.shutdown();
handle.join();
let mut changes = Vec::new();
while let Some((_, slot)) = consumer.try_consume() {
if let JournalEvent::AckPolicyChange(change) = slot.event {
changes.push(change.unwitnessed);
}
}
assert_eq!(
changes,
[false, true],
"staged first, then the forced runtime one"
);
}
fn spawn_parked_reader(
halt: HaltGate,
staged: Option<crate::ack_policy::AckPolicy>,
) -> (
UringReaderHandle<UnixStream>,
ring::Consumer<InputSlot<TestEvent>>,
Arc<crate::policy_control::AckPolicyControl>,
) {
use crate::ack_policy::AckPolicy;
use crate::policy_control::RequestOutcome;
use crate::policy_control::test_support::{control_with_publish_wait, first_change};
let (mut producer, mut consumers) = DisruptorBuilder::<InputSlot<TestEvent>>::new(64)
.add_consumer()
.build(melin_pipeline::wait::WaitStrategy::SpinThenYield);
let consumer = consumers.pop().expect("consumer");
let (control_tx, _) = mpsc::channel();
let control = control_with_publish_wait(AckPolicy::DiskAndRam, Duration::from_secs(30));
let claim =
control
.claim()
.journal_boot(first_change(AckPolicy::DiskAndRam), &mut producer, &halt);
if let Some(policy) = staged {
assert_eq!(control.request(policy), RequestOutcome::Staged);
}
let handle = spawn_reader::<TestApp, UnixStream>(
producer,
Arc::new(ByteDecoder),
halt,
refusal_channel().0,
claim,
control_tx,
0,
None,
None, RingSizing::for_max_connections(4).expect("a supported cap"),
Arc::new(AtomicBool::new(false)),
)
.expect("the reader starts");
std::thread::sleep(Duration::from_millis(50));
(handle, consumer, control)
}
fn refusal_channel() -> (
RefusalSender<TestReport>,
crate::halt::RefusalQueue<TestReport>,
) {
crate::halt::refusal_channel(Arc::new(melin_pipeline::padding::CachePadded::new(
std::sync::atomic::AtomicU64::new(0),
)))
}
fn run_process_frames_with(
conn: &mut ConnectionEntry<UnixStream>,
producer: &mut ring::Producer<InputSlot<TestEvent>>,
halt: &HaltGate,
refusals: &mut RefusalSender<TestReport>,
) -> (bool, mpsc::Receiver<ControlEvent>) {
#[cfg(feature = "latency-trace")]
let mut publish_rec = melin_transport_core::trace::register_stage("test: publish");
#[cfg(feature = "tick-to-trade")]
let mut ingest_rec = melin_transport_core::trace::register_stage("test: ingest");
#[allow(clippy::let_unit_value)] let recv_ts = melin_transport_core::trace::mono_trace_ns();
let (control_tx, control_rx) = mpsc::channel();
let disconnect = process_frames::<TestApp, UnixStream>(
conn,
producer,
&ByteDecoder,
halt,
refusals,
&control_tx,
0xDEAD_BEEF,
recv_ts,
#[cfg(feature = "latency-trace")]
&mut publish_rec,
#[cfg(feature = "tick-to-trade")]
&mut ingest_rec,
);
(disconnect, control_rx)
}
fn busy_events(rx: &mpsc::Receiver<ControlEvent>) -> usize {
let mut n = 0;
while let Ok(event) = rx.try_recv() {
if matches!(event, ControlEvent::PipelineBusy { connection_id } if connection_id == 7) {
n += 1;
}
}
n
}
fn drain(
consumer: &mut ring::Consumer<InputSlot<TestEvent>>,
) -> Vec<(u64, InputSlot<TestEvent>)> {
let mut out = Vec::new();
while let Some(pair) = consumer.try_consume() {
out.push(pair);
}
out
}
fn read_server_busy(peer: &mut UnixStream) -> Option<[u8; 5]> {
let mut buf = [0u8; 5];
match peer.read_exact(&mut buf) {
Ok(()) => Some(buf),
Err(e) if e.kind() == ErrorKind::WouldBlock || e.kind() == ErrorKind::TimedOut => None,
Err(e) => panic!("unexpected peer read error: {e}"),
}
}
#[test]
fn process_frames_publishes_all_frames_after_single_commit() {
let Fixture {
mut conn,
mut producer,
mut consumer,
mut peer,
} = make_fixture(16);
for byte in [0x01, 0x02, 0x03, 0x04, 0x05] {
conn.parse_buf.extend_from_slice(&frame(byte));
}
let (disconnect, control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(!disconnect, "no oversize frame ⇒ no disconnect");
let events = drain(&mut consumer);
assert_eq!(events.len(), 5, "all 5 frames must be visible");
for (i, (seq, slot)) in events.iter().enumerate() {
assert_eq!(*seq, i as u64, "seq monotonic from 0");
assert_eq!(slot.connection_id, 7);
assert_eq!(slot.key_hash, 0xC0FFEE_u64);
let byte = (i + 1) as u8;
assert_eq!(slot.event, JournalEvent::App(TestEvent::Cmd(byte)));
assert_eq!(slot.timestamp_ns, 0xDEAD_BEEF);
}
assert!(conn.parse_buf.is_empty());
assert_eq!(
busy_events(&control_rx),
0,
"no PipelineBusy on the happy path"
);
assert!(
read_server_busy(&mut peer).is_none(),
"reader must never write the client socket"
);
}
#[test]
fn a_halted_reader_refuses_writes_before_publishing_and_still_publishes_queries() {
let Fixture {
mut conn,
mut producer,
mut consumer,
..
} = make_fixture(16);
for byte in [0x01, 0xFF, 0x02, 0xFF] {
conn.parse_buf.extend_from_slice(&frame(byte));
}
let (mut refusals, mut queue) = refusal_channel();
let counted = Arc::new(AtomicU64::new(0));
let (disconnect, control_rx) = run_process_frames_with(
&mut conn,
&mut producer,
&gate_counting_into(Some(0), false, Arc::clone(&counted)),
&mut refusals,
);
assert!(!disconnect);
assert_eq!(busy_events(&control_rx), 0);
assert!(conn.parse_buf.is_empty(), "every frame consumed");
assert_eq!(
refused_count(&counted),
2,
"the two writes are counted, the queries are not"
);
let published: Vec<_> = drain(&mut consumer)
.into_iter()
.map(|(seq, slot)| (seq, slot.event))
.collect();
assert_eq!(
published,
[
(0, JournalEvent::App(TestEvent::Query)),
(1, JournalEvent::App(TestEvent::Query)),
],
"only the queries enter the pipeline"
);
let mut refused = Vec::new();
assert!(queue.sync(), "the refusals are flushed");
queue.release(u64::MAX, |r| refused.push(r));
let expect = |input_seq, byte| crate::halt::Refusal {
connection_id: 7,
input_seq,
report: (TestEvent::Cmd(byte), RejectReason::ReplicaDisconnected),
};
assert_eq!(
refused,
[expect(0, 0x01), expect(1, 0x02)],
"each write is stamped with the sequence of the query that followed it"
);
}
#[test]
fn a_fenced_reader_closes_the_connection() {
let Fixture {
mut conn,
mut producer,
mut consumer,
..
} = make_fixture(16);
for byte in [0x01, 0xFF] {
conn.parse_buf.extend_from_slice(&frame(byte));
}
let (mut refusals, mut queue) = refusal_channel();
let counted = Arc::new(AtomicU64::new(0));
let (disconnect, control_rx) = run_process_frames_with(
&mut conn,
&mut producer,
&gate_counting_into(Some(1), true, Arc::clone(&counted)),
&mut refusals,
);
assert!(disconnect, "the connection is dropped");
assert_eq!(busy_events(&control_rx), 0);
assert!(
drain(&mut consumer).is_empty(),
"not even the query is published"
);
assert!(!queue.sync(), "no refusal is queued");
assert_eq!(refused_count(&counted), 0);
}
#[test]
fn a_full_refusal_queue_is_pipeline_busy() {
let Fixture {
mut conn,
mut producer,
mut consumer,
..
} = make_fixture(16);
conn.parse_buf.extend_from_slice(&frame(0x01));
let (mut refusals, _queue) = refusal_channel();
let filler = crate::halt::Refusal {
connection_id: 1,
input_seq: 0,
report: (TestEvent::Cmd(0), RejectReason::ReplicaDisconnected),
};
while refusals.try_send(filler).is_ok() {}
let counted = Arc::new(AtomicU64::new(0));
let (disconnect, control_rx) = run_process_frames_with(
&mut conn,
&mut producer,
&gate_counting_into(Some(0), false, Arc::clone(&counted)),
&mut refusals,
);
assert!(!disconnect);
assert_eq!(busy_events(&control_rx), 1);
assert!(drain(&mut consumer).is_empty());
assert_eq!(
refused_count(&counted),
1,
"a shed write still counts as refused by the halt"
);
}
#[cfg(feature = "latency-trace")]
#[test]
fn process_frames_stamps_every_slot_with_caller_recv_ts() {
let Fixture {
mut conn,
mut producer,
mut consumer,
..
} = make_fixture(16);
for byte in [0x01, 0x02, 0x03, 0x04] {
conn.parse_buf.extend_from_slice(&frame(byte));
}
const RECV_TS: u64 = 0x5EED_5EED;
let mut publish_rec = melin_transport_core::trace::register_stage("test: publish recv_ts");
#[cfg(feature = "tick-to-trade")]
let mut ingest_rec = melin_transport_core::trace::register_stage("test: ingest recv_ts");
let (control_tx, _control_rx) = mpsc::channel();
let (mut refusals, _queue) = refusal_channel();
let disconnect = process_frames::<TestApp, UnixStream>(
&mut conn,
&mut producer,
&ByteDecoder,
&gate(None, false),
&mut refusals,
&control_tx,
0xDEAD_BEEF,
RECV_TS,
&mut publish_rec,
#[cfg(feature = "tick-to-trade")]
&mut ingest_rec,
);
assert!(!disconnect);
let events = drain(&mut consumer);
assert_eq!(events.len(), 4, "all frames published");
for (_seq, slot) in &events {
assert_eq!(
slot.recv_ts, RECV_TS,
"every frame from one recv shares the single caller-supplied recv_ts"
);
}
}
#[test]
fn process_frames_rotates_batch_at_commit_every_cap() {
let Fixture {
mut conn,
mut producer,
mut consumer,
..
} = make_fixture(64);
const EVENT_COUNT: usize = 32;
for i in 0..EVENT_COUNT {
conn.parse_buf.extend_from_slice(&frame((i + 1) as u8));
}
let (disconnect, _control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(!disconnect, "no oversize / no Full ⇒ no disconnect");
let events = drain(&mut consumer);
assert_eq!(events.len(), EVENT_COUNT, "every event visible");
for (i, (seq, slot)) in events.iter().enumerate() {
assert_eq!(*seq, i as u64, "seq contiguous across batch rotations");
let byte = (i + 1) as u8;
assert_eq!(slot.event, JournalEvent::App(TestEvent::Cmd(byte)));
}
assert!(conn.parse_buf.is_empty());
}
#[test]
fn process_frames_query_event_skips_wall_clock_stamp() {
let Fixture {
mut conn,
mut producer,
mut consumer,
..
} = make_fixture(8);
conn.parse_buf.extend_from_slice(&frame(0xFF));
let (disconnect, _control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(!disconnect);
let events = drain(&mut consumer);
assert_eq!(events.len(), 1);
let (_, slot) = &events[0];
assert_eq!(slot.event, JournalEvent::App(TestEvent::Query));
assert_eq!(
slot.timestamp_ns, 0,
"query events must skip the wall-clock stamp"
);
}
#[test]
fn process_frames_partial_commit_then_server_busy_when_pipeline_full() {
let Fixture {
mut conn,
mut producer,
mut consumer,
mut peer,
} = make_fixture(4);
for byte in [0x01, 0x02, 0x03, 0x04, 0x05, 0x06] {
conn.parse_buf.extend_from_slice(&frame(byte));
}
let (disconnect, control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(!disconnect, "Full does not drop the connection");
let events = drain(&mut consumer);
assert_eq!(
events.len(),
4,
"only the first capacity-many frames are visible"
);
for (i, (_, slot)) in events.iter().enumerate() {
let byte = (i + 1) as u8;
assert_eq!(slot.event, JournalEvent::App(TestEvent::Cmd(byte)));
}
assert_eq!(
busy_events(&control_rx),
1,
"exactly one PipelineBusy routed"
);
assert!(
read_server_busy(&mut peer).is_none(),
"reader must never write the client socket"
);
assert_eq!(
conn.parse_buf,
frame(0x06),
"the 6th frame remains in parse_buf for the next recv-cycle"
);
}
#[test]
fn pipeline_full_never_touches_the_client_socket() {
let Fixture {
mut conn,
mut producer,
mut consumer,
mut peer,
} = make_fixture(2);
let junk = [0u8; 4096];
loop {
let n = unsafe {
libc::send(
conn.fd,
junk.as_ptr().cast(),
junk.len(),
libc::MSG_DONTWAIT,
)
};
if n < 0 {
assert_eq!(
std::io::Error::last_os_error().raw_os_error(),
Some(libc::EAGAIN),
"unexpected errno while filling socket"
);
break;
}
}
conn.parse_buf.extend_from_slice(&frame(0x01));
conn.parse_buf.extend_from_slice(&frame(0x02));
conn.parse_buf.extend_from_slice(&frame(0x03));
let start = std::time::Instant::now();
let (disconnect, control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(
start.elapsed() < Duration::from_secs(1),
"pipeline-full path blocked on a full socket"
);
assert!(!disconnect, "Full does not drop the connection");
assert_eq!(
drain(&mut consumer).len(),
2,
"the frames that fit are still published"
);
assert_eq!(
busy_events(&control_rx),
1,
"busy notice routed to response stage"
);
let mut sink = [0u8; 4096];
loop {
match peer.read(&mut sink) {
Ok(0) => panic!("unexpected EOF"),
Ok(_) => continue,
Err(e) if e.kind() == ErrorKind::WouldBlock || e.kind() == ErrorKind::TimedOut => {
break;
}
Err(e) => panic!("unexpected peer read error: {e}"),
}
}
}
#[test]
fn process_frames_oversize_commits_prior_frames_then_signals_disconnect() {
let Fixture {
mut conn,
mut producer,
mut consumer,
mut peer,
} = make_fixture(16);
conn.parse_buf.extend_from_slice(&frame(0x01));
conn.parse_buf.extend_from_slice(&frame(0x02));
conn.parse_buf.extend_from_slice(&oversize_prefix());
let (disconnect, control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(disconnect, "oversize frame must request disconnect");
let events = drain(&mut consumer);
assert_eq!(
events.len(),
2,
"prior frames are committed before the break"
);
assert_eq!(events[0].1.event, JournalEvent::App(TestEvent::Cmd(0x01)));
assert_eq!(events[1].1.event, JournalEvent::App(TestEvent::Cmd(0x02)));
assert_eq!(
busy_events(&control_rx),
0,
"PipelineBusy is emitted on Full, not on oversize"
);
assert!(
read_server_busy(&mut peer).is_none(),
"reader must never write the client socket"
);
}
#[test]
fn process_frames_filters_denied_and_decode_errors_advance_cursor() {
let Fixture {
mut conn,
mut producer,
mut consumer,
..
} = make_fixture(16);
for byte in [0x01, 0xFC, 0xFD, 0xFE, 0x02] {
conn.parse_buf.extend_from_slice(&frame(byte));
}
let (disconnect, _control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(!disconnect);
let events = drain(&mut consumer);
assert_eq!(events.len(), 2);
assert_eq!(events[0].1.event, JournalEvent::App(TestEvent::Cmd(0x01)));
assert_eq!(events[1].1.event, JournalEvent::App(TestEvent::Cmd(0x02)));
assert!(
conn.parse_buf.is_empty(),
"all bytes advanced past compaction"
);
}
#[test]
fn process_frames_reads_the_request_tag_itself() {
let Fixture {
mut conn,
mut producer,
mut consumer,
..
} = make_fixture(16);
conn.parse_buf.extend_from_slice(&frame(0x01));
conn.parse_buf.extend_from_slice(&0u32.to_le_bytes());
for tag in [0x00, TAG_CHALLENGE_RESPONSE, TAG_RESPONSE_HEARTBEAT, 0x10] {
conn.parse_buf
.extend_from_slice(&request_frame(tag, &[0x03]));
}
conn.parse_buf.extend_from_slice(&frame(0x02));
let (disconnect, _control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(!disconnect, "malformed requests do not drop the connection");
let events = drain(&mut consumer);
assert_eq!(events.len(), 2);
assert_eq!(events[0].1.event, JournalEvent::App(TestEvent::Cmd(0x01)));
assert_eq!(events[1].1.event, JournalEvent::App(TestEvent::Cmd(0x02)));
assert!(conn.parse_buf.is_empty(), "every frame was consumed");
}
#[test]
fn process_frames_preserves_partial_trailing_frame() {
let Fixture {
mut conn,
mut producer,
mut consumer,
..
} = make_fixture(16);
conn.parse_buf.extend_from_slice(&frame(0x42));
conn.parse_buf.extend_from_slice(&[0xDE, 0xAD, 0xBE]);
let (disconnect, _control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(!disconnect);
let events = drain(&mut consumer);
assert_eq!(events.len(), 1);
assert_eq!(events[0].1.event, JournalEvent::App(TestEvent::Cmd(0x42)));
assert_eq!(
conn.parse_buf,
vec![0xDE, 0xAD, 0xBE],
"partial length prefix preserved for next recv-cycle"
);
}
#[test]
fn process_frames_admits_a_frame_at_the_limit_and_drops_one_past_it() {
let Fixture {
mut conn,
mut producer,
mut consumer,
..
} = make_fixture(16);
let mut body = vec![0x05; MAX_REQUEST_BODY];
body[0] = 0x2A;
let at_limit = request_frame(TAG_APP, &body);
assert_eq!(at_limit.len(), 4 + MAX_FRAME_SIZE);
conn.parse_buf.extend_from_slice(&at_limit[..4 + 1]);
let (disconnect, _control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(!disconnect, "a frame at the limit is not oversized");
assert_eq!(conn.parse_buf, at_limit[..4 + 1]);
assert!(drain(&mut consumer).is_empty());
conn.parse_buf.extend_from_slice(&at_limit[4 + 1..]);
let (disconnect, _control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(!disconnect);
let events = drain(&mut consumer);
assert_eq!(events.len(), 1);
assert_eq!(events[0].1.event, JournalEvent::App(TestEvent::Cmd(0x2A)));
assert!(conn.parse_buf.is_empty());
conn.parse_buf.extend_from_slice(&oversize_prefix());
let (disconnect, _control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(disconnect, "oversize prefix alone must request disconnect");
assert!(drain(&mut consumer).is_empty());
}
#[test]
fn process_frames_empty_buffer_is_noop() {
let Fixture {
mut conn,
mut producer,
mut consumer,
mut peer,
} = make_fixture(4);
let (disconnect, control_rx) = run_process_frames(&mut conn, &mut producer);
assert!(!disconnect);
assert_eq!(drain(&mut consumer).len(), 0);
assert!(conn.parse_buf.is_empty());
assert_eq!(busy_events(&control_rx), 0);
assert!(read_server_busy(&mut peer).is_none());
}
#[test]
fn reader_loop_soak_delivers_every_frame_in_order() {
use std::io::Write as _;
const CONNS: u64 = 4;
const FRAMES: u32 = 10_000;
let (producer, mut consumers) = DisruptorBuilder::<InputSlot<TestEvent>>::new(65536)
.add_consumer()
.build(melin_pipeline::wait::WaitStrategy::SpinThenYield);
let mut consumer = consumers.pop().expect("consumer");
let (control_tx, control_rx) = mpsc::channel();
let shutdown = Arc::new(AtomicBool::new(false));
let mut handle = spawn_reader::<TestApp, UnixStream>(
producer,
Arc::new(ByteDecoder),
gate(None, false),
refusal_channel().0,
policy_claim(),
control_tx,
0, None, None, RingSizing::for_max_connections(CONNS).expect("a supported cap"),
Arc::clone(&shutdown),
)
.expect("the reader starts");
let mut writers = Vec::new();
for id in 0..CONNS {
let (client, server_side) = UnixStream::pair().expect("socketpair");
handle.register(ReaderRegistration {
connection_id: ConnectionId(id),
reader: server_side,
addr: "127.0.0.1:1".parse().expect("addr"),
role: ClientRole::Operator,
key_hash: id,
});
writers.push(std::thread::spawn(move || {
let mut client = client;
for i in 0..FRAMES {
let byte = (i % 200 + 1) as u8;
client.write_all(&frame(byte)).expect("client write");
if i % 512 == 0 {
std::thread::sleep(Duration::from_micros(200));
}
}
client }));
}
let deadline = Instant::now() + Duration::from_secs(60);
let mut per_conn_count = [0u32; CONNS as usize];
let mut total: u64 = 0;
while total < CONNS * FRAMES as u64 {
assert!(
Instant::now() < deadline,
"soak stalled: {total} of {} frames after 60s ({per_conn_count:?})",
CONNS * FRAMES as u64
);
while let Some((_seq, slot)) = consumer.try_consume() {
let conn = slot.connection_id as usize;
assert!(conn < CONNS as usize, "unknown connection id");
let expected = (per_conn_count[conn] % 200 + 1) as u8;
assert_eq!(
slot.event,
JournalEvent::App(TestEvent::Cmd(expected)),
"conn {conn} frame {} out of order or corrupted",
per_conn_count[conn]
);
per_conn_count[conn] += 1;
total += 1;
}
std::hint::spin_loop();
}
if let Ok(event) = control_rx.try_recv() {
match event {
ControlEvent::Disconnected { connection_id } => {
panic!("spurious disconnect of connection {connection_id}")
}
ControlEvent::PipelineBusy { connection_id } => {
panic!("spurious pipeline-full for connection {connection_id}")
}
ControlEvent::Connected { .. } => {
unreachable!("reader never sends Connected")
}
}
}
let _clients: Vec<UnixStream> = writers
.into_iter()
.map(|w| w.join().expect("writer thread"))
.collect();
handle.shutdown();
handle.join();
}
#[test]
fn shutdown_quiesces_armed_connections_promptly() {
use std::io::Write as _;
const CONNS: u64 = 3;
let (producer, mut consumers) = DisruptorBuilder::<InputSlot<TestEvent>>::new(1024)
.add_consumer()
.build(melin_pipeline::wait::WaitStrategy::SpinThenYield);
let mut consumer = consumers.pop().expect("consumer");
let (control_tx, _control_rx) = mpsc::channel();
let shutdown = Arc::new(AtomicBool::new(false));
let mut handle = spawn_reader::<TestApp, UnixStream>(
producer,
Arc::new(ByteDecoder),
gate(None, false),
refusal_channel().0,
policy_claim(),
control_tx,
0, None, None, RingSizing::for_max_connections(CONNS).expect("a supported cap"),
Arc::clone(&shutdown),
)
.expect("the reader starts");
let mut clients = Vec::new();
for id in 0..CONNS {
let (mut client, server_side) = UnixStream::pair().expect("socketpair");
handle.register(ReaderRegistration {
connection_id: ConnectionId(id),
reader: server_side,
addr: "127.0.0.1:1".parse().expect("addr"),
role: ClientRole::Operator,
key_hash: id,
});
client.write_all(&frame(7)).expect("client write");
clients.push(client);
}
let deadline = Instant::now() + Duration::from_secs(30);
let mut seen = 0u64;
while seen < CONNS {
assert!(
Instant::now() < deadline,
"frames never arrived ({seen}/{CONNS})"
);
while consumer.try_consume().is_some() {
seen += 1;
}
std::hint::spin_loop();
}
let started = Instant::now();
handle.shutdown();
handle.join();
let elapsed = started.elapsed();
assert!(
elapsed < crate::uring_teardown::DRAIN_TIMEOUT / 4,
"reader teardown took {elapsed:?}, close to the drain deadline — \
armed operations were not woken"
);
drop(clients);
}
#[test]
fn torn_down_connection_cannot_inject_frames_into_a_reused_slot() {
use std::io::Write as _;
let (producer, mut consumers) = DisruptorBuilder::<InputSlot<TestEvent>>::new(65536)
.add_consumer()
.build(melin_pipeline::wait::WaitStrategy::SpinThenYield);
let mut consumer = consumers.pop().expect("consumer");
let (control_tx, control_rx) = mpsc::channel();
let shutdown = Arc::new(AtomicBool::new(false));
let mut handle = spawn_reader::<TestApp, UnixStream>(
producer,
Arc::new(ByteDecoder),
gate(None, false),
refusal_channel().0,
policy_claim(),
control_tx,
0,
None,
None,
RingSizing::for_max_connections(4).expect("a supported cap"),
Arc::clone(&shutdown),
)
.expect("the reader starts");
const A_ID: u64 = 100;
const B_ID: u64 = 200;
let (a_client, a_server) = UnixStream::pair().expect("socketpair");
handle.register(ReaderRegistration {
connection_id: ConnectionId(A_ID),
reader: a_server,
addr: "127.0.0.1:1".parse().expect("addr"),
role: ClientRole::Operator,
key_hash: A_ID,
});
let a_writer = std::thread::spawn(move || {
let mut a = a_client;
for byte in 0x01..=0x08u8 {
let _ = a.write_all(&frame(byte));
}
let _ = a.write_all(&oversize_prefix());
let stop = Instant::now() + Duration::from_millis(300);
while Instant::now() < stop {
let _ = a.write_all(&frame(0x05));
std::thread::sleep(Duration::from_micros(50));
}
a
});
let deadline = Instant::now() + Duration::from_secs(10);
loop {
match control_rx.try_recv() {
Ok(ControlEvent::Disconnected { connection_id }) if connection_id == A_ID => break,
Ok(_) => {}
Err(_) => {
assert!(Instant::now() < deadline, "no Disconnected for A");
std::thread::sleep(Duration::from_millis(1));
}
}
}
const B_FRAMES: u32 = 2_000;
let (b_client, b_server) = UnixStream::pair().expect("socketpair");
handle.register(ReaderRegistration {
connection_id: ConnectionId(B_ID),
reader: b_server,
addr: "127.0.0.1:2".parse().expect("addr"),
role: ClientRole::Operator,
key_hash: B_ID,
});
let b_writer = std::thread::spawn(move || {
let mut b = b_client;
for i in 0..B_FRAMES {
let byte = 0x21 + (i % 0x10) as u8; b.write_all(&frame(byte)).expect("B write");
if i % 128 == 0 {
std::thread::sleep(Duration::from_micros(100));
}
}
b
});
let deadline = Instant::now() + Duration::from_secs(30);
let mut b_seen = 0u32;
while b_seen < B_FRAMES {
assert!(
Instant::now() < deadline,
"stalled: {b_seen}/{B_FRAMES} B frames"
);
while let Some((_seq, slot)) = consumer.try_consume() {
match slot.event {
JournalEvent::App(TestEvent::Cmd(byte)) => {
if slot.connection_id == B_ID {
assert_eq!(slot.key_hash, B_ID, "B slot with foreign key hash");
assert!(
(0x21..=0x30).contains(&byte),
"byte {byte:#x} injected into B's stream"
);
b_seen += 1;
} else {
assert_eq!(slot.connection_id, A_ID, "unknown connection id");
assert!(
(0x01..=0x08).contains(&byte),
"unexpected byte {byte:#x} on A"
);
}
}
ref other => panic!("unexpected event {other:?}"),
}
}
std::hint::spin_loop();
}
let _a = a_writer.join().expect("A writer");
let _b = b_writer.join().expect("B writer");
handle.shutdown();
handle.join();
}
#[test]
fn a_ring_that_cannot_be_created_is_a_startup_error() {
let (producer, _consumers) = DisruptorBuilder::<InputSlot<TestEvent>>::new(64)
.add_consumer()
.build(melin_pipeline::wait::WaitStrategy::SpinThenYield);
let (control_tx, _control_rx) = mpsc::channel();
let result = spawn_reader::<TestApp, UnixStream>(
producer,
Arc::new(ByteDecoder),
gate(None, false),
refusal_channel().0,
policy_claim(),
control_tx,
0,
None,
None,
RingSizing::unchecked(0, 32),
Arc::new(AtomicBool::new(false)),
);
let err = match result {
Ok(_) => panic!("a zero-entry ring cannot be created"),
Err(e) => e,
};
let msg = err.to_string();
assert!(msg.contains("uring-reader"), "{msg}");
assert!(msg.contains("io_uring"), "{msg}");
}
#[test]
fn push_sqe_makes_room_when_the_sq_is_full() {
const ENTRIES: u32 = 8;
const PUSHES: u64 = 3 * ENTRIES as u64;
let mut ring = IoUring::builder()
.setup_cqsize(PUSHES as u32)
.build(ENTRIES)
.expect("create an 8-entry io_uring");
for token in 0..PUSHES {
let nop = opcode::Nop::new().build().user_data(token);
unsafe { push_sqe(&mut ring, &nop, "test NOP") };
}
ring.submit_and_wait(PUSHES as usize)
.expect("submit the remaining NOPs");
let mut seen: Vec<u64> = ring.completion().map(|cqe| cqe.user_data()).collect();
seen.sort_unstable();
assert_eq!(seen, (0..PUSHES).collect::<Vec<_>>());
}
}