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