melin_server_runtime/replication/tcp_sender.rs
1//! TCP replication sender (primary side).
2//!
3//! Listens for replica connections, handles authentication, journal
4//! catch-up, snapshot transfer, and live streaming. The wire protocol
5//! is defined in `super::protocol`.
6
7use std::io::{self, Write};
8use std::net::{TcpListener, TcpStream};
9use std::sync::Arc;
10use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
11
12use tracing::{debug, error, info, warn};
13
14use melin_journal::replication::ReplicationConsumer;
15use melin_pipeline::wait::WaitStrategy;
16
17use super::auth::authenticate_replica;
18use super::{ReplicaCursors, ReplicaGate, ReplicationMetrics, SentHighWater};
19use crate::layout::Placement;
20use crate::uring_teardown::{DrainBackoff, wake_pending_ops};
21use melin_app::Application;
22use melin_transport_core::replication::catchup::{
23 CatchUpResult, bridge_catchup_to_live, can_catch_up_from_journal, catch_up_from_journal,
24 preflight_snapshot_transfer, snapshot_transfer_with,
25};
26use melin_transport_core::replication::protocol::{
27 Ack, MAX_CONTROL_FRAME, ReplicaMessage, decode_replica_message, encode_hash_mismatch,
28 encode_heartbeat, encode_need_snapshot, encode_stream_start, read_frame,
29};
30use melin_transport_core::replication::validate::{
31 HandshakeValidation, validate_replica_handshake_settled,
32};
33use melin_wire_protocol::framing::split_frame_limited;
34
35// --- Replication Sender (Primary side) ---
36
37/// Owned state for the replication sender thread.
38/// A TCP listener guaranteed non-blocking: the only constructor sets
39/// the flag, so the guarantee is carried by the type. [`run_sender`]'s
40/// accept loop polls the shutdown flag between accepts — handed a
41/// blocking listener it would park inside `accept()` and only observe
42/// shutdown when the next replica happened to dial in. This newtype
43/// makes that state unrepresentable instead of documenting it away.
44pub struct ReplicationListener(TcpListener);
45
46impl ReplicationListener {
47 /// Wrap a bound listener, setting it non-blocking.
48 pub fn new(listener: TcpListener) -> std::io::Result<Self> {
49 listener.set_nonblocking(true)?;
50 Ok(Self(listener))
51 }
52}
53
54pub struct Sender {
55 /// Pre-bound replication listener, non-blocking by construction.
56 /// Bound by the caller at startup — not here — so a bad or stolen
57 /// `--replication-bind` port fails the boot with a clear error.
58 /// Binding on this thread meant a failure only killed the sender
59 /// silently, leaving a primary that looked healthy but could never
60 /// satisfy `disk+ram`/`two-disks` acks.
61 pub listener: ReplicationListener,
62 pub repl_consumer_1: ReplicationConsumer,
63 pub repl_consumer_2: ReplicationConsumer,
64 pub replica_slots: Arc<melin_transport_core::ReplicaSlotCursors>,
65 pub journal_path: std::path::PathBuf,
66 pub authorized_keys: Arc<melin_app::auth::AuthorizedKeys>,
67 pub evict_flags: [Arc<AtomicBool>; 2],
68 pub active_flags: [Arc<AtomicBool>; 2],
69 pub metrics: Arc<ReplicationMetrics>,
70 /// Where each replica handler thread runs and how it waits on its
71 /// replication ring, by slot.
72 pub handlers: [Placement; 2],
73 pub batch_size: usize,
74 pub heartbeat_secs: u64,
75 /// Node fencing state. Read to stamp the primary's epoch onto each
76 /// `StreamStart`, and to self-demote when a replica handshakes with a
77 /// higher epoch (this primary has been superseded). See `crate::fence`.
78 pub fence_state: Arc<melin_transport_core::fence::FenceState>,
79 /// The active ack policy, shared with the response gate and the
80 /// admin `ACK-POLICY` command. Stamped on every `StreamStart` and
81 /// `Heartbeat` so replicas judge auto-promotion against the policy
82 /// this primary actually acks under.
83 pub ack_policy: Arc<std::sync::atomic::AtomicU8>,
84}
85
86/// Run the replication sender. Listens for replica connections,
87/// streams journal data batches, processes acks, and updates the
88/// replication cursor.
89///
90/// Runs on a dedicated thread. Blocks until shutdown.
91pub fn run_sender<A: Application>(
92 config: Sender,
93 shutdown: &AtomicBool,
94 replica_ready: &AtomicBool,
95 replicas_connected: &AtomicU32,
96) {
97 let Sender {
98 listener,
99 repl_consumer_1,
100 repl_consumer_2,
101 replica_slots,
102 journal_path,
103 authorized_keys,
104 evict_flags,
105 active_flags,
106 metrics,
107 handlers,
108 batch_size,
109 heartbeat_secs,
110 fence_state,
111 ack_policy,
112 } = config;
113 // Unwrap the invariant-carrying newtype — from here `listener` is a
114 // plain TcpListener, known non-blocking.
115 let ReplicationListener(listener) = listener;
116 match listener.local_addr() {
117 Ok(addr) => info!(%addr, "replication sender listening"),
118 // Best-effort log detail — the pre-bound listener works either way.
119 Err(_) => info!("replication sender listening"),
120 }
121
122 // Two replica slots, each with its own ring consumer and thread handle.
123 // The accept loop fills empty slots. When a replica disconnects, its
124 // slot becomes available for a new connection. The consumer is `None`
125 // while a handler thread owns it (moved into the thread, returned on exit).
126 struct ReplicaSlot {
127 consumer: Option<ReplicationConsumer>,
128 handle: Option<std::thread::JoinHandle<ReplicationConsumer>>,
129 /// The current handler's eviction has been counted. A plain `bool`:
130 /// only this supervisor thread reads or writes it.
131 eviction_counted: bool,
132 }
133
134 let mut slots = [
135 ReplicaSlot {
136 consumer: Some(repl_consumer_1),
137 handle: None,
138 eviction_counted: false,
139 },
140 ReplicaSlot {
141 consumer: Some(repl_consumer_2),
142 handle: None,
143 eviction_counted: false,
144 },
145 ];
146
147 // Single owner of the per-replica progress cursors (per-slot acked
148 // positions and the gate's gauge pair). Shared by both handler
149 // threads — see `ReplicaCursors` for the ordering contract relative
150 // to the active flags.
151 let cursors = Arc::new(ReplicaCursors::new(replica_slots, Arc::clone(&metrics)));
152
153 // Per-slot "this connection authenticated" latch. The halt gate
154 // (`replicas_connected`) is lifted by the handler thread only after auth
155 // succeeds, so the main loop must lower it on teardown only when the
156 // handler got that far — a connection that drops mid-auth never counted.
157 // Per-slot `Arc<AtomicBool>` mirrors `active_flags`/`evict_flags`; shared
158 // with the handler thread, which flips it true post-auth.
159 let authenticated_flags: [Arc<AtomicBool>; 2] = [
160 Arc::new(AtomicBool::new(false)),
161 Arc::new(AtomicBool::new(false)),
162 ];
163
164 loop {
165 if shutdown.load(Ordering::Relaxed) {
166 info!("replication sender shutting down");
167 // Wait for active replica threads to finish.
168 for slot in &mut slots {
169 if let Some(handle) = slot.handle.take()
170 && let Ok(consumer) = handle.join()
171 {
172 slot.consumer = Some(consumer);
173 }
174 }
175 return;
176 }
177
178 // Check eviction flags from the journal stage. When set, the
179 // journal stage timed out publishing to this slot's ring. We need
180 // to reclaim the consumer so the idle drain loop can clear the ring,
181 // allowing the journal stage to resume publishing.
182 for (i, slot) in slots.iter_mut().enumerate() {
183 // The handler watches the same flag and exits on its own; it is
184 // collected below once finished. Until then the flag stays set
185 // across passes, so count and log the eviction once.
186 if take_new_eviction(
187 evict_flags[i].load(Ordering::Acquire),
188 slot.handle.is_some(),
189 &mut slot.eviction_counted,
190 ) {
191 metrics.evictions_total.fetch_add(1, Ordering::Relaxed);
192 warn!(
193 slot = i,
194 "evicting slow replica (ring backpressure timeout)"
195 );
196 }
197 }
198
199 // Disengage a finished slot's shared state. Both the clean-exit and
200 // panic arms below call this, so they cannot drift: a panicked handler
201 // must tear down exactly like a clean exit, or it leaks the halt gate
202 // and — far worse — leaves the slot's progress cursors engaged,
203 // so the primary stops acking client requests even with a healthy
204 // surviving replica. Borrows only atomics (interior
205 // mutability), so it coexists with the `&mut` iteration over `slots`.
206 let disengage_slot = |slot_idx: usize| {
207 // Lower the gate only if this connection authenticated (and so
208 // lifted it); `swap` reads and resets the latch in one op, and the
209 // `join()` ordered the handler's post-auth writes before here.
210 // `lower` warns "halted" if this was the last replica.
211 if authenticated_flags[slot_idx].swap(false, Ordering::AcqRel) {
212 ReplicaGate::new(replicas_connected).lower();
213 }
214 // Disengage cursors BEFORE clearing the active flag — ordering
215 // contract B2 (see `ReplicaCursors`).
216 cursors.clear_on_disconnect(slot_idx);
217 metrics.catching_up[slot_idx].store(false, Ordering::Relaxed);
218 // Journal stage stops publishing to this ring.
219 active_flags[slot_idx].store(false, Ordering::Release);
220 };
221
222 // Collect finished replica threads (disconnected replicas).
223 for (i, slot) in slots.iter_mut().enumerate() {
224 if let Some(ref handle) = slot.handle
225 && handle.is_finished()
226 {
227 let handle = slot.handle.take().expect("just checked is_some");
228 // Re-arm for the slot's next handler.
229 slot.eviction_counted = false;
230 match handle.join() {
231 Ok(mut consumer) => {
232 // Drop any unread entries before the consumer is
233 // stashed back in the slot. The journal stage may
234 // have published batches into the ring that the
235 // evicted handler never got to forward — if we
236 // left them in place, the NEXT handler on this
237 // slot would drain them to its replica and
238 // acknowledge with pre-eviction sequences. Those
239 // acks would stall the slot's acked-progress
240 // cursor at the old position and gate the
241 // primary's response stage at the slow-replica
242 // rate.
243 //
244 // Fast-forward to the producer cursor so the
245 // live-streaming loop starts with a clean ring.
246 consumer.skip_to_producer();
247 slot.consumer = Some(consumer);
248 disengage_slot(i);
249 // Clear eviction flag after reclaiming the consumer.
250 if evict_flags[i].load(Ordering::Relaxed) {
251 evict_flags[i].store(false, Ordering::Release);
252 warn!(slot = i, "evicted replica — ring ready for reconnection");
253 } else {
254 warn!(slot = i, "replica disconnected");
255 }
256 }
257 Err(_) => {
258 error!(slot = i, "replica handler thread panicked");
259 // Consumer is lost — the slot stays handle=None /
260 // consumer=None, so the accept loop skips it. With
261 // independent rings, only this slot's ring is affected;
262 // the other replica continues normally. Still disengage
263 // its shared state, or the gate leaks and the cursor
264 // freezes (see `disengage_slot`).
265 disengage_slot(i);
266 evict_flags[i].store(false, Ordering::Release);
267 }
268 }
269 }
270 }
271
272 // Find a slot that has a consumer available (not in use by a thread).
273 let empty_slot = slots
274 .iter()
275 .position(|s| s.handle.is_none() && s.consumer.is_some());
276
277 // Accept a connection if there's an empty slot.
278 if let Some(slot_idx) = empty_slot {
279 match listener.accept() {
280 Ok((stream, addr)) => {
281 info!(addr = %addr, slot = slot_idx, "replica connected");
282 if let Err(e) = stream.set_nodelay(true) {
283 debug!(error = %e, "failed to set TCP_NODELAY on replica connection");
284 }
285 // SO_BUSY_POLL on the sender side: the per-replica thread
286 // spins on `recv` for ack frames, so kernel busy-polling
287 // removes the softirq->wakeup handoff from the ack path.
288 if let Err(e) =
289 crate::server::set_busy_poll(&stream, crate::server::BUSY_POLL_US)
290 {
291 debug!(error = %e, "failed to set SO_BUSY_POLL on replica connection");
292 }
293
294 // The halt gate is NOT lifted here — a bare connect
295 // hasn't proven a Replication key. The handler lifts it after
296 // auth succeeds (see `handle_replica_connection`).
297
298 // Take the consumer out of the slot for the handler thread.
299 // The slot's consumer becomes None while the thread owns it.
300 let consumer = slots[slot_idx]
301 .consumer
302 .take()
303 .expect("empty_slot check guarantees consumer is Some");
304
305 let slot_cursors = Arc::clone(&cursors);
306 let jpath = journal_path.clone();
307 let auth_keys = Arc::clone(&authorized_keys);
308 let slot_metrics = Arc::clone(&metrics);
309 let slot_active = Arc::clone(&active_flags[slot_idx]);
310 let slot_evict = Arc::clone(&evict_flags[slot_idx]);
311 let slot_fence = Arc::clone(&fence_state);
312 let slot_ack_policy = Arc::clone(&ack_policy);
313 let slot_authenticated = Arc::clone(&authenticated_flags[slot_idx]);
314 let handler = handlers[slot_idx];
315 let shutdown_flag = shutdown as *const AtomicBool as usize;
316 let ready_flag = replica_ready as *const AtomicBool as usize;
317 let connected_flag = replicas_connected as *const AtomicU32 as usize;
318 let handle = std::thread::Builder::new()
319 .name(format!("repl-{slot_idx}"))
320 .spawn(move || {
321 // Pin to a dedicated core if configured (> 0),
322 // otherwise clear inherited affinity from the
323 // accept thread so the OS can schedule freely.
324 //
325 // Both branches only work because the accept
326 // thread is left unpinned (see its spawn in
327 // `server.rs`). A child inherits its creator's
328 // affinity and scheduling policy at creation,
329 // and cannot fix that itself if the parent is
330 // pinned real-time and never yields — fixing
331 // itself requires running. Pin the accept
332 // thread and this handler may never start.
333 if handler.is_pinned() {
334 match melin_app::affinity::pin_to_core(handler.core) {
335 Ok(c) => tracing::info!(
336 core = c,
337 slot = slot_idx,
338 "replication handler pinned to core"
339 ),
340 Err(e) => tracing::warn!(
341 error = e,
342 slot = slot_idx,
343 "failed to pin handler"
344 ),
345 }
346 } else if let Err(e) = melin_app::affinity::clear_affinity() {
347 tracing::warn!(error = e, "failed to clear handler affinity");
348 }
349 // Safety: shutdown and replica_ready outlive this thread
350 // (they're on the parent's stack, which blocks on join
351 // during shutdown).
352 let shutdown_ref = unsafe { &*(shutdown_flag as *const AtomicBool) };
353 let ready_ref = unsafe { &*(ready_flag as *const AtomicBool) };
354 // Safety: replicas_connected outlives this thread —
355 // same lifetime argument as shutdown/replica_ready
356 // above (the parent blocks on join during shutdown).
357 let connected_ref = unsafe { &*(connected_flag as *const AtomicU32) };
358 let ctx = SlotContext {
359 cursors: &slot_cursors,
360 journal_path: &jpath,
361 authorized_keys: &auth_keys,
362 shutdown: shutdown_ref,
363 replica_ready: ready_ref,
364 active_flag: &slot_active,
365 evict_flag: &slot_evict,
366 metrics: &slot_metrics,
367 fence_state: &slot_fence,
368 ack_policy: &slot_ack_policy,
369 slot_idx,
370 batch_size,
371 heartbeat_secs,
372 wait: handler.wait,
373 replicas_connected: connected_ref,
374 authenticated: &slot_authenticated,
375 };
376 run_replica_slot::<A>(stream, consumer, &ctx)
377 })
378 .expect("spawn replica handler thread");
379 slots[slot_idx].handle = Some(handle);
380 }
381 Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
382 // No pending connection.
383 }
384 Err(e) => {
385 error!(error = %e, "replication accept error");
386 }
387 }
388 }
389
390 // No idle drain needed — the journal stage only publishes to
391 // rings where active_flag is true (set by handler on live loop
392 // entry, cleared on disconnect). Idle consumers stay empty.
393
394 // Supervisory cadence. Nothing on the data path runs in this loop
395 // — the ring consumers are moved into the handler threads — so the
396 // only cost is up to 50 ms to notice a new connection or a handler
397 // that has finished, both negligible against the journal catch-up
398 // that follows. In exchange the thread costs ~no CPU when idle,
399 // which is why it needs no core of its own.
400 std::thread::sleep(std::time::Duration::from_millis(50));
401 }
402}
403
404/// Whether the supervisor should count an eviction this pass: the journal
405/// stage has flagged the slot, a handler is still running on it, and that
406/// handler's eviction has not been counted yet. Latches `counted`; the
407/// supervisor clears it when it collects the handler.
408fn take_new_eviction(evict_flag_set: bool, handler_running: bool, counted: &mut bool) -> bool {
409 if evict_flag_set && handler_running && !*counted {
410 *counted = true;
411 true
412 } else {
413 false
414 }
415}
416
417/// Per-slot state shared across the replica handler call chain
418/// (`run_replica_slot` → `handle_replica_connection` → `live_stream_uring`).
419struct SlotContext<'a> {
420 cursors: &'a Arc<ReplicaCursors>,
421 journal_path: &'a std::path::Path,
422 authorized_keys: &'a melin_app::auth::AuthorizedKeys,
423 shutdown: &'a AtomicBool,
424 replica_ready: &'a AtomicBool,
425 active_flag: &'a AtomicBool,
426 evict_flag: &'a AtomicBool,
427 metrics: &'a ReplicationMetrics,
428 fence_state: &'a melin_transport_core::fence::FenceState,
429 /// The primary's active ack policy — stamped on `StreamStart` and
430 /// `Heartbeat` (see `Sender::ack_policy`).
431 ack_policy: &'a std::sync::atomic::AtomicU8,
432 slot_idx: usize,
433 batch_size: usize,
434 heartbeat_secs: u64,
435 wait: WaitStrategy,
436 /// Halt gate. The handler lifts it (`fetch_add`) only after auth
437 /// succeeds; the main loop lowers it on teardown, gated by `authenticated`.
438 replicas_connected: &'a AtomicU32,
439 /// Per-slot latch the handler sets true once this connection passes auth,
440 /// so the main loop knows whether this slot lifted the gate.
441 authenticated: &'a AtomicBool,
442}
443
444/// Handle a single replica connection on a dedicated thread.
445/// Returns the consumer when the connection ends (for slot reuse).
446fn run_replica_slot<A: Application>(
447 stream: TcpStream,
448 mut consumer: ReplicationConsumer,
449 ctx: &SlotContext<'_>,
450) -> ReplicationConsumer {
451 match handle_replica_connection::<A>(stream, &mut consumer, ctx) {
452 Ok(()) => info!("replica disconnected cleanly"),
453 Err(e) => warn!(error = %e, "replica connection error"),
454 }
455 consumer
456}
457
458fn handle_replica_connection<A: Application>(
459 stream: TcpStream,
460 repl_consumer: &mut ReplicationConsumer,
461 ctx: &SlotContext<'_>,
462) -> io::Result<()> {
463 let SlotContext {
464 cursors,
465 journal_path,
466 authorized_keys,
467 shutdown,
468 replica_ready,
469 active_flag,
470 evict_flag: _,
471 metrics,
472 fence_state,
473 ack_policy,
474 slot_idx,
475 batch_size: _,
476 heartbeat_secs,
477 wait: _,
478 replicas_connected,
479 authenticated,
480 } = ctx;
481 let slot_idx = *slot_idx;
482 let heartbeat_secs = *heartbeat_secs;
483
484 let mut reader = stream.try_clone()?;
485 let mut writer = stream;
486
487 // Set a read timeout for the handshake and auth.
488 reader.set_read_timeout(Some(std::time::Duration::from_secs(10)))?;
489
490 // Authenticate before any data exchange. `reader` and `writer` are clones
491 // of the same socket; auth is sequential (write challenge, read response,
492 // write result) so a single handle suffices. Use `reader` — it carries the
493 // read timeout set above.
494 authenticate_replica(&mut reader, authorized_keys)?;
495 info!("replica authenticated");
496
497 // Auth passed — only now lift the halt gate. The main loop lowers
498 // it on slot teardown, reading `authenticated` to know this connection
499 // lifted it. Set the latch after the increment so a post-join read of the
500 // latch implies the increment is visible too.
501 ReplicaGate::new(replicas_connected).lift();
502 authenticated.store(true, Ordering::Release);
503
504 // Read handshake.
505 let handshake_frame = read_frame(&mut reader, MAX_CONTROL_FRAME)?;
506 let handshake = match decode_replica_message(&handshake_frame)? {
507 ReplicaMessage::Handshake(h) => h,
508 ReplicaMessage::Ack(_) => {
509 return Err(io::Error::other("expected Handshake, got Ack"));
510 }
511 };
512
513 info!(
514 last_sequence = handshake.last_sequence,
515 replica_epoch = handshake.epoch,
516 "replica handshake received"
517 );
518
519 // Fence: a replica advertising an epoch higher than ours means a
520 // promotion happened that this node missed — we are a stale ex-primary.
521 // The policy (latch fenced + trigger shutdown) lives on `FenceState` so
522 // the kernel-TCP and DPDK senders cannot drift; the matching stage
523 // halts new client writes, the response gate stops acking, and the node
524 // winds down for the operator to restart it as a replica. Refuse this
525 // connection so we send no stale stream.
526 let our_epoch = fence_state.epoch();
527 if let Some(first_latch) = fence_state.fence_if_superseded(handshake.epoch, shutdown) {
528 if first_latch {
529 error!(
530 replica_epoch = handshake.epoch,
531 our_epoch,
532 "fenced: a replica advertises a higher epoch — this primary has been \
533 superseded; self-demoting and shutting down"
534 );
535 }
536 return Err(io::Error::other("fenced by higher-epoch replica"));
537 }
538
539 // Mark this slot as catching up. Cleared when entering the live loop.
540 metrics.catching_up[slot_idx].store(true, Ordering::Relaxed);
541
542 let mut send_buf = Vec::with_capacity(128);
543
544 // Chain validation: recompute our chain at the replica's claimed
545 // position and compare with the hash it reported. A mismatch (or a
546 // position beyond our tip) means its journal holds divergent
547 // history — streaming on top would silently fork the audit trail.
548 // Send HashMismatch (the replica archives its local journal on
549 // receipt) and route it through snapshot resync on this same
550 // connection.
551 let divergent = match validate_replica_handshake_settled(journal_path, &handshake)? {
552 HandshakeValidation::Ok => false,
553 HandshakeValidation::Divergent(kind) => {
554 // Alertable on /metrics: divergence outside an expected
555 // failover rejoin means corruption or a serious bug.
556 metrics.divergence_total.fetch_add(1, Ordering::Relaxed);
557 warn!(
558 last_sequence = handshake.last_sequence,
559 ?kind,
560 "replica journal divergent at handshake — routing through snapshot resync"
561 );
562 true
563 }
564 };
565
566 // Probe whether journal catch-up is possible before committing to
567 // a protocol path. This avoids sending StreamStart only to discover
568 // the journals are too old.
569 let can_catch_up =
570 !divergent && can_catch_up_from_journal(journal_path, handshake.last_sequence)?;
571
572 let mut publish = |buf: &[u8]| -> io::Result<()> {
573 writer.write_all(buf)?;
574 writer.flush()
575 };
576
577 // Cursor/stream floor for everything downstream. A divergent
578 // replica's claimed position is meaningless (its history forked —
579 // possibly ahead of our tip); the resync rebases it, so seed from 0
580 // like a fresh replica: the ack cursors understate until real acks
581 // arrive (safe), and the sent high-water seeds from the catch-up
582 // end instead of a bogus claim that would mark live ring chunks as
583 // already covered.
584 let stream_base = if divergent {
585 0
586 } else {
587 handshake.last_sequence
588 };
589
590 let catchup_end = if can_catch_up {
591 // The lineage origin (oldest segment's header identity) lets a
592 // fresh replica create a byte-identical journal before
593 // consuming the stream; replicas with local state ignore it.
594 // The lineage's genesis length, from the same header, concerns
595 // every replica: a fresh one records it, one with a journal
596 // checks it against its own.
597 let origin = melin_transport_core::replication::catchup::lineage_origin(journal_path)?;
598 encode_stream_start(
599 handshake.last_sequence,
600 origin.starting_sequence,
601 origin.anchor_hash,
602 origin.genesis_entries,
603 fence_state.epoch(),
604 ack_policy.load(Ordering::Relaxed),
605 &mut send_buf,
606 );
607 publish(&send_buf)?;
608 send_buf.clear();
609
610 match catch_up_from_journal::<A::Event>(
611 journal_path,
612 handshake.last_sequence,
613 &mut writer,
614 shutdown,
615 )? {
616 CatchUpResult::Ok(end) => end,
617 CatchUpResult::NeedSnapshot => {
618 return Err(io::Error::other("catch-up failed unexpectedly after probe"));
619 }
620 }
621 } else {
622 // Resync verdict precedes the snapshot data: `HashMismatch`
623 // tells the replica its journal is divergent (it archives its
624 // local lineage on receipt); plain `NeedSnapshot` is the
625 // too-far-behind rebase. The receiver expects `SnapshotBegin`
626 // as the very next frame after the verdict, so
627 // `snapshot_transfer_with` must follow immediately.
628 //
629 // Pre-flight before the verdict goes on the wire: the replica
630 // archives its lineage on receipt, so promising a snapshot we
631 // cannot produce (none configured, boundary segment pruned)
632 // would strand it with its history moved aside. Failing here
633 // drops the connection with the replica's journal intact.
634 preflight_snapshot_transfer(journal_path)?;
635 if divergent {
636 encode_hash_mismatch(&mut send_buf);
637 } else {
638 encode_need_snapshot(&mut send_buf);
639 }
640 publish(&send_buf)?;
641 send_buf.clear();
642 match snapshot_transfer_with::<A::Event>(
643 journal_path,
644 &mut publish,
645 shutdown,
646 ack_policy.load(Ordering::Relaxed),
647 )? {
648 CatchUpResult::Ok(end) => end,
649 CatchUpResult::NeedSnapshot => {
650 return Err(io::Error::other(
651 "catch-up failed even after snapshot transfer",
652 ));
653 }
654 }
655 };
656
657 // Engage this slot's cursors and seed the gauge pair the response
658 // gate's `evaluate_durability` reads. Must happen BEFORE the
659 // bridge's active_flag Release — see `ReplicaCursors` for the
660 // ordering contract.
661 cursors.seed_on_handshake(slot_idx, stream_base);
662
663 // Bridge into live streaming: activates the ring, re-reads from the
664 // journal the entries that fell into the activation window, then
665 // drains the ring into sequence-contiguity (back-filling from disk
666 // if a skipped entry hasn't flushed yet) before going live. Returns
667 // the slot's sent high-water — the bound the ack-sanity invariant
668 // checks acks against, advanced by every chunk the live loop
669 // forwards. The bridge closes the catch-up→live gap under load (the
670 // receiver's contiguity gate backstops only the rare quiescent
671 // corner). See `bridge_catchup_to_live`.
672 let mut sent = {
673 let mut publish = |buf: &[u8]| -> io::Result<()> {
674 writer.write_all(buf)?;
675 writer.flush()
676 };
677 bridge_catchup_to_live::<A::Event>(
678 journal_path,
679 stream_base,
680 catchup_end,
681 active_flag,
682 repl_consumer,
683 &mut publish,
684 shutdown,
685 )?
686 };
687
688 // Catch-up complete — replica is entering the live streaming loop.
689 metrics.catching_up[slot_idx].store(false, Ordering::Relaxed);
690
691 // Signal that this replica is ready to consume from the replication
692 // ring. A primary that began a new history waits on this (its
693 // bring-up gate) before serving clients, so it must be set only
694 // AFTER catch-up and overlap drain complete: the first client write
695 // then finds a replica streaming live rather than one still catching
696 // up.
697 replica_ready.store(true, Ordering::Release);
698
699 let heartbeat_interval = std::time::Duration::from_secs(heartbeat_secs);
700 let mut last_send = std::time::Instant::now();
701
702 live_stream_uring(
703 writer,
704 repl_consumer,
705 ctx,
706 heartbeat_interval,
707 &mut send_buf,
708 &mut last_send,
709 &mut sent,
710 )
711}
712
713/// Quiescing owner of the live-stream ring: the ring itself, the RECV
714/// landing buffer the kernel writes ack bytes into, the SEND source
715/// buffer the kernel reads out of, and the number of operations the
716/// kernel still owes CQEs for.
717///
718/// Exists for its `Drop`. Returning from the loop — shutdown, eviction,
719/// an I/O error, or a panic's unwind — while a RECV or SEND is still in
720/// flight would release those buffers (and drop the ring) with the
721/// kernel holding pointers into them. Ring teardown on fd close is
722/// *asynchronous*: the kernel can complete a pending RECV after the
723/// allocation has been returned to the allocator, writing replica-ack
724/// bytes into freed heap. That corruption surfaced as malloc-metadata
725/// segfaults at sender-thread exit in the failover tests, and being a
726/// kernel-side store it is invisible to AddressSanitizer. The guard
727/// closes every exit path: wake pending operations by shutting the
728/// socket down, then reap CQEs until the kernel owes none, and only
729/// then let the buffers go.
730struct InflightUring<'a> {
731 ring: io_uring::IoUring,
732 /// RECV landing buffer for ack frames. `Vec` so the heap pointer
733 /// handed to the kernel stays stable however the guard itself
734 /// moves; sized 4 KiB (acks are 13 bytes, the kernel may coalesce
735 /// many).
736 recv_buf: Vec<u8>,
737 /// SEND source buffer, borrowed from `handle_replica_connection`
738 /// (it outlives the live-stream loop so it can be reused across
739 /// reconnects). Held by the guard purely so the drain covers it:
740 /// the kernel *reads* it for an in-flight SEND, so releasing it
741 /// early would splice freed heap into the replication stream.
742 send_buf: &'a mut Vec<u8>,
743 /// SQEs pushed minus CQEs reaped. Between a push and the following
744 /// `submit` the kernel does not yet owe those CQEs, so this is an
745 /// upper bound rather than an exact debt; [`Self::drain`] submits
746 /// before waiting, which makes the two converge. `u32` matches the
747 /// ring-depth scale (at most a RECV and a SEND are ever
748 /// outstanding here).
749 in_flight: u32,
750 /// Raw socket fd, for the `shutdown(2)` wake in `Drop`. Always
751 /// outlives the guard: the owning `TcpStream` is a parameter of
752 /// [`live_stream_uring`] and parameters drop after locals.
753 tcp_fd: std::os::unix::io::RawFd,
754}
755
756impl Drop for InflightUring<'_> {
757 fn drop(&mut self) {
758 wake_pending_ops(self.tcp_fd);
759 if !self.drain() {
760 // The drain could not be proven complete, so the kernel may
761 // still write into `recv_buf` or read out of `send_buf`.
762 // Leak both: a few KiB lost on a dying connection is safe,
763 // returning memory the kernel still touches to the
764 // allocator is not. `take` leaves each `Vec` empty, so the
765 // caller's `send_buf` stays a valid (if reallocating) Vec.
766 std::mem::forget(std::mem::take(&mut self.recv_buf));
767 std::mem::forget(std::mem::take(self.send_buf));
768 warn!(
769 in_flight = self.in_flight,
770 "io_uring teardown drain did not complete; leaking replication buffers"
771 );
772 }
773 }
774}
775
776impl InflightUring<'_> {
777 /// Reap CQEs until the kernel owes none, bounded by
778 /// [`crate::uring_teardown::DRAIN_TIMEOUT`]. `true` means
779 /// `in_flight` reached zero and the
780 /// buffers are provably free of kernel references; `false` means
781 /// the caller must not release them. See [`crate::uring_teardown`]
782 /// for why this polls rather than blocking in `submit_and_wait`.
783 fn drain(&mut self) -> bool {
784 // Flush SQEs that were pushed but never submitted: until they
785 // reach the kernel, no CQE is ever coming for them.
786 loop {
787 match self.ring.submit() {
788 Ok(_) => break,
789 Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
790 Err(_) => return false,
791 }
792 }
793 let mut backoff = DrainBackoff::new();
794 while self.in_flight > 0 {
795 for _ in self.ring.completion() {
796 self.in_flight = self.in_flight.saturating_sub(1);
797 }
798 if self.in_flight > 0 && !backoff.wait() {
799 return false;
800 }
801 }
802 true
803 }
804}
805
806/// Live streaming loop using async RECV/SEND via io_uring. A single RECV is always
807/// in-flight for ack frames; SEND is submitted when the replication ring
808/// has data. Both complete via the memory-mapped CQ with zero syscalls
809/// in the hot path. In-flight operations and both kernel-visible
810/// buffers are owned by an [`InflightUring`] guard, so no exit path can
811/// release memory the kernel is still reading or writing.
812fn live_stream_uring(
813 writer: TcpStream,
814 repl_consumer: &mut ReplicationConsumer,
815 ctx: &SlotContext<'_>,
816 heartbeat_interval: std::time::Duration,
817 send_buf: &mut Vec<u8>,
818 last_send: &mut std::time::Instant,
819 sent: &mut SentHighWater,
820) -> io::Result<()> {
821 let SlotContext {
822 cursors,
823 shutdown,
824 evict_flag,
825 metrics,
826 slot_idx,
827 batch_size,
828 wait,
829 ack_policy,
830 // Only used during handshake/catch-up (handle_replica_connection).
831 journal_path: _,
832 authorized_keys: _,
833 replica_ready: _,
834 active_flag: _,
835 heartbeat_secs: _,
836 fence_state: _,
837 replicas_connected: _,
838 authenticated: _,
839 } = ctx;
840 let slot_idx = *slot_idx;
841 let batch_size = *batch_size;
842 let wait = *wait;
843
844 use io_uring::{IoUring, opcode, types};
845 use std::os::unix::io::AsRawFd;
846
847 const TOKEN_RECV: u64 = 0;
848 const TOKEN_SEND: u64 = 1;
849
850 let tcp_fd = writer.as_raw_fd();
851
852 let ring: IoUring = IoUring::builder()
853 .setup_single_issuer()
854 .build(8)
855 .map_err(|e| io::Error::other(format!("io_uring init failed: {e}")))?;
856
857 ring.submitter()
858 .register_files(&[tcp_fd])
859 .map_err(|e| io::Error::other(format!("io_uring register_files: {e}")))?;
860
861 // Pin io-wq workers to core 0 (keep them off pipeline cores).
862 {
863 let mut cpuset: libc::cpu_set_t = unsafe { std::mem::zeroed() };
864 unsafe { libc::CPU_SET(0, &mut cpuset) };
865 let _ = ring.submitter().register_iowq_aff(&cpuset);
866 }
867
868 // Ring, both kernel-visible buffers, and the in-flight count live
869 // in the quiescing guard from here on: every exit below (shutdown,
870 // eviction, `?`, unwind) must drain the kernel before any of these
871 // allocations can be released. The loop keeps using them through
872 // disjoint field borrows.
873 let mut uring = InflightUring {
874 ring,
875 recv_buf: vec![0u8; 4096],
876 send_buf,
877 in_flight: 0,
878 tcp_fd,
879 };
880 let (ring, recv_buf, send_buf) = (&mut uring.ring, &mut uring.recv_buf, &mut *uring.send_buf);
881 // Accumulation buffer for partial ack frame parsing.
882 let mut parse_buf: Vec<u8> = Vec::with_capacity(MAX_CONTROL_FRAME + 4);
883 // RECV is always resubmitted after CQE processing — no explicit
884 // tracking needed. The io_uring kernel guarantees ordering.
885 let mut send_in_flight = false;
886 let mut send_offset: usize = 0;
887 let mut waiter = wait.waiter();
888 let mut heartbeat_timer = melin_app::amortized_timer::AmortizedTimer::new();
889
890 // Diagnostic (RUST_LOG=debug): per-slot TCP_INFO snapshot once a
891 // second, slow-SEND detection (CQE elapsed >= threshold), and a
892 // TCP_INFO capture at the evict-exit point. Amortized so the
893 // per-iteration cost is a single `AND` + predictable branch.
894 let mut info_log_timer = melin_app::amortized_timer::AmortizedTimer::new();
895 let mut send_submit_ts: Option<std::time::Instant> = None;
896 const SLOW_SEND_THRESHOLD_MS: u128 = 5;
897
898 // Submit initial RECV.
899 let sqe = opcode::Recv::new(
900 types::Fixed(0),
901 recv_buf.as_mut_ptr(),
902 recv_buf.len() as u32,
903 )
904 .build()
905 .user_data(TOKEN_RECV);
906 // SAFETY: `recv_buf` is owned by the `InflightUring` guard, which
907 // drains this operation's CQE before the buffer can be freed on any
908 // exit path. The ring is single-threaded — only this task
909 // pushes/reaps on it.
910 unsafe { ring.submission().push(&sqe).expect("SQ full") };
911 uring.in_flight += 1;
912
913 loop {
914 // --- Check flags ---
915 if shutdown.load(Ordering::Relaxed) {
916 return Ok(());
917 }
918 if evict_flag.load(Ordering::Relaxed) {
919 // Capture the TCP state at the moment of eviction — the
920 // critical frame for comparing an evicted slot's teardown
921 // against the still-live slot's socket state when
922 // diagnosing post-eviction regressions.
923 super::log_tcp_info(tcp_fd, "evict_exit", slot_idx);
924 info!(slot = slot_idx, "handler exiting: evicted by journal stage");
925 return Ok(());
926 }
927
928 // --- Drain replication ring into send_buf (memory, non-blocking) ---
929 //
930 // Ring chunks are wire-ready `InputBatch` frames produced by the
931 // journal stage (phase 3 of the unified-pipeline plan), so the
932 // sender is a passthrough — no decode + re-encode here.
933 if !send_in_flight {
934 let mut coalesced = 0;
935 while coalesced < batch_size {
936 if let Some((meta, data)) = repl_consumer.try_read() {
937 send_buf.extend_from_slice(data);
938 repl_consumer.commit();
939 sent.advance(meta.end_sequence);
940 coalesced += 1;
941 } else {
942 break;
943 }
944 }
945
946 if coalesced > 0 {
947 // Submit SEND for the coalesced buffer.
948 let sqe =
949 opcode::Send::new(types::Fixed(0), send_buf.as_ptr(), send_buf.len() as u32)
950 .build()
951 .user_data(TOKEN_SEND);
952 // SAFETY: `send_buf` is owned by this task and pinned
953 // by `send_in_flight = true` below — it isn't mutated
954 // until the matching TOKEN_SEND CQE clears the flag,
955 // and the `InflightUring` guard drains the CQE before
956 // any exit path can release it. The ring is
957 // single-threaded.
958 unsafe { ring.submission().push(&sqe).expect("SQ full") };
959 uring.in_flight += 1;
960 send_in_flight = true;
961 send_offset = 0;
962 *last_send = std::time::Instant::now();
963 send_submit_ts = Some(*last_send);
964 waiter.reset();
965 heartbeat_timer = melin_app::amortized_timer::AmortizedTimer::new();
966 } else {
967 // Heartbeat check: amortized when spinning (mask keeps the
968 // clock read at ~10/s at 10M iter/s). In yield mode the loop
969 // already pays a syscall per iteration, so the clock read is
970 // free and must not be skipped — see AmortizedTimer docs.
971 if heartbeat_timer
972 .tick(heartbeat_interval, waiter.spinning())
973 .is_some()
974 {
975 encode_heartbeat(sent.get(), ack_policy.load(Ordering::Relaxed), send_buf);
976 let sqe = opcode::Send::new(
977 types::Fixed(0),
978 send_buf.as_ptr(),
979 send_buf.len() as u32,
980 )
981 .build()
982 .user_data(TOKEN_SEND);
983 // SAFETY: same as the coalesced send above —
984 // `send_buf` is owned and pinned by `send_in_flight`
985 // until the TOKEN_SEND CQE clears it, and the guard
986 // drains the CQE on every exit path.
987 unsafe { ring.submission().push(&sqe).expect("SQ full") };
988 uring.in_flight += 1;
989 send_in_flight = true;
990 send_offset = 0;
991 *last_send = std::time::Instant::now();
992 send_submit_ts = Some(*last_send);
993 }
994 }
995 }
996
997 // Periodic TCP_INFO dump — debug level. Amortized so the
998 // per-iteration cost is a single `AND` + predictable branch.
999 if info_log_timer
1000 .tick(std::time::Duration::from_secs(1), waiter.spinning())
1001 .is_some()
1002 {
1003 super::log_tcp_info(tcp_fd, "live_stream", slot_idx);
1004 }
1005
1006 // --- Submit SQEs to kernel (non-blocking) ---
1007 // Skip the syscall when no new SQEs were pushed — an empty
1008 // io_uring_enter still costs ~200 ns of mode-switch overhead and
1009 // showed up as 6 % of total CPU in profiles of the sender loop.
1010 let pending = ring.submission().len();
1011 if pending > 0 {
1012 ring.submit()
1013 .map_err(|e| io::Error::other(format!("io_uring submit: {e}")))?;
1014 }
1015
1016 // --- Collect CQEs (must drain before pushing new SQEs) ---
1017 // Collecting into a small stack array avoids the CQ borrow
1018 // conflicting with SQ pushes during processing.
1019 let mut cqes: [(u64, i32); 4] = [(0, 0); 4];
1020 let mut cqe_count = 0;
1021 for cqe in ring.completion() {
1022 // Every reaped CQE settles one operation the kernel owed —
1023 // counted even if it overflows the recording array (it
1024 // cannot: at most a RECV and a SEND are ever in flight).
1025 // `saturating_sub` keeps a miscount from wrapping into a
1026 // drain that never terminates; the assert catches the
1027 // miscount itself rather than letting it hide in release.
1028 debug_assert!(
1029 uring.in_flight > 0,
1030 "CQE reaped with no operation recorded in flight"
1031 );
1032 uring.in_flight = uring.in_flight.saturating_sub(1);
1033 if cqe_count < cqes.len() {
1034 cqes[cqe_count] = (cqe.user_data(), cqe.result());
1035 cqe_count += 1;
1036 }
1037 }
1038
1039 let any_cqe = cqe_count > 0;
1040 for &(token, result) in &cqes[..cqe_count] {
1041 waiter.reset();
1042 match token {
1043 TOKEN_RECV => {
1044 if result <= 0 {
1045 return Err(io::Error::other(format!(
1046 "replica disconnected (recv returned {result})"
1047 )));
1048 }
1049 let n = result as usize;
1050 parse_buf.extend_from_slice(&recv_buf[..n]);
1051
1052 drain_ack_frames(&mut parse_buf, |ack| {
1053 // Eviction on violation: returning Err tears the
1054 // connection down; the accept loop's cleanup
1055 // disengages the cursors and frees the slot.
1056 cursors
1057 .record_ack(slot_idx, &ack, sent.get())
1058 .map_err(io::Error::other)?;
1059 metrics.ack_latency_us[slot_idx]
1060 .store(last_send.elapsed().as_micros() as u64, Ordering::Relaxed);
1061 Ok(())
1062 })?;
1063
1064 // Resubmit RECV.
1065 let sqe = opcode::Recv::new(
1066 types::Fixed(0),
1067 recv_buf.as_mut_ptr(),
1068 recv_buf.len() as u32,
1069 )
1070 .build()
1071 .user_data(TOKEN_RECV);
1072 // SAFETY: `recv_buf` is owned by the guard; only
1073 // one RECV is ever in flight (we just consumed the
1074 // previous CQE), and the guard drains this one on
1075 // every exit path. Ring is single-threaded.
1076 unsafe { ring.submission().push(&sqe).expect("SQ full") };
1077 uring.in_flight += 1;
1078 }
1079
1080 TOKEN_SEND => {
1081 if result < 0 {
1082 return Err(io::Error::other(format!("send error (returned {result})")));
1083 }
1084 let sent = result as usize;
1085 send_offset += sent;
1086 if send_offset >= send_buf.len() {
1087 // Fully sent. Measure end-to-end SEND latency: a
1088 // healthy io_uring TCP SEND completes in
1089 // microseconds; > threshold implies the kernel
1090 // waited on cwnd / peer window / retransmit.
1091 if let Some(ts) = send_submit_ts.take() {
1092 let elapsed = ts.elapsed();
1093 if elapsed.as_millis() >= SLOW_SEND_THRESHOLD_MS {
1094 tracing::debug!(
1095 slot = slot_idx,
1096 elapsed_us = elapsed.as_micros() as u64,
1097 bytes = send_buf.len(),
1098 "slow SEND completion"
1099 );
1100 super::log_tcp_info(tcp_fd, "slow_send", slot_idx);
1101 }
1102 }
1103 metrics.bytes_sent[slot_idx]
1104 .fetch_add(send_buf.len() as u64, Ordering::Relaxed);
1105 send_buf.clear();
1106 send_offset = 0;
1107 send_in_flight = false;
1108 } else {
1109 // Partial send — resubmit remainder.
1110 let sqe = opcode::Send::new(
1111 types::Fixed(0),
1112 send_buf[send_offset..].as_ptr(),
1113 (send_buf.len() - send_offset) as u32,
1114 )
1115 .build()
1116 .user_data(TOKEN_SEND);
1117 // SAFETY: `send_buf` is owned and still pinned
1118 // (`send_in_flight` stays true across partial
1119 // sends), and the guard drains the CQE on every
1120 // exit path; ring is single-threaded.
1121 unsafe { ring.submission().push(&sqe).expect("SQ full") };
1122 uring.in_flight += 1;
1123 }
1124 }
1125
1126 _ => {}
1127 }
1128 }
1129
1130 // --- Idle wait ---
1131 if !any_cqe && send_buf.is_empty() {
1132 waiter.idle();
1133 }
1134 }
1135}
1136
1137/// Split the replica's complete frames off the front of `parse_buf`, hand
1138/// every `Ack` among them to `on_ack`, and compact the consumed bytes
1139/// away; a trailing partial frame stays for the next RECV.
1140///
1141/// A length prefix above [`MAX_CONTROL_FRAME`], or of zero (every replica
1142/// message carries a tag), is an `Err` as soon as its 4 bytes are in, as
1143/// is an `Err` from `on_ack`; either ends the connection, so nothing is
1144/// compacted on that path. A frame that is not a decodable `Ack` is
1145/// skipped.
1146///
1147/// Not [`super::receiver_transport::try_extract_frame`], which applies
1148/// the same rule: the error names the length the replica declared.
1149fn drain_ack_frames(
1150 parse_buf: &mut Vec<u8>,
1151 mut on_ack: impl FnMut(Ack) -> io::Result<()>,
1152) -> io::Result<()> {
1153 let mut cursor = 0;
1154 loop {
1155 let frame = match split_frame_limited(&parse_buf[cursor..], MAX_CONTROL_FRAME) {
1156 Ok(Some(frame)) if !frame.payload.is_empty() => frame,
1157 Ok(Some(_)) => return Err(io::Error::other("invalid ack frame length: 0")),
1158 Ok(None) => break,
1159 Err(too_large) => {
1160 return Err(io::Error::other(format!(
1161 "invalid ack frame length: {}",
1162 too_large.declared
1163 )));
1164 }
1165 };
1166 if let Ok(ReplicaMessage::Ack(ack)) = decode_replica_message(frame.payload) {
1167 on_ack(ack)?;
1168 }
1169 cursor += frame.consumed();
1170 }
1171 if cursor > 0 {
1172 let remaining = parse_buf.len() - cursor;
1173 parse_buf.copy_within(cursor.., 0);
1174 parse_buf.truncate(remaining);
1175 }
1176 Ok(())
1177}
1178
1179#[cfg(test)]
1180mod tests {
1181 use super::*;
1182 use crate::uring_teardown::DRAIN_TIMEOUT;
1183 use io_uring::{IoUring, opcode, types};
1184
1185 /// Connected AF_UNIX stream pair standing in for the replication
1186 /// socket. Raw fds rather than `UnixStream` because the guard takes
1187 /// a `RawFd` and the test needs to outlive it.
1188 fn socketpair() -> (libc::c_int, libc::c_int) {
1189 let mut fds = [0 as libc::c_int; 2];
1190 // SAFETY: `fds` is a valid two-element array for the kernel to
1191 // fill in; the fds are owned by the caller from here.
1192 let rc = unsafe { libc::socketpair(libc::AF_UNIX, libc::SOCK_STREAM, 0, fds.as_mut_ptr()) };
1193 assert_eq!(rc, 0, "socketpair: {}", io::Error::last_os_error());
1194 (fds[0], fds[1])
1195 }
1196
1197 /// SAFETY: both fds are owned by the caller and not used again.
1198 fn close_pair(a: libc::c_int, b: libc::c_int) {
1199 // Return value dropped deliberately: nothing can be done about
1200 // a failing close on a test teardown path.
1201 unsafe {
1202 libc::close(a);
1203 libc::close(b);
1204 }
1205 }
1206
1207 /// Guard over `fd` with a RECV already submitted.
1208 ///
1209 /// Panics rather than skipping if the host denies io_uring:
1210 /// [`live_stream_uring`] is the only live-stream path there is, so
1211 /// a host that cannot build this ring cannot replicate at all. A
1212 /// test that quietly passed there would report coverage for a
1213 /// teardown it never ran.
1214 fn guard_with_pending_recv(fd: libc::c_int, send_buf: &mut Vec<u8>) -> InflightUring<'_> {
1215 let ring = IoUring::builder()
1216 .setup_single_issuer()
1217 .build(8)
1218 .expect("io_uring is a hard requirement for replication");
1219 ring.submitter()
1220 .register_files(&[fd])
1221 .expect("register_files");
1222 let mut uring = InflightUring {
1223 ring,
1224 recv_buf: vec![0u8; 4096],
1225 send_buf,
1226 in_flight: 0,
1227 tcp_fd: fd,
1228 };
1229 let sqe = opcode::Recv::new(
1230 types::Fixed(0),
1231 uring.recv_buf.as_mut_ptr(),
1232 uring.recv_buf.len() as u32,
1233 )
1234 .build()
1235 .user_data(0);
1236 // SAFETY: `recv_buf` belongs to the guard being built here, and
1237 // the guard drains this operation before releasing it. The ring
1238 // is only ever touched from this thread.
1239 unsafe { uring.ring.submission().push(&sqe).expect("SQ full") };
1240 uring.in_flight = 1;
1241 uring.ring.submit().expect("submit initial RECV");
1242 uring
1243 }
1244
1245 /// The assumption the whole teardown design rests on: a
1246 /// `shutdown(2)` — not an AsyncCancel, which some hosts filter —
1247 /// is enough to make the kernel hand back a RECV that nobody is
1248 /// ever going to write to.
1249 #[test]
1250 fn drain_reaps_a_pending_recv_after_shutdown() {
1251 let (a, b) = socketpair();
1252 let mut send_buf = Vec::new();
1253 let mut uring = guard_with_pending_recv(a, &mut send_buf);
1254 // The peer never writes, so the only thing that can complete
1255 // this RECV is the shutdown.
1256 unsafe { libc::shutdown(a, libc::SHUT_RDWR) };
1257 assert!(
1258 uring.drain(),
1259 "drain hit the deadline with a RECV in flight"
1260 );
1261 assert_eq!(uring.in_flight, 0, "kernel still owes a CQE");
1262 drop(uring);
1263 close_pair(a, b);
1264 }
1265
1266 /// `Drop` must quiesce the ring on its own — including issuing the
1267 /// wake — and must return promptly rather than burning the whole
1268 /// [`DRAIN_TIMEOUT`], which is what a teardown that silently failed
1269 /// to wake the operation would look like.
1270 #[test]
1271 fn drop_quiesces_a_pending_recv_without_hanging() {
1272 let (a, b) = socketpair();
1273 let mut send_buf = vec![0u8; 64];
1274 let uring = guard_with_pending_recv(a, &mut send_buf);
1275 let started = std::time::Instant::now();
1276 drop(uring);
1277 let elapsed = started.elapsed();
1278 assert!(
1279 elapsed < DRAIN_TIMEOUT / 4,
1280 "teardown took {elapsed:?}, close to the {DRAIN_TIMEOUT:?} deadline"
1281 );
1282 // A clean drain releases the buffers; only a drain that cannot
1283 // prove the kernel is done leaks them.
1284 assert_eq!(send_buf.len(), 64, "send_buf leaked despite a clean drain");
1285 close_pair(a, b);
1286 }
1287
1288 /// An eviction is counted once, however many supervisor passes see
1289 /// the flag before the handler exits, and counted again for the next
1290 /// handler once the supervisor has re-armed the slot.
1291 #[test]
1292 fn eviction_is_counted_once_per_handler() {
1293 let mut counted = false;
1294 // No flag, or no handler to evict: nothing to count.
1295 assert!(!take_new_eviction(false, true, &mut counted));
1296 assert!(!take_new_eviction(true, false, &mut counted));
1297 assert!(!counted);
1298
1299 // First pass with the flag set counts; later passes, while the
1300 // handler is still draining, do not.
1301 assert!(take_new_eviction(true, true, &mut counted));
1302 for _ in 0..5 {
1303 assert!(!take_new_eviction(true, true, &mut counted));
1304 }
1305
1306 // The supervisor collects the handler and re-arms the slot.
1307 counted = false;
1308 assert!(take_new_eviction(true, true, &mut counted));
1309 }
1310
1311 // --- drain_ack_frames ---
1312
1313 fn ack_frame(seq: u64) -> Vec<u8> {
1314 let mut buf = Vec::new();
1315 melin_transport_core::replication::protocol::encode_ack(
1316 &Ack {
1317 acked_sequence: seq,
1318 in_memory_sequence: seq,
1319 },
1320 &mut buf,
1321 );
1322 buf
1323 }
1324
1325 /// `drain_ack_frames` over `buf`, returning the acked sequences it
1326 /// handed over alongside its result.
1327 fn drain(buf: &mut Vec<u8>) -> (io::Result<()>, Vec<u64>) {
1328 let mut seen = Vec::new();
1329 let result = drain_ack_frames(buf, |ack| {
1330 seen.push(ack.acked_sequence);
1331 Ok(())
1332 });
1333 (result, seen)
1334 }
1335
1336 #[test]
1337 fn ack_frames_are_handed_over_and_a_partial_one_is_kept() {
1338 let mut buf = ack_frame(1);
1339 buf.extend_from_slice(&ack_frame(2));
1340 let third = ack_frame(3);
1341 buf.extend_from_slice(&third[..third.len() - 1]);
1342
1343 let (result, seen) = drain(&mut buf);
1344 result.unwrap();
1345 assert_eq!(seen, [1, 2]);
1346 assert_eq!(
1347 buf,
1348 third[..third.len() - 1],
1349 "partial frame compacted to the front"
1350 );
1351
1352 buf.push(*third.last().unwrap());
1353 let (result, seen) = drain(&mut buf);
1354 result.unwrap();
1355 assert_eq!(seen, [3]);
1356 assert!(buf.is_empty());
1357 }
1358
1359 #[test]
1360 fn a_frame_that_is_not_an_ack_is_skipped() {
1361 let mut buf = vec![1, 0, 0, 0, 0xEE];
1362 buf.extend_from_slice(&ack_frame(4));
1363 let (result, seen) = drain(&mut buf);
1364 result.unwrap();
1365 assert_eq!(seen, [4]);
1366 assert!(buf.is_empty());
1367 }
1368
1369 #[test]
1370 fn a_frame_at_the_cap_waits_for_its_payload() {
1371 let mut buf = (MAX_CONTROL_FRAME as u32).to_le_bytes().to_vec();
1372 let (result, seen) = drain(&mut buf);
1373 result.unwrap();
1374 assert!(seen.is_empty());
1375 assert_eq!(buf.len(), 4, "prefix kept for the next RECV");
1376 }
1377
1378 /// Refused on the prefix alone, after the acks in front of it.
1379 #[test]
1380 fn an_oversized_prefix_fails_before_its_payload() {
1381 let mut buf = ack_frame(5);
1382 buf.extend_from_slice(&((MAX_CONTROL_FRAME as u32) + 1).to_le_bytes());
1383 let (result, seen) = drain(&mut buf);
1384 let err = result.unwrap_err();
1385 assert_eq!(
1386 err.to_string(),
1387 format!("invalid ack frame length: {}", MAX_CONTROL_FRAME + 1)
1388 );
1389 assert_eq!(seen, [5]);
1390 }
1391
1392 #[test]
1393 fn a_zero_length_frame_fails() {
1394 let mut buf = 0u32.to_le_bytes().to_vec();
1395 let (result, seen) = drain(&mut buf);
1396 assert_eq!(
1397 result.unwrap_err().to_string(),
1398 "invalid ack frame length: 0"
1399 );
1400 assert!(seen.is_empty());
1401 }
1402
1403 #[test]
1404 fn an_on_ack_error_stops_the_drain() {
1405 let mut buf = ack_frame(6);
1406 buf.extend_from_slice(&ack_frame(7));
1407 let mut calls = 0;
1408 let result = drain_ack_frames(&mut buf, |_| {
1409 calls += 1;
1410 Err(io::Error::other("cursor violation"))
1411 });
1412 assert_eq!(result.unwrap_err().to_string(), "cursor violation");
1413 assert_eq!(calls, 1);
1414 }
1415}