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::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}