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