Skip to main content

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}