use std::io::{self, Write};
use std::net::{TcpListener, TcpStream};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use tracing::{debug, error, info, warn};
use melin_journal::replication::ReplicationConsumer;
use melin_pipeline::wait::WaitStrategy;
use melin_transport_core::replica_presence::{FirstAck, ReplicaPresence};
use super::auth::authenticate_replica;
use super::{ReplicaCursors, ReplicationMetrics, SentHighWater};
use crate::layout::Placement;
use crate::uring_teardown::{DrainBackoff, wake_pending_ops};
use melin_app::Application;
use melin_transport_core::replication::catchup::{
CatchUpResult, bridge_catchup_to_live, can_catch_up_from_journal, catch_up_from_journal,
preflight_snapshot_transfer, snapshot_transfer_with,
};
use melin_transport_core::replication::protocol::{
Ack, MAX_CONTROL_FRAME, ReplicaMessage, decode_replica_message, encode_hash_mismatch,
encode_heartbeat, encode_need_snapshot, encode_stream_start, read_frame,
};
use melin_transport_core::replication::validate::{
HandshakeValidation, validate_replica_handshake_settled,
};
use melin_wire_protocol::framing::split_frame_limited;
pub struct ReplicationListener(TcpListener);
impl ReplicationListener {
pub fn new(listener: TcpListener) -> std::io::Result<Self> {
listener.set_nonblocking(true)?;
Ok(Self(listener))
}
}
pub struct Sender {
pub listener: ReplicationListener,
pub repl_consumer_1: ReplicationConsumer,
pub repl_consumer_2: ReplicationConsumer,
pub replica_slots: Arc<melin_transport_core::ReplicaSlotCursors>,
pub journal_path: std::path::PathBuf,
pub authorized_keys: Arc<melin_app::auth::AuthorizedKeys>,
pub evict_flags: [Arc<AtomicBool>; 2],
pub active_flags: [Arc<AtomicBool>; 2],
pub metrics: Arc<ReplicationMetrics>,
pub handlers: [Placement; 2],
pub batch_size: usize,
pub heartbeat_secs: u64,
pub fence_state: Arc<melin_transport_core::fence::FenceState>,
}
pub fn run_sender<A: Application>(
config: Sender,
shutdown: &AtomicBool,
replica_ready: &AtomicBool,
replicas_connected: &ReplicaPresence,
) {
let Sender {
listener,
repl_consumer_1,
repl_consumer_2,
replica_slots,
journal_path,
authorized_keys,
evict_flags,
active_flags,
metrics,
handlers,
batch_size,
heartbeat_secs,
fence_state,
} = config;
let ReplicationListener(listener) = listener;
match listener.local_addr() {
Ok(addr) => info!(%addr, "replication sender listening"),
Err(_) => info!("replication sender listening"),
}
struct ReplicaSlot {
consumer: Option<ReplicationConsumer>,
handle: Option<std::thread::JoinHandle<ReplicationConsumer>>,
eviction_counted: bool,
}
let mut slots = [
ReplicaSlot {
consumer: Some(repl_consumer_1),
handle: None,
eviction_counted: false,
},
ReplicaSlot {
consumer: Some(repl_consumer_2),
handle: None,
eviction_counted: false,
},
];
let cursors = Arc::new(ReplicaCursors::new(replica_slots, Arc::clone(&metrics)));
let counted_flags: [Arc<AtomicBool>; 2] = [
Arc::new(AtomicBool::new(false)),
Arc::new(AtomicBool::new(false)),
];
loop {
if shutdown.load(Ordering::Relaxed) {
info!("replication sender shutting down");
for slot in &mut slots {
if let Some(handle) = slot.handle.take()
&& let Ok(consumer) = handle.join()
{
slot.consumer = Some(consumer);
}
}
return;
}
for (i, slot) in slots.iter_mut().enumerate() {
if take_new_eviction(
evict_flags[i].load(Ordering::Acquire),
slot.handle.is_some(),
&mut slot.eviction_counted,
) {
metrics.evictions_total.fetch_add(1, Ordering::Relaxed);
warn!(
slot = i,
"evicting slow replica (ring backpressure timeout)"
);
}
}
let disengage_slot = |slot_idx: usize| {
if counted_flags[slot_idx].swap(false, Ordering::AcqRel) {
replicas_connected.lower();
}
cursors.clear_on_disconnect(slot_idx);
metrics.catching_up[slot_idx].store(false, Ordering::Relaxed);
active_flags[slot_idx].store(false, Ordering::Release);
};
for (i, slot) in slots.iter_mut().enumerate() {
if let Some(ref handle) = slot.handle
&& handle.is_finished()
{
let handle = slot.handle.take().expect("just checked is_some");
slot.eviction_counted = false;
match handle.join() {
Ok(mut consumer) => {
consumer.skip_to_producer();
slot.consumer = Some(consumer);
disengage_slot(i);
if evict_flags[i].load(Ordering::Relaxed) {
evict_flags[i].store(false, Ordering::Release);
warn!(slot = i, "evicted replica — ring ready for reconnection");
} else {
warn!(slot = i, "replica disconnected");
}
}
Err(_) => {
error!(slot = i, "replica handler thread panicked");
disengage_slot(i);
evict_flags[i].store(false, Ordering::Release);
}
}
}
}
let empty_slot = slots
.iter()
.position(|s| s.handle.is_none() && s.consumer.is_some());
if let Some(slot_idx) = empty_slot {
match listener.accept() {
Ok((stream, addr)) => {
info!(addr = %addr, slot = slot_idx, "replica connected");
if let Err(e) = stream.set_nodelay(true) {
debug!(error = %e, "failed to set TCP_NODELAY on replica connection");
}
if let Err(e) =
crate::server::set_busy_poll(&stream, crate::server::BUSY_POLL_US)
{
debug!(error = %e, "failed to set SO_BUSY_POLL on replica connection");
}
let consumer = slots[slot_idx]
.consumer
.take()
.expect("empty_slot check guarantees consumer is Some");
let slot_cursors = Arc::clone(&cursors);
let jpath = journal_path.clone();
let auth_keys = Arc::clone(&authorized_keys);
let slot_metrics = Arc::clone(&metrics);
let slot_active = Arc::clone(&active_flags[slot_idx]);
let slot_evict = Arc::clone(&evict_flags[slot_idx]);
let slot_fence = Arc::clone(&fence_state);
let slot_counted = Arc::clone(&counted_flags[slot_idx]);
let handler = handlers[slot_idx];
let shutdown_flag = shutdown as *const AtomicBool as usize;
let ready_flag = replica_ready as *const AtomicBool as usize;
let connected_flag = replicas_connected as *const ReplicaPresence as usize;
let spawned = melin_app::affinity::spawn_placed(
&format!("repl-{slot_idx}"),
handler.core,
move || {
let shutdown_ref = unsafe { &*(shutdown_flag as *const AtomicBool) };
let ready_ref = unsafe { &*(ready_flag as *const AtomicBool) };
let connected_ref =
unsafe { &*(connected_flag as *const ReplicaPresence) };
let ctx = SlotContext {
cursors: &slot_cursors,
journal_path: &jpath,
authorized_keys: &auth_keys,
shutdown: shutdown_ref,
replica_ready: ready_ref,
active_flag: &slot_active,
evict_flag: &slot_evict,
metrics: &slot_metrics,
fence_state: &slot_fence,
slot_idx,
batch_size,
heartbeat_secs,
wait: handler.wait,
replicas_connected: connected_ref,
counted: &slot_counted,
};
run_replica_slot::<A>(stream, consumer, &ctx)
},
);
let handle = spawned
.or_else(melin_app::affinity::SpawnPlacedError::keep_stuck)
.expect("spawn replica handler thread");
slots[slot_idx].handle = Some(handle);
}
Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
}
Err(e) => {
error!(error = %e, "replication accept error");
}
}
}
std::thread::sleep(std::time::Duration::from_millis(50));
}
}
fn take_new_eviction(evict_flag_set: bool, handler_running: bool, counted: &mut bool) -> bool {
if evict_flag_set && handler_running && !*counted {
*counted = true;
true
} else {
false
}
}
struct SlotContext<'a> {
cursors: &'a Arc<ReplicaCursors>,
journal_path: &'a std::path::Path,
authorized_keys: &'a melin_app::auth::AuthorizedKeys,
shutdown: &'a AtomicBool,
replica_ready: &'a AtomicBool,
active_flag: &'a AtomicBool,
evict_flag: &'a AtomicBool,
metrics: &'a ReplicationMetrics,
fence_state: &'a melin_transport_core::fence::FenceState,
slot_idx: usize,
batch_size: usize,
heartbeat_secs: u64,
wait: WaitStrategy,
replicas_connected: &'a ReplicaPresence,
counted: &'a AtomicBool,
}
fn run_replica_slot<A: Application>(
stream: TcpStream,
mut consumer: ReplicationConsumer,
ctx: &SlotContext<'_>,
) -> ReplicationConsumer {
match handle_replica_connection::<A>(stream, &mut consumer, ctx) {
Ok(()) => info!("replica disconnected cleanly"),
Err(e) => warn!(error = %e, "replica connection error"),
}
consumer
}
fn handle_replica_connection<A: Application>(
stream: TcpStream,
repl_consumer: &mut ReplicationConsumer,
ctx: &SlotContext<'_>,
) -> io::Result<()> {
let SlotContext {
cursors,
journal_path,
authorized_keys,
shutdown,
replica_ready,
active_flag,
evict_flag: _,
metrics,
fence_state,
slot_idx,
batch_size: _,
heartbeat_secs,
wait: _,
replicas_connected,
counted,
} = ctx;
let slot_idx = *slot_idx;
let heartbeat_secs = *heartbeat_secs;
let mut reader = stream.try_clone()?;
let mut writer = stream;
reader.set_read_timeout(Some(std::time::Duration::from_secs(10)))?;
authenticate_replica(&mut reader, authorized_keys)?;
info!("replica authenticated");
let handshake_frame = read_frame(&mut reader, MAX_CONTROL_FRAME)?;
let handshake = match decode_replica_message(&handshake_frame)? {
ReplicaMessage::Handshake(h) => h,
ReplicaMessage::Ack(_) => {
return Err(io::Error::other("expected Handshake, got Ack"));
}
};
info!(
last_sequence = handshake.last_sequence,
replica_epoch = handshake.epoch,
"replica handshake received"
);
let our_epoch = fence_state.epoch();
if let Some(first_latch) = fence_state.fence_if_superseded(handshake.epoch, shutdown) {
if first_latch {
error!(
replica_epoch = handshake.epoch,
our_epoch,
"fenced: a replica advertises a higher epoch — this primary has been \
superseded; self-demoting and shutting down"
);
}
return Err(io::Error::other("fenced by higher-epoch replica"));
}
metrics.catching_up[slot_idx].store(true, Ordering::Relaxed);
let mut send_buf = Vec::with_capacity(128);
let divergent = match validate_replica_handshake_settled(journal_path, &handshake)? {
HandshakeValidation::Ok => false,
HandshakeValidation::Divergent(kind) => {
metrics.divergence_total.fetch_add(1, Ordering::Relaxed);
warn!(
last_sequence = handshake.last_sequence,
?kind,
"replica journal divergent at handshake — routing through snapshot resync"
);
true
}
};
let can_catch_up =
!divergent && can_catch_up_from_journal(journal_path, handshake.last_sequence)?;
let mut first_ack = replicas_connected.lift();
counted.store(true, Ordering::Release);
let mut publish = |buf: &[u8]| -> io::Result<()> {
writer.write_all(buf)?;
writer.flush()
};
let stream_base = if divergent {
0
} else {
handshake.last_sequence
};
let catchup_end = if can_catch_up {
let origin = melin_transport_core::replication::catchup::lineage_origin(journal_path)?;
encode_stream_start(
handshake.last_sequence,
origin.starting_sequence,
origin.anchor_hash,
origin.genesis_entries,
fence_state.epoch(),
&mut send_buf,
);
publish(&send_buf)?;
send_buf.clear();
match catch_up_from_journal::<A::Event>(
journal_path,
handshake.last_sequence,
&mut writer,
shutdown,
)? {
CatchUpResult::Ok(end) => end,
CatchUpResult::NeedSnapshot => {
return Err(io::Error::other("catch-up failed unexpectedly after probe"));
}
}
} else {
preflight_snapshot_transfer(journal_path)?;
if divergent {
encode_hash_mismatch(&mut send_buf);
} else {
encode_need_snapshot(&mut send_buf);
}
publish(&send_buf)?;
send_buf.clear();
match snapshot_transfer_with::<A::Event>(journal_path, &mut publish, shutdown)? {
CatchUpResult::Ok(end) => end,
CatchUpResult::NeedSnapshot => {
return Err(io::Error::other(
"catch-up failed even after snapshot transfer",
));
}
}
};
cursors.seed_on_handshake(slot_idx, stream_base);
let mut sent = {
let mut publish = |buf: &[u8]| -> io::Result<()> {
writer.write_all(buf)?;
writer.flush()
};
bridge_catchup_to_live::<A::Event>(
journal_path,
stream_base,
catchup_end,
active_flag,
repl_consumer,
&mut publish,
shutdown,
)?
};
metrics.catching_up[slot_idx].store(false, Ordering::Relaxed);
super::log_replica_live(slot_idx, &sent);
replica_ready.store(true, Ordering::Release);
let heartbeat_interval = std::time::Duration::from_secs(heartbeat_secs);
let mut last_send = std::time::Instant::now();
live_stream_uring(
writer,
repl_consumer,
ctx,
heartbeat_interval,
&mut send_buf,
&mut last_send,
&mut sent,
&mut first_ack,
)
}
struct InflightUring<'a> {
ring: io_uring::IoUring,
recv_buf: Vec<u8>,
send_buf: &'a mut Vec<u8>,
in_flight: u32,
tcp_fd: std::os::unix::io::RawFd,
}
impl Drop for InflightUring<'_> {
fn drop(&mut self) {
wake_pending_ops(self.tcp_fd);
if !self.drain() {
std::mem::forget(std::mem::take(&mut self.recv_buf));
std::mem::forget(std::mem::take(self.send_buf));
warn!(
in_flight = self.in_flight,
"io_uring teardown drain did not complete; leaking replication buffers"
);
}
}
}
impl InflightUring<'_> {
fn drain(&mut self) -> bool {
loop {
match self.ring.submit() {
Ok(_) => break,
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(_) => return false,
}
}
let mut backoff = DrainBackoff::new();
while self.in_flight > 0 {
for _ in self.ring.completion() {
self.in_flight = self.in_flight.saturating_sub(1);
}
if self.in_flight > 0 && !backoff.wait() {
return false;
}
}
true
}
}
fn live_stream_uring(
writer: TcpStream,
repl_consumer: &mut ReplicationConsumer,
ctx: &SlotContext<'_>,
heartbeat_interval: std::time::Duration,
send_buf: &mut Vec<u8>,
last_send: &mut std::time::Instant,
sent: &mut SentHighWater,
first_ack: &mut FirstAck,
) -> io::Result<()> {
let SlotContext {
cursors,
shutdown,
evict_flag,
metrics,
slot_idx,
batch_size,
wait,
replicas_connected,
journal_path: _,
authorized_keys: _,
replica_ready: _,
active_flag: _,
heartbeat_secs: _,
fence_state: _,
counted: _,
} = ctx;
let slot_idx = *slot_idx;
let batch_size = *batch_size;
let wait = *wait;
use io_uring::{IoUring, opcode, types};
use std::os::unix::io::AsRawFd;
const TOKEN_RECV: u64 = 0;
const TOKEN_SEND: u64 = 1;
let tcp_fd = writer.as_raw_fd();
let ring: IoUring = IoUring::builder()
.setup_single_issuer()
.build(8)
.map_err(|e| io::Error::other(format!("io_uring init failed: {e}")))?;
ring.submitter()
.register_files(&[tcp_fd])
.map_err(|e| io::Error::other(format!("io_uring register_files: {e}")))?;
{
let mut cpuset: libc::cpu_set_t = unsafe { std::mem::zeroed() };
unsafe { libc::CPU_SET(0, &mut cpuset) };
let _ = ring.submitter().register_iowq_aff(&cpuset);
}
let mut uring = InflightUring {
ring,
recv_buf: vec![0u8; 4096],
send_buf,
in_flight: 0,
tcp_fd,
};
let (ring, recv_buf, send_buf) = (&mut uring.ring, &mut uring.recv_buf, &mut *uring.send_buf);
let mut parse_buf: Vec<u8> = Vec::with_capacity(MAX_CONTROL_FRAME + 4);
let mut send_in_flight = false;
let mut send_offset: usize = 0;
let mut waiter = wait.waiter();
let mut heartbeat_timer = melin_app::amortized_timer::AmortizedTimer::new();
let mut persist_watch = melin_transport_core::replication::PersistWatch::default();
let mut info_log_timer = melin_app::amortized_timer::AmortizedTimer::new();
let mut send_submit_ts: Option<std::time::Instant> = None;
const SLOW_SEND_THRESHOLD_MS: u128 = 5;
let sqe = opcode::Recv::new(
types::Fixed(0),
recv_buf.as_mut_ptr(),
recv_buf.len() as u32,
)
.build()
.user_data(TOKEN_RECV);
unsafe { ring.submission().push(&sqe).expect("SQ full") };
uring.in_flight += 1;
loop {
if shutdown.load(Ordering::Relaxed) {
return Ok(());
}
if evict_flag.load(Ordering::Relaxed) {
super::log_tcp_info(tcp_fd, "evict_exit", slot_idx);
info!(slot = slot_idx, "handler exiting: evicted by journal stage");
return Ok(());
}
if !send_in_flight {
let mut coalesced = 0;
while coalesced < batch_size {
if let Some(batch) = repl_consumer.peek() {
send_buf.extend_from_slice(batch.data());
let end_sequence = batch.meta().end_sequence;
batch.commit();
sent.advance(end_sequence);
coalesced += 1;
} else {
break;
}
}
if coalesced > 0 {
let sqe =
opcode::Send::new(types::Fixed(0), send_buf.as_ptr(), send_buf.len() as u32)
.build()
.user_data(TOKEN_SEND);
unsafe { ring.submission().push(&sqe).expect("SQ full") };
uring.in_flight += 1;
send_in_flight = true;
send_offset = 0;
*last_send = std::time::Instant::now();
send_submit_ts = Some(*last_send);
waiter.reset();
heartbeat_timer = melin_app::amortized_timer::AmortizedTimer::new();
} else {
if heartbeat_timer
.tick(heartbeat_interval, waiter.spinning())
.is_some()
{
encode_heartbeat(sent.get(), send_buf);
let sqe = opcode::Send::new(
types::Fixed(0),
send_buf.as_ptr(),
send_buf.len() as u32,
)
.build()
.user_data(TOKEN_SEND);
unsafe { ring.submission().push(&sqe).expect("SQ full") };
uring.in_flight += 1;
send_in_flight = true;
send_offset = 0;
*last_send = std::time::Instant::now();
send_submit_ts = Some(*last_send);
}
}
}
if info_log_timer
.tick(std::time::Duration::from_secs(1), waiter.spinning())
.is_some()
{
super::log_tcp_info(tcp_fd, "live_stream", slot_idx);
}
let pending = ring.submission().len();
if pending > 0 {
ring.submit()
.map_err(|e| io::Error::other(format!("io_uring submit: {e}")))?;
}
let mut cqes: [(u64, i32); 4] = [(0, 0); 4];
let mut cqe_count = 0;
for cqe in ring.completion() {
debug_assert!(
uring.in_flight > 0,
"CQE reaped with no operation recorded in flight"
);
uring.in_flight = uring.in_flight.saturating_sub(1);
if cqe_count < cqes.len() {
cqes[cqe_count] = (cqe.user_data(), cqe.result());
cqe_count += 1;
}
}
let any_cqe = cqe_count > 0;
for &(token, result) in &cqes[..cqe_count] {
waiter.reset();
match token {
TOKEN_RECV => {
if result <= 0 {
return Err(io::Error::other(format!(
"replica disconnected (recv returned {result})"
)));
}
let n = result as usize;
parse_buf.extend_from_slice(&recv_buf[..n]);
let received = std::time::Instant::now();
drain_ack_frames(&mut parse_buf, |ack| {
cursors
.record_ack(slot_idx, &ack, sent.get())
.map_err(io::Error::other)?;
first_ack.acked(replicas_connected);
metrics.ack_latency_us[slot_idx]
.store(last_send.elapsed().as_micros() as u64, Ordering::Relaxed);
if let Some(event) = persist_watch.observe(&ack, received) {
event.log(slot_idx);
}
Ok(())
})?;
let sqe = opcode::Recv::new(
types::Fixed(0),
recv_buf.as_mut_ptr(),
recv_buf.len() as u32,
)
.build()
.user_data(TOKEN_RECV);
unsafe { ring.submission().push(&sqe).expect("SQ full") };
uring.in_flight += 1;
}
TOKEN_SEND => {
if result < 0 {
return Err(io::Error::other(format!("send error (returned {result})")));
}
let sent = result as usize;
send_offset += sent;
if send_offset >= send_buf.len() {
if let Some(ts) = send_submit_ts.take() {
let elapsed = ts.elapsed();
if elapsed.as_millis() >= SLOW_SEND_THRESHOLD_MS {
tracing::debug!(
slot = slot_idx,
elapsed_us = elapsed.as_micros() as u64,
bytes = send_buf.len(),
"slow SEND completion"
);
super::log_tcp_info(tcp_fd, "slow_send", slot_idx);
}
}
metrics.bytes_sent[slot_idx]
.fetch_add(send_buf.len() as u64, Ordering::Relaxed);
send_buf.clear();
send_offset = 0;
send_in_flight = false;
} else {
let sqe = opcode::Send::new(
types::Fixed(0),
send_buf[send_offset..].as_ptr(),
(send_buf.len() - send_offset) as u32,
)
.build()
.user_data(TOKEN_SEND);
unsafe { ring.submission().push(&sqe).expect("SQ full") };
uring.in_flight += 1;
}
}
_ => {}
}
}
if !any_cqe && send_buf.is_empty() {
waiter.idle();
}
}
}
fn drain_ack_frames(
parse_buf: &mut Vec<u8>,
mut on_ack: impl FnMut(Ack) -> io::Result<()>,
) -> io::Result<()> {
let mut cursor = 0;
loop {
let frame = match split_frame_limited(&parse_buf[cursor..], MAX_CONTROL_FRAME) {
Ok(Some(frame)) if !frame.payload.is_empty() => frame,
Ok(Some(_)) => return Err(io::Error::other("invalid ack frame length: 0")),
Ok(None) => break,
Err(too_large) => {
return Err(io::Error::other(format!(
"invalid ack frame length: {}",
too_large.declared
)));
}
};
if let Ok(ReplicaMessage::Ack(ack)) = decode_replica_message(frame.payload) {
on_ack(ack)?;
}
cursor += frame.consumed();
}
if cursor > 0 {
let remaining = parse_buf.len() - cursor;
parse_buf.copy_within(cursor.., 0);
parse_buf.truncate(remaining);
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::uring_teardown::DRAIN_TIMEOUT;
use io_uring::{IoUring, opcode, types};
fn socketpair() -> (libc::c_int, libc::c_int) {
let mut fds = [0 as libc::c_int; 2];
let rc = unsafe { libc::socketpair(libc::AF_UNIX, libc::SOCK_STREAM, 0, fds.as_mut_ptr()) };
assert_eq!(rc, 0, "socketpair: {}", io::Error::last_os_error());
(fds[0], fds[1])
}
fn close_pair(a: libc::c_int, b: libc::c_int) {
unsafe {
libc::close(a);
libc::close(b);
}
}
fn guard_with_pending_recv(fd: libc::c_int, send_buf: &mut Vec<u8>) -> InflightUring<'_> {
let ring = IoUring::builder()
.setup_single_issuer()
.build(8)
.expect("io_uring is a hard requirement for replication");
ring.submitter()
.register_files(&[fd])
.expect("register_files");
let mut uring = InflightUring {
ring,
recv_buf: vec![0u8; 4096],
send_buf,
in_flight: 0,
tcp_fd: fd,
};
let sqe = opcode::Recv::new(
types::Fixed(0),
uring.recv_buf.as_mut_ptr(),
uring.recv_buf.len() as u32,
)
.build()
.user_data(0);
unsafe { uring.ring.submission().push(&sqe).expect("SQ full") };
uring.in_flight = 1;
uring.ring.submit().expect("submit initial RECV");
uring
}
#[test]
fn drain_reaps_a_pending_recv_after_shutdown() {
let (a, b) = socketpair();
let mut send_buf = Vec::new();
let mut uring = guard_with_pending_recv(a, &mut send_buf);
unsafe { libc::shutdown(a, libc::SHUT_RDWR) };
assert!(
uring.drain(),
"drain hit the deadline with a RECV in flight"
);
assert_eq!(uring.in_flight, 0, "kernel still owes a CQE");
drop(uring);
close_pair(a, b);
}
#[test]
fn drop_quiesces_a_pending_recv_without_hanging() {
let (a, b) = socketpair();
let mut send_buf = vec![0u8; 64];
let uring = guard_with_pending_recv(a, &mut send_buf);
let started = std::time::Instant::now();
drop(uring);
let elapsed = started.elapsed();
assert!(
elapsed < DRAIN_TIMEOUT / 4,
"teardown took {elapsed:?}, close to the {DRAIN_TIMEOUT:?} deadline"
);
assert_eq!(send_buf.len(), 64, "send_buf leaked despite a clean drain");
close_pair(a, b);
}
#[test]
fn eviction_is_counted_once_per_handler() {
let mut counted = false;
assert!(!take_new_eviction(false, true, &mut counted));
assert!(!take_new_eviction(true, false, &mut counted));
assert!(!counted);
assert!(take_new_eviction(true, true, &mut counted));
for _ in 0..5 {
assert!(!take_new_eviction(true, true, &mut counted));
}
counted = false;
assert!(take_new_eviction(true, true, &mut counted));
}
fn ack_frame(seq: u64) -> Vec<u8> {
let mut buf = Vec::new();
melin_transport_core::replication::protocol::encode_ack(
&Ack {
acked_sequence: seq,
in_memory_sequence: seq,
},
&mut buf,
);
buf
}
fn drain(buf: &mut Vec<u8>) -> (io::Result<()>, Vec<u64>) {
let mut seen = Vec::new();
let result = drain_ack_frames(buf, |ack| {
seen.push(ack.acked_sequence);
Ok(())
});
(result, seen)
}
#[test]
fn ack_frames_are_handed_over_and_a_partial_one_is_kept() {
let mut buf = ack_frame(1);
buf.extend_from_slice(&ack_frame(2));
let third = ack_frame(3);
buf.extend_from_slice(&third[..third.len() - 1]);
let (result, seen) = drain(&mut buf);
result.unwrap();
assert_eq!(seen, [1, 2]);
assert_eq!(
buf,
third[..third.len() - 1],
"partial frame compacted to the front"
);
buf.push(*third.last().unwrap());
let (result, seen) = drain(&mut buf);
result.unwrap();
assert_eq!(seen, [3]);
assert!(buf.is_empty());
}
#[test]
fn a_frame_that_is_not_an_ack_is_skipped() {
let mut buf = vec![1, 0, 0, 0, 0xEE];
buf.extend_from_slice(&ack_frame(4));
let (result, seen) = drain(&mut buf);
result.unwrap();
assert_eq!(seen, [4]);
assert!(buf.is_empty());
}
#[test]
fn a_frame_at_the_cap_waits_for_its_payload() {
let mut buf = (MAX_CONTROL_FRAME as u32).to_le_bytes().to_vec();
let (result, seen) = drain(&mut buf);
result.unwrap();
assert!(seen.is_empty());
assert_eq!(buf.len(), 4, "prefix kept for the next RECV");
}
#[test]
fn an_oversized_prefix_fails_before_its_payload() {
let mut buf = ack_frame(5);
buf.extend_from_slice(&((MAX_CONTROL_FRAME as u32) + 1).to_le_bytes());
let (result, seen) = drain(&mut buf);
let err = result.unwrap_err();
assert_eq!(
err.to_string(),
format!("invalid ack frame length: {}", MAX_CONTROL_FRAME + 1)
);
assert_eq!(seen, [5]);
}
#[test]
fn a_zero_length_frame_fails() {
let mut buf = 0u32.to_le_bytes().to_vec();
let (result, seen) = drain(&mut buf);
assert_eq!(
result.unwrap_err().to_string(),
"invalid ack frame length: 0"
);
assert!(seen.is_empty());
}
#[test]
fn an_on_ack_error_stops_the_drain() {
let mut buf = ack_frame(6);
buf.extend_from_slice(&ack_frame(7));
let mut calls = 0;
let result = drain_ack_frames(&mut buf, |_| {
calls += 1;
Err(io::Error::other("cursor violation"))
});
assert_eq!(result.unwrap_err().to_string(), "cursor violation");
assert_eq!(calls, 1);
}
}