melin_server_runtime/replication/

tcp_sender.rs

//! TCP replication sender (primary side).
//!
//! Listens for replica connections, handles authentication, journal
//! catch-up, snapshot transfer, and live streaming. The wire protocol
//! is defined in `super::protocol`.

use std::io::{self, Write};
use std::net::{TcpListener, TcpStream};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};

use tracing::{debug, error, info, warn};

use melin_journal::replication::ReplicationConsumer;

use super::auth::authenticate_replica;
use super::{ReplicaCursors, ReplicaGate, ReplicationMetrics, SentHighWater};
use crate::uring_teardown::{DrainBackoff, wake_pending_ops};
use melin_app::Application;
use melin_transport_core::replication::catchup::{
    CatchUpResult, bridge_catchup_to_live, can_catch_up_from_journal, catch_up_from_journal,
    preflight_snapshot_transfer, snapshot_transfer_with,
};
use melin_transport_core::replication::protocol::{
    MAX_CONTROL_FRAME, ReplicaMessage, decode_replica_message, encode_hash_mismatch,
    encode_heartbeat, encode_need_snapshot, encode_stream_start, read_frame,
};
use melin_transport_core::replication::validate::{
    HandshakeValidation, validate_replica_handshake_settled,
};

// --- Replication Sender (Primary side) ---

/// Owned state for the replication sender thread.
/// A TCP listener guaranteed non-blocking: the only constructor sets
/// the flag, so the guarantee is carried by the type. [`run_sender`]'s
/// accept loop polls the shutdown flag between accepts — handed a
/// blocking listener it would park inside `accept()` and only observe
/// shutdown when the next replica happened to dial in. This newtype
/// makes that state unrepresentable instead of documenting it away.
pub struct ReplicationListener(TcpListener);

impl ReplicationListener {
    /// Wrap a bound listener, setting it non-blocking.
    pub fn new(listener: TcpListener) -> std::io::Result<Self> {
        listener.set_nonblocking(true)?;
        Ok(Self(listener))
    }
}

pub struct Sender {
    /// Pre-bound replication listener, non-blocking by construction.
    /// Bound by the caller at startup — not here — so a bad or stolen
    /// `--replication-bind` port fails the boot with a clear error.
    /// Binding on this thread meant a failure only killed the sender
    /// silently, leaving a primary that looked healthy but could never
    /// satisfy hybrid/durably-replicated acks.
    pub listener: ReplicationListener,
    pub repl_consumer_1: ReplicationConsumer,
    pub repl_consumer_2: ReplicationConsumer,
    pub replica_slots: Arc<melin_transport_core::ReplicaSlotCursors>,
    pub journal_path: std::path::PathBuf,
    pub authorized_keys: Arc<melin_app::auth::AuthorizedKeys>,
    pub evict_flags: [Arc<AtomicBool>; 2],
    pub active_flags: [Arc<AtomicBool>; 2],
    pub metrics: Arc<ReplicationMetrics>,
    pub handler_cores: [usize; 2],
    pub batch_size: usize,
    pub heartbeat_secs: u64,
    pub busy_spin: bool,
    /// Node fencing state. Read to stamp the primary's epoch onto each
    /// `StreamStart`, and to self-demote when a replica handshakes with a
    /// higher epoch (this primary has been superseded). See `crate::fence`.
    pub fence_state: Arc<melin_transport_core::fence::FenceState>,
    /// The active durability (acking) mode, shared with the response
    /// gate and the admin `DURABILITY` command. Stamped on every
    /// `StreamStart` and `Heartbeat` so replicas judge auto-promotion
    /// against the mode this primary actually acks under.
    pub durability_mode: Arc<std::sync::atomic::AtomicU8>,
}

/// Run the replication sender. Listens for replica connections,
/// streams journal data batches, processes acks, and updates the
/// replication cursor.
///
/// Runs on a dedicated thread. Blocks until shutdown.
pub fn run_sender<A: Application>(
    config: Sender,
    shutdown: &AtomicBool,
    replica_ready: &AtomicBool,
    replicas_connected: &AtomicU32,
) {
    let Sender {
        listener,
        repl_consumer_1,
        repl_consumer_2,
        replica_slots,
        journal_path,
        authorized_keys,
        evict_flags,
        active_flags,
        metrics,
        handler_cores,
        batch_size,
        heartbeat_secs,
        busy_spin,
        fence_state,
        durability_mode,
    } = config;
    // Unwrap the invariant-carrying newtype — from here `listener` is a
    // plain TcpListener, known non-blocking.
    let ReplicationListener(listener) = listener;
    match listener.local_addr() {
        Ok(addr) => info!(%addr, "replication sender listening"),
        // Best-effort log detail — the pre-bound listener works either way.
        Err(_) => info!("replication sender listening"),
    }

    // Two replica slots, each with its own ring consumer and thread handle.
    // The accept loop fills empty slots. When a replica disconnects, its
    // slot becomes available for a new connection. The consumer is `None`
    // while a handler thread owns it (moved into the thread, returned on exit).
    struct ReplicaSlot {
        consumer: Option<ReplicationConsumer>,
        handle: Option<std::thread::JoinHandle<ReplicationConsumer>>,
    }

    let mut slots = [
        ReplicaSlot {
            consumer: Some(repl_consumer_1),
            handle: None,
        },
        ReplicaSlot {
            consumer: Some(repl_consumer_2),
            handle: None,
        },
    ];

    // Single owner of the per-replica progress cursors (per-slot acked
    // positions and the gate's gauge pair). Shared by both handler
    // threads — see `ReplicaCursors` for the ordering contract relative
    // to the active flags.
    let cursors = Arc::new(ReplicaCursors::new(replica_slots, Arc::clone(&metrics)));

    // Per-slot "this connection authenticated" latch. The trading-halt gate
    // (`replicas_connected`) is lifted by the handler thread only after auth
    // succeeds, so the main loop must lower it on teardown only when the
    // handler got that far — a connection that drops mid-auth never counted.
    // Per-slot `Arc<AtomicBool>` mirrors `active_flags`/`evict_flags`; shared
    // with the handler thread, which flips it true post-auth.
    let authenticated_flags: [Arc<AtomicBool>; 2] = [
        Arc::new(AtomicBool::new(false)),
        Arc::new(AtomicBool::new(false)),
    ];

    loop {
        if shutdown.load(Ordering::Relaxed) {
            info!("replication sender shutting down");
            // Wait for active replica threads to finish.
            for slot in &mut slots {
                if let Some(handle) = slot.handle.take()
                    && let Ok(consumer) = handle.join()
                {
                    slot.consumer = Some(consumer);
                }
            }
            return;
        }

        // Check eviction flags from the journal stage. When set, the
        // journal stage timed out publishing to this slot's ring. We need
        // to reclaim the consumer so the idle drain loop can clear the ring,
        // allowing the journal stage to resume publishing.
        for (i, slot) in slots.iter_mut().enumerate() {
            if evict_flags[i].load(Ordering::Acquire) && slot.handle.is_some() {
                metrics.evictions_total.fetch_add(1, Ordering::Relaxed);
                warn!(
                    slot = i,
                    "evicting slow replica (ring backpressure timeout)"
                );
                // The handler thread checks shutdown — we can't signal it
                // individually without adding per-slot flags. Instead, join
                // the thread with a short timeout by checking is_finished.
                // The handler's TCP read timeout (5s) will cause it to exit
                // on the next iteration when it checks shutdown. But we
                // want faster eviction, so we shutdown the TCP stream to
                // unblock the read.
                //
                // For now, mark the slot and let it be collected below
                // when the handler finishes naturally (TCP timeout or
                // next send failure after the ring stops being fed).
            }
        }

        // Disengage a finished slot's shared state. Both the clean-exit and
        // panic arms below call this, so they cannot drift: a panicked handler
        // must tear down exactly like a clean exit, or it leaks the trading-halt
        // gate and — far worse — leaves the slot's progress cursors engaged,
        // so the primary stops acking client requests even with a healthy
        // surviving replica. Borrows only atomics (interior
        // mutability), so it coexists with the `&mut` iteration over `slots`.
        let disengage_slot = |slot_idx: usize| {
            // Lower the gate only if this connection authenticated (and so
            // lifted it); `swap` reads and resets the latch in one op, and the
            // `join()` ordered the handler's post-auth writes before here.
            // `lower` warns "trading halted" if this was the last replica.
            if authenticated_flags[slot_idx].swap(false, Ordering::AcqRel) {
                ReplicaGate::new(replicas_connected).lower();
            }
            // Disengage cursors BEFORE clearing the active flag — ordering
            // contract B2 (see `ReplicaCursors`).
            cursors.clear_on_disconnect(slot_idx);
            metrics.catching_up[slot_idx].store(false, Ordering::Relaxed);
            // Journal stage stops publishing to this ring.
            active_flags[slot_idx].store(false, Ordering::Release);
        };

        // Collect finished replica threads (disconnected replicas).
        for (i, slot) in slots.iter_mut().enumerate() {
            if let Some(ref handle) = slot.handle
                && handle.is_finished()
            {
                let handle = slot.handle.take().expect("just checked is_some");
                match handle.join() {
                    Ok(mut consumer) => {
                        // Drop any unread entries before the consumer is
                        // stashed back in the slot. The journal stage may
                        // have published batches into the ring that the
                        // evicted handler never got to forward — if we
                        // left them in place, the NEXT handler on this
                        // slot would drain them to its replica and
                        // acknowledge with pre-eviction sequences. Those
                        // acks would stall the slot's acked-progress
                        // cursor at the old position and gate the
                        // primary's response stage at the slow-replica
                        // rate.
                        //
                        // Fast-forward to the producer cursor so the
                        // live-streaming loop starts with a clean ring.
                        consumer.skip_to_producer();
                        slot.consumer = Some(consumer);
                        disengage_slot(i);
                        // Clear eviction flag after reclaiming the consumer.
                        if evict_flags[i].load(Ordering::Relaxed) {
                            evict_flags[i].store(false, Ordering::Release);
                            warn!(slot = i, "evicted replica — ring ready for reconnection");
                        } else {
                            warn!(slot = i, "replica disconnected");
                        }
                    }
                    Err(_) => {
                        error!(slot = i, "replica handler thread panicked");
                        // Consumer is lost — the slot stays handle=None /
                        // consumer=None, so the accept loop skips it. With
                        // independent rings, only this slot's ring is affected;
                        // the other replica continues normally. Still disengage
                        // its shared state, or the gate leaks and the cursor
                        // freezes (see `disengage_slot`).
                        disengage_slot(i);
                        evict_flags[i].store(false, Ordering::Release);
                    }
                }
            }
        }

        // Find a slot that has a consumer available (not in use by a thread).
        let empty_slot = slots
            .iter()
            .position(|s| s.handle.is_none() && s.consumer.is_some());

        // Accept a connection if there's an empty slot.
        if let Some(slot_idx) = empty_slot {
            match listener.accept() {
                Ok((stream, addr)) => {
                    info!(addr = %addr, slot = slot_idx, "replica connected");
                    if let Err(e) = stream.set_nodelay(true) {
                        debug!(error = %e, "failed to set TCP_NODELAY on replica connection");
                    }
                    // SO_BUSY_POLL on the sender side: the per-replica thread
                    // spins on `recv` for ack frames, so kernel busy-polling
                    // removes the softirq->wakeup handoff from the ack path.
                    if let Err(e) =
                        crate::server::set_busy_poll(&stream, crate::server::BUSY_POLL_US)
                    {
                        debug!(error = %e, "failed to set SO_BUSY_POLL on replica connection");
                    }

                    // The trading-halt gate is NOT lifted here — a bare connect
                    // hasn't proven a Replication key. The handler lifts it after
                    // auth succeeds (see `handle_replica_connection`).

                    // Take the consumer out of the slot for the handler thread.
                    // The slot's consumer becomes None while the thread owns it.
                    let consumer = slots[slot_idx]
                        .consumer
                        .take()
                        .expect("empty_slot check guarantees consumer is Some");

                    let slot_cursors = Arc::clone(&cursors);
                    let jpath = journal_path.clone();
                    let auth_keys = Arc::clone(&authorized_keys);
                    let slot_metrics = Arc::clone(&metrics);
                    let slot_active = Arc::clone(&active_flags[slot_idx]);
                    let slot_evict = Arc::clone(&evict_flags[slot_idx]);
                    let slot_fence = Arc::clone(&fence_state);
                    let slot_durability = Arc::clone(&durability_mode);
                    let slot_authenticated = Arc::clone(&authenticated_flags[slot_idx]);
                    let handler_core = handler_cores[slot_idx];
                    let shutdown_flag = shutdown as *const AtomicBool as usize;
                    let ready_flag = replica_ready as *const AtomicBool as usize;
                    let connected_flag = replicas_connected as *const AtomicU32 as usize;
                    let handle = std::thread::Builder::new()
                        .name(format!("repl-{slot_idx}"))
                        .spawn(move || {
                            // Pin to a dedicated core if configured (> 0),
                            // otherwise clear inherited affinity from the
                            // sender thread so the OS can schedule freely.
                            if handler_core > 0 {
                                match melin_app::affinity::pin_to_core(handler_core) {
                                    Ok(c) => tracing::info!(
                                        core = c,
                                        slot = slot_idx,
                                        "replication handler pinned to core"
                                    ),
                                    Err(e) => tracing::warn!(
                                        error = e,
                                        slot = slot_idx,
                                        "failed to pin handler"
                                    ),
                                }
                            } else if let Err(e) = melin_app::affinity::clear_affinity() {
                                tracing::warn!(error = e, "failed to clear handler affinity");
                            }
                            // Safety: shutdown and replica_ready outlive this thread
                            // (they're on the parent's stack, which blocks on join
                            // during shutdown).
                            let shutdown_ref = unsafe { &*(shutdown_flag as *const AtomicBool) };
                            let ready_ref = unsafe { &*(ready_flag as *const AtomicBool) };
                            // Safety: replicas_connected outlives this thread —
                            // same lifetime argument as shutdown/replica_ready
                            // above (the parent blocks on join during shutdown).
                            let connected_ref = unsafe { &*(connected_flag as *const AtomicU32) };
                            let ctx = SlotContext {
                                cursors: &slot_cursors,
                                journal_path: &jpath,
                                authorized_keys: &auth_keys,
                                shutdown: shutdown_ref,
                                replica_ready: ready_ref,
                                active_flag: &slot_active,
                                evict_flag: &slot_evict,
                                metrics: &slot_metrics,
                                fence_state: &slot_fence,
                                durability_mode: &slot_durability,
                                slot_idx,
                                batch_size,
                                heartbeat_secs,
                                busy_spin,
                                replicas_connected: connected_ref,
                                authenticated: &slot_authenticated,
                            };
                            run_replica_slot::<A>(stream, consumer, &ctx)
                        })
                        .expect("spawn replica handler thread");
                    slots[slot_idx].handle = Some(handle);
                }
                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                    // No pending connection.
                }
                Err(e) => {
                    error!(error = %e, "replication accept error");
                }
            }
        }

        // No idle drain needed — the journal stage only publishes to
        // rings where active_flag is true (set by handler on live loop
        // entry, cleared on disconnect). Idle consumers stay empty.

        std::thread::sleep(std::time::Duration::from_millis(50));
    }
}

/// Per-slot state shared across the replica handler call chain
/// (`run_replica_slot` → `handle_replica_connection` → `live_stream_uring`).
struct SlotContext<'a> {
    cursors: &'a Arc<ReplicaCursors>,
    journal_path: &'a std::path::Path,
    authorized_keys: &'a melin_app::auth::AuthorizedKeys,
    shutdown: &'a AtomicBool,
    replica_ready: &'a AtomicBool,
    active_flag: &'a AtomicBool,
    evict_flag: &'a AtomicBool,
    metrics: &'a ReplicationMetrics,
    fence_state: &'a melin_transport_core::fence::FenceState,
    /// The primary's active acking mode — stamped on `StreamStart` and
    /// `Heartbeat` (see `Sender::durability_mode`).
    durability_mode: &'a std::sync::atomic::AtomicU8,
    slot_idx: usize,
    batch_size: usize,
    heartbeat_secs: u64,
    busy_spin: bool,
    /// Trading-halt gate. The handler lifts it (`fetch_add`) only after auth
    /// succeeds; the main loop lowers it on teardown, gated by `authenticated`.
    replicas_connected: &'a AtomicU32,
    /// Per-slot latch the handler sets true once this connection passes auth,
    /// so the main loop knows whether this slot lifted the gate.
    authenticated: &'a AtomicBool,
}

/// Handle a single replica connection on a dedicated thread.
/// Returns the consumer when the connection ends (for slot reuse).
fn run_replica_slot<A: Application>(
    stream: TcpStream,
    mut consumer: ReplicationConsumer,
    ctx: &SlotContext<'_>,
) -> ReplicationConsumer {
    match handle_replica_connection::<A>(stream, &mut consumer, ctx) {
        Ok(()) => info!("replica disconnected cleanly"),
        Err(e) => warn!(error = %e, "replica connection error"),
    }
    consumer
}

fn handle_replica_connection<A: Application>(
    stream: TcpStream,
    repl_consumer: &mut ReplicationConsumer,
    ctx: &SlotContext<'_>,
) -> io::Result<()> {
    let SlotContext {
        cursors,
        journal_path,
        authorized_keys,
        shutdown,
        replica_ready,
        active_flag,
        evict_flag: _,
        metrics,
        fence_state,
        durability_mode,
        slot_idx,
        batch_size: _,
        heartbeat_secs,
        busy_spin: _,
        replicas_connected,
        authenticated,
    } = ctx;
    let slot_idx = *slot_idx;
    let heartbeat_secs = *heartbeat_secs;

    let mut reader = stream.try_clone()?;
    let mut writer = stream;

    // Set a read timeout for the handshake and auth.
    reader.set_read_timeout(Some(std::time::Duration::from_secs(10)))?;

    // Authenticate before any data exchange. `reader` and `writer` are clones
    // of the same socket; auth is sequential (write challenge, read response,
    // write result) so a single handle suffices. Use `reader` — it carries the
    // read timeout set above.
    authenticate_replica(&mut reader, authorized_keys)?;
    info!("replica authenticated");

    // Auth passed — only now lift the trading-halt gate. The main loop lowers
    // it on slot teardown, reading `authenticated` to know this connection
    // lifted it. Set the latch after the increment so a post-join read of the
    // latch implies the increment is visible too.
    ReplicaGate::new(replicas_connected).lift();
    authenticated.store(true, Ordering::Release);

    // Read handshake.
    let handshake_frame = read_frame(&mut reader, MAX_CONTROL_FRAME)?;
    let handshake = match decode_replica_message(&handshake_frame)? {
        ReplicaMessage::Handshake(h) => h,
        ReplicaMessage::Ack(_) => {
            return Err(io::Error::other("expected Handshake, got Ack"));
        }
    };

    info!(
        last_sequence = handshake.last_sequence,
        replica_epoch = handshake.epoch,
        "replica handshake received"
    );

    // Fence: a replica advertising an epoch higher than ours means a
    // promotion happened that this node missed — we are a stale ex-primary.
    // The policy (latch fenced + trigger shutdown) lives on `FenceState` so
    // the kernel-TCP and DPDK senders cannot drift; the matching stage
    // halts new client writes, the response gate stops acking, and the node
    // winds down for the operator to restart it as a replica. Refuse this
    // connection so we send no stale stream.
    let our_epoch = fence_state.epoch();
    if let Some(first_latch) = fence_state.fence_if_superseded(handshake.epoch, shutdown) {
        if first_latch {
            error!(
                replica_epoch = handshake.epoch,
                our_epoch,
                "fenced: a replica advertises a higher epoch — this primary has been \
                 superseded; self-demoting and shutting down"
            );
        }
        return Err(io::Error::other("fenced by higher-epoch replica"));
    }

    // Mark this slot as catching up. Cleared when entering the live loop.
    metrics.catching_up[slot_idx].store(true, Ordering::Relaxed);

    let mut send_buf = Vec::with_capacity(128);

    // Chain validation: recompute our chain at the replica's claimed
    // position and compare with the hash it reported. A mismatch (or a
    // position beyond our tip) means its journal holds divergent
    // history — streaming on top would silently fork the audit trail.
    // Send HashMismatch (the replica archives its local journal on
    // receipt) and route it through snapshot resync on this same
    // connection.
    let divergent = match validate_replica_handshake_settled(journal_path, &handshake)? {
        HandshakeValidation::Ok => false,
        HandshakeValidation::Divergent(kind) => {
            // Alertable on /metrics: divergence outside an expected
            // failover rejoin means corruption or a serious bug.
            metrics.divergence_total.fetch_add(1, Ordering::Relaxed);
            warn!(
                last_sequence = handshake.last_sequence,
                ?kind,
                "replica journal divergent at handshake — routing through snapshot resync"
            );
            true
        }
    };

    // Probe whether journal catch-up is possible before committing to
    // a protocol path. This avoids sending StreamStart only to discover
    // the journals are too old.
    let can_catch_up =
        !divergent && can_catch_up_from_journal(journal_path, handshake.last_sequence)?;

    let mut publish = |buf: &[u8]| -> io::Result<()> {
        writer.write_all(buf)?;
        writer.flush()
    };

    // Cursor/stream floor for everything downstream. A divergent
    // replica's claimed position is meaningless (its history forked —
    // possibly ahead of our tip); the resync rebases it, so seed from 0
    // like a fresh replica: the ack cursors understate until real acks
    // arrive (safe), and the sent high-water seeds from the catch-up
    // end instead of a bogus claim that would mark live ring chunks as
    // already covered.
    let stream_base = if divergent {
        0
    } else {
        handshake.last_sequence
    };

    let catchup_end = if can_catch_up {
        // The lineage origin (oldest segment's header identity) lets a
        // fresh replica create a byte-identical journal before
        // consuming the stream. Replicas with local state ignore it.
        let (lineage_start, lineage_anchor) =
            melin_transport_core::replication::catchup::lineage_origin(journal_path)?;
        encode_stream_start(
            handshake.last_sequence,
            lineage_start,
            lineage_anchor,
            fence_state.epoch(),
            durability_mode.load(Ordering::Relaxed),
            &mut send_buf,
        );
        publish(&send_buf)?;
        send_buf.clear();

        match catch_up_from_journal::<A::Event>(
            journal_path,
            handshake.last_sequence,
            &mut writer,
            shutdown,
        )? {
            CatchUpResult::Ok(end) => end,
            CatchUpResult::NeedSnapshot => {
                return Err(io::Error::other("catch-up failed unexpectedly after probe"));
            }
        }
    } else {
        // Resync verdict precedes the snapshot data: `HashMismatch`
        // tells the replica its journal is divergent (it archives its
        // local lineage on receipt); plain `NeedSnapshot` is the
        // too-far-behind rebase. The receiver expects `SnapshotBegin`
        // as the very next frame after the verdict, so
        // `snapshot_transfer_with` must follow immediately.
        //
        // Pre-flight before the verdict goes on the wire: the replica
        // archives its lineage on receipt, so promising a snapshot we
        // cannot produce (none configured, boundary segment pruned)
        // would strand it with its history moved aside. Failing here
        // drops the connection with the replica's journal intact.
        preflight_snapshot_transfer(journal_path)?;
        if divergent {
            encode_hash_mismatch(&mut send_buf);
        } else {
            encode_need_snapshot(&mut send_buf);
        }
        publish(&send_buf)?;
        send_buf.clear();
        match snapshot_transfer_with::<A::Event>(
            journal_path,
            &mut publish,
            shutdown,
            durability_mode.load(Ordering::Relaxed),
        )? {
            CatchUpResult::Ok(end) => end,
            CatchUpResult::NeedSnapshot => {
                return Err(io::Error::other(
                    "catch-up failed even after snapshot transfer",
                ));
            }
        }
    };

    // Engage this slot's cursors and seed the gauge pair the response
    // gate's `evaluate_durability` reads. Must happen BEFORE the
    // bridge's active_flag Release — see `ReplicaCursors` for the
    // ordering contract.
    cursors.seed_on_handshake(slot_idx, stream_base);

    // Bridge into live streaming: activates the ring, re-reads from the
    // journal the entries that fell into the activation window, then
    // drains the ring into sequence-contiguity (back-filling from disk
    // if a skipped entry hasn't flushed yet) before going live. Returns
    // the slot's sent high-water — the bound the ack-sanity invariant
    // checks acks against, advanced by every chunk the live loop
    // forwards. The bridge closes the catch-up→live gap under load (the
    // receiver's contiguity gate backstops only the rare quiescent
    // corner). See `bridge_catchup_to_live`.
    let mut sent = {
        let mut publish = |buf: &[u8]| -> io::Result<()> {
            writer.write_all(buf)?;
            writer.flush()
        };
        bridge_catchup_to_live::<A::Event>(
            journal_path,
            stream_base,
            catchup_end,
            active_flag,
            repl_consumer,
            &mut publish,
            shutdown,
        )?
    };

    // Catch-up complete — replica is entering the live streaming loop.
    metrics.catching_up[slot_idx].store(false, Ordering::Relaxed);

    // Signal that this replica is ready to consume from the replication
    // ring. The main thread waits on this before seeding test data.
    // Must happen AFTER catch-up and overlap drain complete — otherwise
    // seeding fills the replication ring faster than we can drain it,
    // deadlocking the journal stage.
    replica_ready.store(true, Ordering::Release);

    let heartbeat_interval = std::time::Duration::from_secs(heartbeat_secs);
    let mut last_send = std::time::Instant::now();

    live_stream_uring(
        writer,
        repl_consumer,
        ctx,
        heartbeat_interval,
        &mut send_buf,
        &mut last_send,
        &mut sent,
    )
}

/// Quiescing owner of the live-stream ring: the ring itself, the RECV
/// landing buffer the kernel writes ack bytes into, the SEND source
/// buffer the kernel reads out of, and the number of operations the
/// kernel still owes CQEs for.
///
/// Exists for its `Drop`. Returning from the loop — shutdown, eviction,
/// an I/O error, or a panic's unwind — while a RECV or SEND is still in
/// flight would release those buffers (and drop the ring) with the
/// kernel holding pointers into them. Ring teardown on fd close is
/// *asynchronous*: the kernel can complete a pending RECV after the
/// allocation has been returned to the allocator, writing replica-ack
/// bytes into freed heap. That corruption surfaced as malloc-metadata
/// segfaults at sender-thread exit in the failover tests, and being a
/// kernel-side store it is invisible to AddressSanitizer. The guard
/// closes every exit path: wake pending operations by shutting the
/// socket down, then reap CQEs until the kernel owes none, and only
/// then let the buffers go.
struct InflightUring<'a> {
    ring: io_uring::IoUring,
    /// RECV landing buffer for ack frames. `Vec` so the heap pointer
    /// handed to the kernel stays stable however the guard itself
    /// moves; sized 4 KiB (acks are 13 bytes, the kernel may coalesce
    /// many).
    recv_buf: Vec<u8>,
    /// SEND source buffer, borrowed from `handle_replica_connection`
    /// (it outlives the live-stream loop so it can be reused across
    /// reconnects). Held by the guard purely so the drain covers it:
    /// the kernel *reads* it for an in-flight SEND, so releasing it
    /// early would splice freed heap into the replication stream.
    send_buf: &'a mut Vec<u8>,
    /// SQEs pushed minus CQEs reaped. Between a push and the following
    /// `submit` the kernel does not yet owe those CQEs, so this is an
    /// upper bound rather than an exact debt; [`Self::drain`] submits
    /// before waiting, which makes the two converge. `u32` matches the
    /// ring-depth scale (at most a RECV and a SEND are ever
    /// outstanding here).
    in_flight: u32,
    /// Raw socket fd, for the `shutdown(2)` wake in `Drop`. Always
    /// outlives the guard: the owning `TcpStream` is a parameter of
    /// [`live_stream_uring`] and parameters drop after locals.
    tcp_fd: std::os::unix::io::RawFd,
}

impl Drop for InflightUring<'_> {
    fn drop(&mut self) {
        wake_pending_ops(self.tcp_fd);
        if !self.drain() {
            // The drain could not be proven complete, so the kernel may
            // still write into `recv_buf` or read out of `send_buf`.
            // Leak both: a few KiB lost on a dying connection is safe,
            // returning memory the kernel still touches to the
            // allocator is not. `take` leaves each `Vec` empty, so the
            // caller's `send_buf` stays a valid (if reallocating) Vec.
            std::mem::forget(std::mem::take(&mut self.recv_buf));
            std::mem::forget(std::mem::take(self.send_buf));
            warn!(
                in_flight = self.in_flight,
                "io_uring teardown drain did not complete; leaking replication buffers"
            );
        }
    }
}

impl InflightUring<'_> {
    /// Reap CQEs until the kernel owes none, bounded by
    /// [`crate::uring_teardown::DRAIN_TIMEOUT`]. `true` means
    /// `in_flight` reached zero and the
    /// buffers are provably free of kernel references; `false` means
    /// the caller must not release them. See [`crate::uring_teardown`]
    /// for why this polls rather than blocking in `submit_and_wait`.
    fn drain(&mut self) -> bool {
        // Flush SQEs that were pushed but never submitted: until they
        // reach the kernel, no CQE is ever coming for them.
        loop {
            match self.ring.submit() {
                Ok(_) => break,
                Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
                Err(_) => return false,
            }
        }
        let mut backoff = DrainBackoff::new();
        while self.in_flight > 0 {
            for _ in self.ring.completion() {
                self.in_flight = self.in_flight.saturating_sub(1);
            }
            if self.in_flight > 0 && !backoff.wait() {
                return false;
            }
        }
        true
    }
}

/// Live streaming loop using async RECV/SEND via io_uring. A single RECV is always
/// in-flight for ack frames; SEND is submitted when the replication ring
/// has data. Both complete via the memory-mapped CQ with zero syscalls
/// in the hot path. In-flight operations and both kernel-visible
/// buffers are owned by an [`InflightUring`] guard, so no exit path can
/// release memory the kernel is still reading or writing.
fn live_stream_uring(
    writer: TcpStream,
    repl_consumer: &mut ReplicationConsumer,
    ctx: &SlotContext<'_>,
    heartbeat_interval: std::time::Duration,
    send_buf: &mut Vec<u8>,
    last_send: &mut std::time::Instant,
    sent: &mut SentHighWater,
) -> io::Result<()> {
    let SlotContext {
        cursors,
        shutdown,
        evict_flag,
        metrics,
        slot_idx,
        batch_size,
        busy_spin,
        durability_mode,
        // Only used during handshake/catch-up (handle_replica_connection).
        journal_path: _,
        authorized_keys: _,
        replica_ready: _,
        active_flag: _,
        heartbeat_secs: _,
        fence_state: _,
        replicas_connected: _,
        authenticated: _,
    } = ctx;
    let slot_idx = *slot_idx;
    let batch_size = *batch_size;
    let busy_spin = *busy_spin;

    use io_uring::{IoUring, opcode, types};
    use std::os::unix::io::AsRawFd;

    const TOKEN_RECV: u64 = 0;
    const TOKEN_SEND: u64 = 1;

    let tcp_fd = writer.as_raw_fd();

    let ring: IoUring = IoUring::builder()
        .setup_single_issuer()
        .build(8)
        .map_err(|e| io::Error::other(format!("io_uring init failed: {e}")))?;

    ring.submitter()
        .register_files(&[tcp_fd])
        .map_err(|e| io::Error::other(format!("io_uring register_files: {e}")))?;

    // Pin io-wq workers to core 0 (keep them off pipeline cores).
    {
        let mut cpuset: libc::cpu_set_t = unsafe { std::mem::zeroed() };
        unsafe { libc::CPU_SET(0, &mut cpuset) };
        let _ = ring.submitter().register_iowq_aff(&cpuset);
    }

    // Ring, both kernel-visible buffers, and the in-flight count live
    // in the quiescing guard from here on: every exit below (shutdown,
    // eviction, `?`, unwind) must drain the kernel before any of these
    // allocations can be released. The loop keeps using them through
    // disjoint field borrows.
    let mut uring = InflightUring {
        ring,
        recv_buf: vec![0u8; 4096],
        send_buf,
        in_flight: 0,
        tcp_fd,
    };
    let (ring, recv_buf, send_buf) = (&mut uring.ring, &mut uring.recv_buf, &mut *uring.send_buf);
    // Accumulation buffer for partial ack frame parsing.
    let mut parse_buf: Vec<u8> = Vec::with_capacity(MAX_CONTROL_FRAME + 4);
    // RECV is always resubmitted after CQE processing — no explicit
    // tracking needed. The io_uring kernel guarantees ordering.
    let mut send_in_flight = false;
    let mut send_offset: usize = 0;
    let mut idle_spins: u32 = 0;
    let mut heartbeat_timer = melin_app::amortized_timer::AmortizedTimer::new();

    // Diagnostic (RUST_LOG=debug): per-slot TCP_INFO snapshot once a
    // second, slow-SEND detection (CQE elapsed >= threshold), and a
    // TCP_INFO capture at the evict-exit point. Amortized so the
    // per-iteration cost is a single `AND` + predictable branch.
    let mut info_log_timer = melin_app::amortized_timer::AmortizedTimer::new();
    let mut send_submit_ts: Option<std::time::Instant> = None;
    const SLOW_SEND_THRESHOLD_MS: u128 = 5;

    // Submit initial RECV.
    let sqe = opcode::Recv::new(
        types::Fixed(0),
        recv_buf.as_mut_ptr(),
        recv_buf.len() as u32,
    )
    .build()
    .user_data(TOKEN_RECV);
    // SAFETY: `recv_buf` is owned by the `InflightUring` guard, which
    // drains this operation's CQE before the buffer can be freed on any
    // exit path. The ring is single-threaded — only this task
    // pushes/reaps on it.
    unsafe { ring.submission().push(&sqe).expect("SQ full") };
    uring.in_flight += 1;

    loop {
        // --- Check flags ---
        if shutdown.load(Ordering::Relaxed) {
            return Ok(());
        }
        if evict_flag.load(Ordering::Relaxed) {
            // Capture the TCP state at the moment of eviction — the
            // critical frame for comparing an evicted slot's teardown
            // against the still-live slot's socket state when
            // diagnosing post-eviction regressions.
            super::log_tcp_info(tcp_fd, "evict_exit", slot_idx);
            info!(slot = slot_idx, "handler exiting: evicted by journal stage");
            return Ok(());
        }

        // --- Drain replication ring into send_buf (memory, non-blocking) ---
        //
        // Ring chunks are wire-ready `InputBatch` frames produced by the
        // journal stage (phase 3 of the unified-pipeline plan), so the
        // sender is a passthrough — no decode + re-encode here.
        if !send_in_flight {
            let mut coalesced = 0;
            while coalesced < batch_size {
                if let Some((meta, data)) = repl_consumer.try_read() {
                    send_buf.extend_from_slice(data);
                    repl_consumer.commit();
                    sent.advance(meta.end_sequence);
                    coalesced += 1;
                } else {
                    break;
                }
            }

            if coalesced > 0 {
                // Submit SEND for the coalesced buffer.
                let sqe =
                    opcode::Send::new(types::Fixed(0), send_buf.as_ptr(), send_buf.len() as u32)
                        .build()
                        .user_data(TOKEN_SEND);
                // SAFETY: `send_buf` is owned by this task and pinned
                // by `send_in_flight = true` below — it isn't mutated
                // until the matching TOKEN_SEND CQE clears the flag,
                // and the `InflightUring` guard drains the CQE before
                // any exit path can release it. The ring is
                // single-threaded.
                unsafe { ring.submission().push(&sqe).expect("SQ full") };
                uring.in_flight += 1;
                send_in_flight = true;
                send_offset = 0;
                *last_send = std::time::Instant::now();
                send_submit_ts = Some(*last_send);
                idle_spins = 0;
                heartbeat_timer = melin_app::amortized_timer::AmortizedTimer::new();
            } else {
                // Heartbeat check: amortized when spinning (mask keeps the
                // clock read at ~10/s at 10M iter/s). In yield mode the loop
                // already pays a syscall per iteration, so the clock read is
                // free and must not be skipped — see AmortizedTimer docs.
                let spinning = busy_spin || idle_spins < 1000;
                if heartbeat_timer.tick(heartbeat_interval, spinning).is_some() {
                    encode_heartbeat(
                        sent.get(),
                        durability_mode.load(Ordering::Relaxed),
                        send_buf,
                    );
                    let sqe = opcode::Send::new(
                        types::Fixed(0),
                        send_buf.as_ptr(),
                        send_buf.len() as u32,
                    )
                    .build()
                    .user_data(TOKEN_SEND);
                    // SAFETY: same as the coalesced send above —
                    // `send_buf` is owned and pinned by `send_in_flight`
                    // until the TOKEN_SEND CQE clears it, and the guard
                    // drains the CQE on every exit path.
                    unsafe { ring.submission().push(&sqe).expect("SQ full") };
                    uring.in_flight += 1;
                    send_in_flight = true;
                    send_offset = 0;
                    *last_send = std::time::Instant::now();
                    send_submit_ts = Some(*last_send);
                }
            }
        }

        // Periodic TCP_INFO dump — debug level. Amortized so the
        // per-iteration cost is a single `AND` + predictable branch.
        if info_log_timer
            .tick(
                std::time::Duration::from_secs(1),
                busy_spin || idle_spins < 1000,
            )
            .is_some()
        {
            super::log_tcp_info(tcp_fd, "live_stream", slot_idx);
        }

        // --- Submit SQEs to kernel (non-blocking) ---
        // Skip the syscall when no new SQEs were pushed — an empty
        // io_uring_enter still costs ~200 ns of mode-switch overhead and
        // showed up as 6 % of total CPU in profiles of the sender loop.
        let pending = ring.submission().len();
        if pending > 0 {
            ring.submit()
                .map_err(|e| io::Error::other(format!("io_uring submit: {e}")))?;
        }

        // --- Collect CQEs (must drain before pushing new SQEs) ---
        // Collecting into a small stack array avoids the CQ borrow
        // conflicting with SQ pushes during processing.
        let mut cqes: [(u64, i32); 4] = [(0, 0); 4];
        let mut cqe_count = 0;
        for cqe in ring.completion() {
            // Every reaped CQE settles one operation the kernel owed —
            // counted even if it overflows the recording array (it
            // cannot: at most a RECV and a SEND are ever in flight).
            // `saturating_sub` keeps a miscount from wrapping into a
            // drain that never terminates; the assert catches the
            // miscount itself rather than letting it hide in release.
            debug_assert!(
                uring.in_flight > 0,
                "CQE reaped with no operation recorded in flight"
            );
            uring.in_flight = uring.in_flight.saturating_sub(1);
            if cqe_count < cqes.len() {
                cqes[cqe_count] = (cqe.user_data(), cqe.result());
                cqe_count += 1;
            }
        }

        let any_cqe = cqe_count > 0;
        for &(token, result) in &cqes[..cqe_count] {
            idle_spins = 0;
            match token {
                TOKEN_RECV => {
                    if result <= 0 {
                        return Err(io::Error::other(format!(
                            "replica disconnected (recv returned {result})"
                        )));
                    }
                    let n = result as usize;
                    parse_buf.extend_from_slice(&recv_buf[..n]);

                    // Extract complete ack frames from parse_buf.
                    let mut cursor = 0;
                    while cursor + 4 <= parse_buf.len() {
                        let frame_len = u32::from_le_bytes(
                            parse_buf[cursor..cursor + 4]
                                .try_into()
                                .expect("bounds checked: 4-byte slice"),
                        ) as usize;
                        if frame_len == 0 || frame_len > MAX_CONTROL_FRAME {
                            return Err(io::Error::other(format!(
                                "invalid ack frame length: {frame_len}"
                            )));
                        }
                        if cursor + 4 + frame_len > parse_buf.len() {
                            break; // Incomplete frame.
                        }
                        let payload = &parse_buf[cursor + 4..cursor + 4 + frame_len];
                        if let Ok(ReplicaMessage::Ack(ack)) = decode_replica_message(payload) {
                            // Eviction on violation: returning Err tears the
                            // connection down; the accept loop's cleanup
                            // disengages the cursors and frees the slot.
                            cursors
                                .record_ack(slot_idx, &ack, sent.get())
                                .map_err(io::Error::other)?;
                            metrics.ack_latency_us[slot_idx]
                                .store(last_send.elapsed().as_micros() as u64, Ordering::Relaxed);
                        }
                        cursor += 4 + frame_len;
                    }
                    // Compact parse_buf.
                    if cursor > 0 {
                        let remaining = parse_buf.len() - cursor;
                        parse_buf.copy_within(cursor.., 0);
                        parse_buf.truncate(remaining);
                    }

                    // Resubmit RECV.
                    let sqe = opcode::Recv::new(
                        types::Fixed(0),
                        recv_buf.as_mut_ptr(),
                        recv_buf.len() as u32,
                    )
                    .build()
                    .user_data(TOKEN_RECV);
                    // SAFETY: `recv_buf` is owned by the guard; only
                    // one RECV is ever in flight (we just consumed the
                    // previous CQE), and the guard drains this one on
                    // every exit path. Ring is single-threaded.
                    unsafe { ring.submission().push(&sqe).expect("SQ full") };
                    uring.in_flight += 1;
                }

                TOKEN_SEND => {
                    if result < 0 {
                        return Err(io::Error::other(format!("send error (returned {result})")));
                    }
                    let sent = result as usize;
                    send_offset += sent;
                    if send_offset >= send_buf.len() {
                        // Fully sent. Measure end-to-end SEND latency: a
                        // healthy io_uring TCP SEND completes in
                        // microseconds; > threshold implies the kernel
                        // waited on cwnd / peer window / retransmit.
                        if let Some(ts) = send_submit_ts.take() {
                            let elapsed = ts.elapsed();
                            if elapsed.as_millis() >= SLOW_SEND_THRESHOLD_MS {
                                tracing::debug!(
                                    slot = slot_idx,
                                    elapsed_us = elapsed.as_micros() as u64,
                                    bytes = send_buf.len(),
                                    "slow SEND completion"
                                );
                                super::log_tcp_info(tcp_fd, "slow_send", slot_idx);
                            }
                        }
                        metrics.bytes_sent[slot_idx]
                            .fetch_add(send_buf.len() as u64, Ordering::Relaxed);
                        send_buf.clear();
                        send_offset = 0;
                        send_in_flight = false;
                    } else {
                        // Partial send — resubmit remainder.
                        let sqe = opcode::Send::new(
                            types::Fixed(0),
                            send_buf[send_offset..].as_ptr(),
                            (send_buf.len() - send_offset) as u32,
                        )
                        .build()
                        .user_data(TOKEN_SEND);
                        // SAFETY: `send_buf` is owned and still pinned
                        // (`send_in_flight` stays true across partial
                        // sends), and the guard drains the CQE on every
                        // exit path; ring is single-threaded.
                        unsafe { ring.submission().push(&sqe).expect("SQ full") };
                        uring.in_flight += 1;
                    }
                }

                _ => {}
            }
        }

        // --- Idle wait ---
        if !any_cqe && send_buf.is_empty() {
            if busy_spin || idle_spins < 1000 {
                idle_spins = idle_spins.wrapping_add(1);
                std::hint::spin_loop();
            } else {
                std::thread::yield_now();
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::uring_teardown::DRAIN_TIMEOUT;
    use io_uring::{IoUring, opcode, types};

    /// Connected AF_UNIX stream pair standing in for the replication
    /// socket. Raw fds rather than `UnixStream` because the guard takes
    /// a `RawFd` and the test needs to outlive it.
    fn socketpair() -> (libc::c_int, libc::c_int) {
        let mut fds = [0 as libc::c_int; 2];
        // SAFETY: `fds` is a valid two-element array for the kernel to
        // fill in; the fds are owned by the caller from here.
        let rc = unsafe { libc::socketpair(libc::AF_UNIX, libc::SOCK_STREAM, 0, fds.as_mut_ptr()) };
        assert_eq!(rc, 0, "socketpair: {}", io::Error::last_os_error());
        (fds[0], fds[1])
    }

    /// SAFETY: both fds are owned by the caller and not used again.
    fn close_pair(a: libc::c_int, b: libc::c_int) {
        // Return value dropped deliberately: nothing can be done about
        // a failing close on a test teardown path.
        unsafe {
            libc::close(a);
            libc::close(b);
        }
    }

    /// Guard over `fd` with a RECV already submitted.
    ///
    /// Panics rather than skipping if the host denies io_uring:
    /// [`live_stream_uring`] is the only live-stream path there is, so
    /// a host that cannot build this ring cannot replicate at all. A
    /// test that quietly passed there would report coverage for a
    /// teardown it never ran.
    fn guard_with_pending_recv(fd: libc::c_int, send_buf: &mut Vec<u8>) -> InflightUring<'_> {
        let ring = IoUring::builder()
            .setup_single_issuer()
            .build(8)
            .expect("io_uring is a hard requirement for replication");
        ring.submitter()
            .register_files(&[fd])
            .expect("register_files");
        let mut uring = InflightUring {
            ring,
            recv_buf: vec![0u8; 4096],
            send_buf,
            in_flight: 0,
            tcp_fd: fd,
        };
        let sqe = opcode::Recv::new(
            types::Fixed(0),
            uring.recv_buf.as_mut_ptr(),
            uring.recv_buf.len() as u32,
        )
        .build()
        .user_data(0);
        // SAFETY: `recv_buf` belongs to the guard being built here, and
        // the guard drains this operation before releasing it. The ring
        // is only ever touched from this thread.
        unsafe { uring.ring.submission().push(&sqe).expect("SQ full") };
        uring.in_flight = 1;
        uring.ring.submit().expect("submit initial RECV");
        uring
    }

    /// The assumption the whole teardown design rests on: a
    /// `shutdown(2)` — not an AsyncCancel, which some hosts filter —
    /// is enough to make the kernel hand back a RECV that nobody is
    /// ever going to write to.
    #[test]
    fn drain_reaps_a_pending_recv_after_shutdown() {
        let (a, b) = socketpair();
        let mut send_buf = Vec::new();
        let mut uring = guard_with_pending_recv(a, &mut send_buf);
        // The peer never writes, so the only thing that can complete
        // this RECV is the shutdown.
        unsafe { libc::shutdown(a, libc::SHUT_RDWR) };
        assert!(
            uring.drain(),
            "drain hit the deadline with a RECV in flight"
        );
        assert_eq!(uring.in_flight, 0, "kernel still owes a CQE");
        drop(uring);
        close_pair(a, b);
    }

    /// `Drop` must quiesce the ring on its own — including issuing the
    /// wake — and must return promptly rather than burning the whole
    /// [`DRAIN_TIMEOUT`], which is what a teardown that silently failed
    /// to wake the operation would look like.
    #[test]
    fn drop_quiesces_a_pending_recv_without_hanging() {
        let (a, b) = socketpair();
        let mut send_buf = vec![0u8; 64];
        let uring = guard_with_pending_recv(a, &mut send_buf);
        let started = std::time::Instant::now();
        drop(uring);
        let elapsed = started.elapsed();
        assert!(
            elapsed < DRAIN_TIMEOUT / 4,
            "teardown took {elapsed:?}, close to the {DRAIN_TIMEOUT:?} deadline"
        );
        // A clean drain releases the buffers; only a drain that cannot
        // prove the kernel is done leaks them.
        assert_eq!(send_buf.len(), 64, "send_buf leaked despite a clean drain");
        close_pair(a, b);
    }
}