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::{SocketAddr, TcpListener, TcpStream};
9use std::sync::Arc;
10use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering};
11
12use tracing::{debug, error, info, warn};
13
14use melin_journal::replication::ReplicationConsumer;
15
16use super::auth::authenticate_replica;
17use super::{ReplicaCursors, ReplicaGate, ReplicationMetrics, SentHighWater};
18use melin_app::Application;
19use melin_transport_core::replication::catchup::{
20    CatchUpResult, bridge_catchup_to_live, can_catch_up_from_journal, catch_up_from_journal,
21    preflight_snapshot_transfer, snapshot_transfer_with,
22};
23use melin_transport_core::replication::protocol::{
24    MAX_CONTROL_FRAME, ReplicaMessage, decode_replica_message, encode_hash_mismatch,
25    encode_heartbeat, encode_need_snapshot, encode_stream_start, read_frame,
26};
27use melin_transport_core::replication::validate::{
28    HandshakeValidation, validate_replica_handshake_settled,
29};
30
31// --- Replication Sender (Primary side) ---
32
33/// Owned state for the replication sender thread.
34pub struct Sender {
35    pub bind_addr: SocketAddr,
36    pub repl_consumer_1: ReplicationConsumer,
37    pub repl_consumer_2: ReplicationConsumer,
38    pub replication_cursor: Arc<AtomicU64>,
39    pub fastest_replica_cursor: Arc<AtomicU64>,
40    pub journal_path: std::path::PathBuf,
41    pub authorized_keys: Arc<melin_app::auth::AuthorizedKeys>,
42    pub evict_flags: [Arc<AtomicBool>; 2],
43    pub active_flags: [Arc<AtomicBool>; 2],
44    pub metrics: Arc<ReplicationMetrics>,
45    pub handler_cores: [usize; 2],
46    pub batch_size: usize,
47    pub heartbeat_secs: u64,
48    pub busy_spin: bool,
49    /// Node fencing state. Read to stamp the primary's epoch onto each
50    /// `StreamStart`, and to self-demote when a replica handshakes with a
51    /// higher epoch (this primary has been superseded). See `crate::fence`.
52    pub fence_state: Arc<melin_transport_core::fence::FenceState>,
53}
54
55/// Run the replication sender. Listens for replica connections,
56/// streams journal data batches, processes acks, and updates the
57/// replication cursor.
58///
59/// Runs on a dedicated thread. Blocks until shutdown.
60pub fn run_sender<A: Application>(
61    config: Sender,
62    shutdown: &AtomicBool,
63    replica_ready: &AtomicBool,
64    replicas_connected: &AtomicU32,
65) {
66    let Sender {
67        bind_addr,
68        repl_consumer_1,
69        repl_consumer_2,
70        replication_cursor,
71        fastest_replica_cursor,
72        journal_path,
73        authorized_keys,
74        evict_flags,
75        active_flags,
76        metrics,
77        handler_cores,
78        batch_size,
79        heartbeat_secs,
80        busy_spin,
81        fence_state,
82    } = config;
83    let listener = match TcpListener::bind(bind_addr) {
84        Ok(l) => l,
85        Err(e) => {
86            error!(addr = %bind_addr, error = %e, "failed to bind replication listener");
87            return;
88        }
89    };
90    // Non-blocking accept so we can check shutdown.
91    if let Err(e) = listener.set_nonblocking(true) {
92        error!(error = %e, "failed to set non-blocking on replication listener");
93        return;
94    }
95
96    info!(addr = %bind_addr, "replication sender listening");
97
98    // Two replica slots, each with its own ring consumer and thread handle.
99    // The accept loop fills empty slots. When a replica disconnects, its
100    // slot becomes available for a new connection. The consumer is `None`
101    // while a handler thread owns it (moved into the thread, returned on exit).
102    struct ReplicaSlot {
103        consumer: Option<ReplicationConsumer>,
104        handle: Option<std::thread::JoinHandle<ReplicationConsumer>>,
105    }
106
107    let mut slots = [
108        ReplicaSlot {
109            consumer: Some(repl_consumer_1),
110            handle: None,
111        },
112        ReplicaSlot {
113            consumer: Some(repl_consumer_2),
114            handle: None,
115        },
116    ];
117
118    // Single owner of the per-replica progress cursors (per-slot acked
119    // positions, shared min/max, and the gate's gauge pair). Shared by
120    // both handler threads — see `ReplicaCursors` for the ordering
121    // contract relative to the active flags.
122    let cursors = Arc::new(ReplicaCursors::new(
123        Arc::clone(&replication_cursor),
124        Arc::clone(&fastest_replica_cursor),
125        Arc::clone(&metrics),
126    ));
127
128    // Per-slot "this connection authenticated" latch. The trading-halt gate
129    // (`replicas_connected`) is lifted by the handler thread only after auth
130    // succeeds, so the main loop must lower it on teardown only when the
131    // handler got that far — a connection that drops mid-auth never counted.
132    // Per-slot `Arc<AtomicBool>` mirrors `active_flags`/`evict_flags`; shared
133    // with the handler thread, which flips it true post-auth.
134    let authenticated_flags: [Arc<AtomicBool>; 2] = [
135        Arc::new(AtomicBool::new(false)),
136        Arc::new(AtomicBool::new(false)),
137    ];
138
139    loop {
140        if shutdown.load(Ordering::Relaxed) {
141            info!("replication sender shutting down");
142            // Wait for active replica threads to finish.
143            for slot in &mut slots {
144                if let Some(handle) = slot.handle.take()
145                    && let Ok(consumer) = handle.join()
146                {
147                    slot.consumer = Some(consumer);
148                }
149            }
150            return;
151        }
152
153        // Check eviction flags from the journal stage. When set, the
154        // journal stage timed out publishing to this slot's ring. We need
155        // to reclaim the consumer so the idle drain loop can clear the ring,
156        // allowing the journal stage to resume publishing.
157        for (i, slot) in slots.iter_mut().enumerate() {
158            if evict_flags[i].load(Ordering::Acquire) && slot.handle.is_some() {
159                metrics.evictions_total.fetch_add(1, Ordering::Relaxed);
160                warn!(
161                    slot = i,
162                    "evicting slow replica (ring backpressure timeout)"
163                );
164                // The handler thread checks shutdown — we can't signal it
165                // individually without adding per-slot flags. Instead, join
166                // the thread with a short timeout by checking is_finished.
167                // The handler's TCP read timeout (5s) will cause it to exit
168                // on the next iteration when it checks shutdown. But we
169                // want faster eviction, so we shutdown the TCP stream to
170                // unblock the read.
171                //
172                // For now, mark the slot and let it be collected below
173                // when the handler finishes naturally (TCP timeout or
174                // next send failure after the ring stops being fed).
175            }
176        }
177
178        // Disengage a finished slot's shared state. Both the clean-exit and
179        // panic arms below call this, so they cannot drift: a panicked handler
180        // must tear down exactly like a clean exit, or it leaks the trading-halt
181        // gate and — far worse — leaves the slot's cursor engaged, freezing the
182        // shared min so the primary stops acking client requests even with a
183        // healthy surviving replica. Borrows only atomics (interior
184        // mutability), so it coexists with the `&mut` iteration over `slots`.
185        let disengage_slot = |slot_idx: usize| {
186            // Lower the gate only if this connection authenticated (and so
187            // lifted it); `swap` reads and resets the latch in one op, and the
188            // `join()` ordered the handler's post-auth writes before here.
189            // `lower` warns "trading halted" if this was the last replica.
190            if authenticated_flags[slot_idx].swap(false, Ordering::AcqRel) {
191                ReplicaGate::new(replicas_connected).lower();
192            }
193            // Disengage cursors BEFORE clearing the active flag — ordering
194            // contract B2 (see `ReplicaCursors`).
195            cursors.clear_on_disconnect(slot_idx);
196            metrics.catching_up[slot_idx].store(false, Ordering::Relaxed);
197            // Journal stage stops publishing to this ring.
198            active_flags[slot_idx].store(false, Ordering::Release);
199        };
200
201        // Collect finished replica threads (disconnected replicas).
202        for (i, slot) in slots.iter_mut().enumerate() {
203            if let Some(ref handle) = slot.handle
204                && handle.is_finished()
205            {
206                let handle = slot.handle.take().expect("just checked is_some");
207                match handle.join() {
208                    Ok(mut consumer) => {
209                        // Drop any unread entries before the consumer is
210                        // stashed back in the slot. The journal stage may
211                        // have published batches into the ring that the
212                        // evicted handler never got to forward — if we
213                        // left them in place, the NEXT handler on this
214                        // slot would drain them to its replica and
215                        // acknowledge with pre-eviction sequences. Those
216                        // acks would stall `replication_cursor` at the
217                        // old position and gate the primary's response
218                        // stage at the slow-replica rate.
219                        //
220                        // Fast-forward to the producer cursor so the
221                        // live-streaming loop starts with a clean ring.
222                        consumer.skip_to_producer();
223                        slot.consumer = Some(consumer);
224                        disengage_slot(i);
225                        // Clear eviction flag after reclaiming the consumer.
226                        if evict_flags[i].load(Ordering::Relaxed) {
227                            evict_flags[i].store(false, Ordering::Release);
228                            warn!(slot = i, "evicted replica — ring ready for reconnection");
229                        } else {
230                            warn!(slot = i, "replica disconnected");
231                        }
232                    }
233                    Err(_) => {
234                        error!(slot = i, "replica handler thread panicked");
235                        // Consumer is lost — the slot stays handle=None /
236                        // consumer=None, so the accept loop skips it. With
237                        // independent rings, only this slot's ring is affected;
238                        // the other replica continues normally. Still disengage
239                        // its shared state, or the gate leaks and the cursor
240                        // freezes (see `disengage_slot`).
241                        disengage_slot(i);
242                        evict_flags[i].store(false, Ordering::Release);
243                    }
244                }
245            }
246        }
247
248        // Find a slot that has a consumer available (not in use by a thread).
249        let empty_slot = slots
250            .iter()
251            .position(|s| s.handle.is_none() && s.consumer.is_some());
252
253        // Accept a connection if there's an empty slot.
254        if let Some(slot_idx) = empty_slot {
255            match listener.accept() {
256                Ok((stream, addr)) => {
257                    info!(addr = %addr, slot = slot_idx, "replica connected");
258                    if let Err(e) = stream.set_nodelay(true) {
259                        debug!(error = %e, "failed to set TCP_NODELAY on replica connection");
260                    }
261                    // SO_BUSY_POLL on the sender side: the per-replica thread
262                    // spins on `recv` for ack frames, so kernel busy-polling
263                    // removes the softirq->wakeup handoff from the ack path.
264                    if let Err(e) =
265                        crate::server::set_busy_poll(&stream, crate::server::BUSY_POLL_US)
266                    {
267                        debug!(error = %e, "failed to set SO_BUSY_POLL on replica connection");
268                    }
269
270                    // The trading-halt gate is NOT lifted here — a bare connect
271                    // hasn't proven a Replication key. The handler lifts it after
272                    // auth succeeds (see `handle_replica_connection`).
273
274                    // Take the consumer out of the slot for the handler thread.
275                    // The slot's consumer becomes None while the thread owns it.
276                    let consumer = slots[slot_idx]
277                        .consumer
278                        .take()
279                        .expect("empty_slot check guarantees consumer is Some");
280
281                    let slot_cursors = Arc::clone(&cursors);
282                    let jpath = journal_path.clone();
283                    let auth_keys = Arc::clone(&authorized_keys);
284                    let slot_metrics = Arc::clone(&metrics);
285                    let slot_active = Arc::clone(&active_flags[slot_idx]);
286                    let slot_evict = Arc::clone(&evict_flags[slot_idx]);
287                    let slot_fence = Arc::clone(&fence_state);
288                    let slot_authenticated = Arc::clone(&authenticated_flags[slot_idx]);
289                    let handler_core = handler_cores[slot_idx];
290                    let shutdown_flag = shutdown as *const AtomicBool as usize;
291                    let ready_flag = replica_ready as *const AtomicBool as usize;
292                    let connected_flag = replicas_connected as *const AtomicU32 as usize;
293                    let handle = std::thread::Builder::new()
294                        .name(format!("repl-{slot_idx}"))
295                        .spawn(move || {
296                            // Pin to a dedicated core if configured (> 0),
297                            // otherwise clear inherited affinity from the
298                            // sender thread so the OS can schedule freely.
299                            if handler_core > 0 {
300                                match melin_app::affinity::pin_to_core(handler_core) {
301                                    Ok(c) => tracing::info!(
302                                        core = c,
303                                        slot = slot_idx,
304                                        "replication handler pinned to core"
305                                    ),
306                                    Err(e) => tracing::warn!(
307                                        error = e,
308                                        slot = slot_idx,
309                                        "failed to pin handler"
310                                    ),
311                                }
312                            } else if let Err(e) = melin_app::affinity::clear_affinity() {
313                                tracing::warn!(error = e, "failed to clear handler affinity");
314                            }
315                            // Safety: shutdown and replica_ready outlive this thread
316                            // (they're on the parent's stack, which blocks on join
317                            // during shutdown).
318                            let shutdown_ref = unsafe { &*(shutdown_flag as *const AtomicBool) };
319                            let ready_ref = unsafe { &*(ready_flag as *const AtomicBool) };
320                            // Safety: replicas_connected outlives this thread —
321                            // same lifetime argument as shutdown/replica_ready
322                            // above (the parent blocks on join during shutdown).
323                            let connected_ref = unsafe { &*(connected_flag as *const AtomicU32) };
324                            let ctx = SlotContext {
325                                cursors: &slot_cursors,
326                                journal_path: &jpath,
327                                authorized_keys: &auth_keys,
328                                shutdown: shutdown_ref,
329                                replica_ready: ready_ref,
330                                active_flag: &slot_active,
331                                evict_flag: &slot_evict,
332                                metrics: &slot_metrics,
333                                fence_state: &slot_fence,
334                                slot_idx,
335                                batch_size,
336                                heartbeat_secs,
337                                busy_spin,
338                                replicas_connected: connected_ref,
339                                authenticated: &slot_authenticated,
340                            };
341                            run_replica_slot::<A>(stream, consumer, &ctx)
342                        })
343                        .expect("spawn replica handler thread");
344                    slots[slot_idx].handle = Some(handle);
345                }
346                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
347                    // No pending connection.
348                }
349                Err(e) => {
350                    error!(error = %e, "replication accept error");
351                }
352            }
353        }
354
355        // No idle drain needed — the journal stage only publishes to
356        // rings where active_flag is true (set by handler on live loop
357        // entry, cleared on disconnect). Idle consumers stay empty.
358
359        std::thread::sleep(std::time::Duration::from_millis(50));
360    }
361}
362
363/// Per-slot state shared across the replica handler call chain
364/// (`run_replica_slot` → `handle_replica_connection` → `live_stream_uring`).
365struct SlotContext<'a> {
366    cursors: &'a Arc<ReplicaCursors>,
367    journal_path: &'a std::path::Path,
368    authorized_keys: &'a melin_app::auth::AuthorizedKeys,
369    shutdown: &'a AtomicBool,
370    replica_ready: &'a AtomicBool,
371    active_flag: &'a AtomicBool,
372    evict_flag: &'a AtomicBool,
373    metrics: &'a ReplicationMetrics,
374    fence_state: &'a melin_transport_core::fence::FenceState,
375    slot_idx: usize,
376    batch_size: usize,
377    heartbeat_secs: u64,
378    busy_spin: bool,
379    /// Trading-halt gate. The handler lifts it (`fetch_add`) only after auth
380    /// succeeds; the main loop lowers it on teardown, gated by `authenticated`.
381    replicas_connected: &'a AtomicU32,
382    /// Per-slot latch the handler sets true once this connection passes auth,
383    /// so the main loop knows whether this slot lifted the gate.
384    authenticated: &'a AtomicBool,
385}
386
387/// Handle a single replica connection on a dedicated thread.
388/// Returns the consumer when the connection ends (for slot reuse).
389fn run_replica_slot<A: Application>(
390    stream: TcpStream,
391    mut consumer: ReplicationConsumer,
392    ctx: &SlotContext<'_>,
393) -> ReplicationConsumer {
394    match handle_replica_connection::<A>(stream, &mut consumer, ctx) {
395        Ok(()) => info!("replica disconnected cleanly"),
396        Err(e) => warn!(error = %e, "replica connection error"),
397    }
398    consumer
399}
400
401fn handle_replica_connection<A: Application>(
402    stream: TcpStream,
403    repl_consumer: &mut ReplicationConsumer,
404    ctx: &SlotContext<'_>,
405) -> io::Result<()> {
406    let SlotContext {
407        cursors,
408        journal_path,
409        authorized_keys,
410        shutdown,
411        replica_ready,
412        active_flag,
413        evict_flag: _,
414        metrics,
415        fence_state,
416        slot_idx,
417        batch_size: _,
418        heartbeat_secs,
419        busy_spin: _,
420        replicas_connected,
421        authenticated,
422    } = ctx;
423    let slot_idx = *slot_idx;
424    let heartbeat_secs = *heartbeat_secs;
425
426    let mut reader = stream.try_clone()?;
427    let mut writer = stream;
428
429    // Set a read timeout for the handshake and auth.
430    reader.set_read_timeout(Some(std::time::Duration::from_secs(10)))?;
431
432    // Authenticate before any data exchange. `reader` and `writer` are clones
433    // of the same socket; auth is sequential (write challenge, read response,
434    // write result) so a single handle suffices. Use `reader` — it carries the
435    // read timeout set above.
436    authenticate_replica(&mut reader, authorized_keys)?;
437    info!("replica authenticated");
438
439    // Auth passed — only now lift the trading-halt gate. The main loop lowers
440    // it on slot teardown, reading `authenticated` to know this connection
441    // lifted it. Set the latch after the increment so a post-join read of the
442    // latch implies the increment is visible too.
443    ReplicaGate::new(replicas_connected).lift();
444    authenticated.store(true, Ordering::Release);
445
446    // Read handshake.
447    let handshake_frame = read_frame(&mut reader, MAX_CONTROL_FRAME)?;
448    let handshake = match decode_replica_message(&handshake_frame)? {
449        ReplicaMessage::Handshake(h) => h,
450        ReplicaMessage::Ack(_) => {
451            return Err(io::Error::other("expected Handshake, got Ack"));
452        }
453    };
454
455    info!(
456        last_sequence = handshake.last_sequence,
457        replica_epoch = handshake.epoch,
458        "replica handshake received"
459    );
460
461    // Fence: a replica advertising an epoch higher than ours means a
462    // promotion happened that this node missed — we are a stale ex-primary.
463    // The policy (latch fenced + trigger shutdown) lives on `FenceState` so
464    // the kernel-TCP and DPDK senders cannot drift; the matching stage
465    // halts new client writes, the response gate stops acking, and the node
466    // winds down for the operator to restart it as a replica. Refuse this
467    // connection so we send no stale stream.
468    let our_epoch = fence_state.epoch();
469    if let Some(first_latch) = fence_state.fence_if_superseded(handshake.epoch, shutdown) {
470        if first_latch {
471            error!(
472                replica_epoch = handshake.epoch,
473                our_epoch,
474                "fenced: a replica advertises a higher epoch — this primary has been \
475                 superseded; self-demoting and shutting down"
476            );
477        }
478        return Err(io::Error::other("fenced by higher-epoch replica"));
479    }
480
481    // Mark this slot as catching up. Cleared when entering the live loop.
482    metrics.catching_up[slot_idx].store(true, Ordering::Relaxed);
483
484    let mut send_buf = Vec::with_capacity(128);
485
486    // Chain validation: recompute our chain at the replica's claimed
487    // position and compare with the hash it reported. A mismatch (or a
488    // position beyond our tip) means its journal holds divergent
489    // history — streaming on top would silently fork the audit trail.
490    // Send HashMismatch (the replica archives its local journal on
491    // receipt) and route it through snapshot resync on this same
492    // connection.
493    let divergent = match validate_replica_handshake_settled(journal_path, &handshake)? {
494        HandshakeValidation::Ok => false,
495        HandshakeValidation::Divergent(kind) => {
496            // Alertable on /metrics: divergence outside an expected
497            // failover rejoin means corruption or a serious bug.
498            metrics.divergence_total.fetch_add(1, Ordering::Relaxed);
499            warn!(
500                last_sequence = handshake.last_sequence,
501                ?kind,
502                "replica journal divergent at handshake — routing through snapshot resync"
503            );
504            true
505        }
506    };
507
508    // Probe whether journal catch-up is possible before committing to
509    // a protocol path. This avoids sending StreamStart only to discover
510    // the journals are too old.
511    let can_catch_up =
512        !divergent && can_catch_up_from_journal(journal_path, handshake.last_sequence)?;
513
514    let mut publish = |buf: &[u8]| -> io::Result<()> {
515        writer.write_all(buf)?;
516        writer.flush()
517    };
518
519    // Cursor/stream floor for everything downstream. A divergent
520    // replica's claimed position is meaningless (its history forked —
521    // possibly ahead of our tip); the resync rebases it, so seed from 0
522    // like a fresh replica: the ack cursors understate until real acks
523    // arrive (safe), and the sent high-water seeds from the catch-up
524    // end instead of a bogus claim that would mark live ring chunks as
525    // already covered.
526    let stream_base = if divergent {
527        0
528    } else {
529        handshake.last_sequence
530    };
531
532    let catchup_end = if can_catch_up {
533        // The lineage origin (oldest segment's header identity) lets a
534        // fresh replica create a byte-identical journal before
535        // consuming the stream. Replicas with local state ignore it.
536        let (lineage_start, lineage_anchor) =
537            melin_transport_core::replication::catchup::lineage_origin(journal_path)?;
538        encode_stream_start(
539            handshake.last_sequence,
540            lineage_start,
541            lineage_anchor,
542            fence_state.epoch(),
543            &mut send_buf,
544        );
545        publish(&send_buf)?;
546        send_buf.clear();
547
548        match catch_up_from_journal::<A::Event>(
549            journal_path,
550            handshake.last_sequence,
551            &mut writer,
552            shutdown,
553        )? {
554            CatchUpResult::Ok(end) => end,
555            CatchUpResult::NeedSnapshot => {
556                return Err(io::Error::other("catch-up failed unexpectedly after probe"));
557            }
558        }
559    } else {
560        // Resync verdict precedes the snapshot data: `HashMismatch`
561        // tells the replica its journal is divergent (it archives its
562        // local lineage on receipt); plain `NeedSnapshot` is the
563        // too-far-behind rebase. The receiver expects `SnapshotBegin`
564        // as the very next frame after the verdict, so
565        // `snapshot_transfer_with` must follow immediately.
566        //
567        // Pre-flight before the verdict goes on the wire: the replica
568        // archives its lineage on receipt, so promising a snapshot we
569        // cannot produce (none configured, boundary segment pruned)
570        // would strand it with its history moved aside. Failing here
571        // drops the connection with the replica's journal intact.
572        preflight_snapshot_transfer(journal_path)?;
573        if divergent {
574            encode_hash_mismatch(&mut send_buf);
575        } else {
576            encode_need_snapshot(&mut send_buf);
577        }
578        publish(&send_buf)?;
579        send_buf.clear();
580        match snapshot_transfer_with::<A::Event>(journal_path, &mut publish, shutdown)? {
581            CatchUpResult::Ok(end) => end,
582            CatchUpResult::NeedSnapshot => {
583                return Err(io::Error::other(
584                    "catch-up failed even after snapshot transfer",
585                ));
586            }
587        }
588    };
589
590    // Engage this slot's cursors and seed the gauge pair the response
591    // gate's `evaluate_durability` reads. Must happen BEFORE the
592    // bridge's active_flag Release — see `ReplicaCursors` for the
593    // ordering contract.
594    cursors.seed_on_handshake(slot_idx, stream_base);
595
596    // Bridge into live streaming: activates the ring, re-reads from the
597    // journal the entries that fell into the activation window, then
598    // drains the ring into sequence-contiguity (back-filling from disk
599    // if a skipped entry hasn't flushed yet) before going live. Returns
600    // the slot's sent high-water — the bound the ack-sanity invariant
601    // checks acks against, advanced by every chunk the live loop
602    // forwards. The bridge closes the catch-up→live gap under load (the
603    // receiver's contiguity gate backstops only the rare quiescent
604    // corner). See `bridge_catchup_to_live`.
605    let mut sent = {
606        let mut publish = |buf: &[u8]| -> io::Result<()> {
607            writer.write_all(buf)?;
608            writer.flush()
609        };
610        bridge_catchup_to_live::<A::Event>(
611            journal_path,
612            stream_base,
613            catchup_end,
614            active_flag,
615            repl_consumer,
616            &mut publish,
617            shutdown,
618        )?
619    };
620
621    // Catch-up complete — replica is entering the live streaming loop.
622    metrics.catching_up[slot_idx].store(false, Ordering::Relaxed);
623
624    // Signal that this replica is ready to consume from the replication
625    // ring. The main thread waits on this before seeding test data.
626    // Must happen AFTER catch-up and overlap drain complete — otherwise
627    // seeding fills the replication ring faster than we can drain it,
628    // deadlocking the journal stage.
629    replica_ready.store(true, Ordering::Release);
630
631    let heartbeat_interval = std::time::Duration::from_secs(heartbeat_secs);
632    let mut last_send = std::time::Instant::now();
633
634    live_stream_uring(
635        writer,
636        repl_consumer,
637        ctx,
638        heartbeat_interval,
639        &mut send_buf,
640        &mut last_send,
641        &mut sent,
642    )
643}
644
645/// io_uring live streaming loop for the primary replication handler.
646///
647/// Live streaming loop using async RECV/SEND via io_uring. A single RECV is always
648/// in-flight for ack frames; SEND is submitted when the replication ring
649/// has data. Both complete via the memory-mapped CQ with zero syscalls
650/// in the hot path.
651fn live_stream_uring(
652    writer: TcpStream,
653    repl_consumer: &mut ReplicationConsumer,
654    ctx: &SlotContext<'_>,
655    heartbeat_interval: std::time::Duration,
656    send_buf: &mut Vec<u8>,
657    last_send: &mut std::time::Instant,
658    sent: &mut SentHighWater,
659) -> io::Result<()> {
660    let SlotContext {
661        cursors,
662        shutdown,
663        evict_flag,
664        metrics,
665        slot_idx,
666        batch_size,
667        busy_spin,
668        // Only used during handshake/catch-up (handle_replica_connection).
669        journal_path: _,
670        authorized_keys: _,
671        replica_ready: _,
672        active_flag: _,
673        heartbeat_secs: _,
674        fence_state: _,
675        replicas_connected: _,
676        authenticated: _,
677    } = ctx;
678    let slot_idx = *slot_idx;
679    let batch_size = *batch_size;
680    let busy_spin = *busy_spin;
681
682    use io_uring::{IoUring, opcode, types};
683    use std::os::unix::io::AsRawFd;
684
685    const TOKEN_RECV: u64 = 0;
686    const TOKEN_SEND: u64 = 1;
687
688    let tcp_fd = writer.as_raw_fd();
689
690    let mut ring: IoUring = IoUring::builder()
691        .setup_single_issuer()
692        .build(8)
693        .map_err(|e| io::Error::other(format!("io_uring init failed: {e}")))?;
694
695    ring.submitter()
696        .register_files(&[tcp_fd])
697        .map_err(|e| io::Error::other(format!("io_uring register_files: {e}")))?;
698
699    // Pin io-wq workers to core 0 (keep them off pipeline cores).
700    {
701        let mut cpuset: libc::cpu_set_t = unsafe { std::mem::zeroed() };
702        unsafe { libc::CPU_SET(0, &mut cpuset) };
703        let _ = ring.submitter().register_iowq_aff(&cpuset);
704    }
705
706    // RECV buffer for ack frames (13 bytes each, but kernel may
707    // coalesce multiple). 4 KiB is plenty.
708    let mut recv_buf = vec![0u8; 4096];
709    // Accumulation buffer for partial ack frame parsing.
710    let mut parse_buf: Vec<u8> = Vec::with_capacity(MAX_CONTROL_FRAME + 4);
711    // RECV is always resubmitted after CQE processing — no explicit
712    // tracking needed. The io_uring kernel guarantees ordering.
713    let mut send_in_flight = false;
714    let mut send_offset: usize = 0;
715    let mut idle_spins: u32 = 0;
716    let mut heartbeat_timer = melin_app::amortized_timer::AmortizedTimer::new();
717
718    // Diagnostic (RUST_LOG=debug): per-slot TCP_INFO snapshot once a
719    // second, slow-SEND detection (CQE elapsed >= threshold), and a
720    // TCP_INFO capture at the evict-exit point. Amortized so the
721    // per-iteration cost is a single `AND` + predictable branch.
722    let mut info_log_timer = melin_app::amortized_timer::AmortizedTimer::new();
723    let mut send_submit_ts: Option<std::time::Instant> = None;
724    const SLOW_SEND_THRESHOLD_MS: u128 = 5;
725
726    // Submit initial RECV.
727    let sqe = opcode::Recv::new(
728        types::Fixed(0),
729        recv_buf.as_mut_ptr(),
730        recv_buf.len() as u32,
731    )
732    .build()
733    .user_data(TOKEN_RECV);
734    // SAFETY: `recv_buf` is owned by this task and lives across the
735    // event loop until the corresponding CQE is drained. The ring is
736    // single-threaded — only this task pushes/reaps on it.
737    unsafe { ring.submission().push(&sqe).expect("SQ full") };
738
739    loop {
740        // --- Check flags ---
741        if shutdown.load(Ordering::Relaxed) {
742            return Ok(());
743        }
744        if evict_flag.load(Ordering::Relaxed) {
745            // Capture the TCP state at the moment of eviction — the
746            // critical frame for comparing an evicted slot's teardown
747            // against the still-live slot's socket state when
748            // diagnosing post-eviction regressions.
749            super::log_tcp_info(tcp_fd, "evict_exit", slot_idx);
750            info!(slot = slot_idx, "handler exiting: evicted by journal stage");
751            return Ok(());
752        }
753
754        // --- Drain replication ring into send_buf (memory, non-blocking) ---
755        //
756        // Ring chunks are wire-ready `InputBatch` frames produced by the
757        // journal stage (phase 3 of the unified-pipeline plan), so the
758        // sender is a passthrough — no decode + re-encode here.
759        if !send_in_flight {
760            let mut coalesced = 0;
761            while coalesced < batch_size {
762                if let Some((meta, data)) = repl_consumer.try_read() {
763                    send_buf.extend_from_slice(data);
764                    repl_consumer.commit();
765                    sent.advance(meta.end_sequence);
766                    coalesced += 1;
767                } else {
768                    break;
769                }
770            }
771
772            if coalesced > 0 {
773                // Submit SEND for the coalesced buffer.
774                let sqe =
775                    opcode::Send::new(types::Fixed(0), send_buf.as_ptr(), send_buf.len() as u32)
776                        .build()
777                        .user_data(TOKEN_SEND);
778                // SAFETY: `send_buf` is owned by this task and pinned
779                // by `send_in_flight = true` below — it isn't mutated
780                // until the matching TOKEN_SEND CQE clears the flag.
781                // The ring is single-threaded.
782                unsafe { ring.submission().push(&sqe).expect("SQ full") };
783                send_in_flight = true;
784                send_offset = 0;
785                *last_send = std::time::Instant::now();
786                send_submit_ts = Some(*last_send);
787                idle_spins = 0;
788                heartbeat_timer = melin_app::amortized_timer::AmortizedTimer::new();
789            } else {
790                // Heartbeat check: amortized when spinning (mask keeps the
791                // clock read at ~10/s at 10M iter/s). In yield mode the loop
792                // already pays a syscall per iteration, so the clock read is
793                // free and must not be skipped — see AmortizedTimer docs.
794                let spinning = busy_spin || idle_spins < 1000;
795                if heartbeat_timer.tick(heartbeat_interval, spinning).is_some() {
796                    encode_heartbeat(sent.get(), send_buf);
797                    let sqe = opcode::Send::new(
798                        types::Fixed(0),
799                        send_buf.as_ptr(),
800                        send_buf.len() as u32,
801                    )
802                    .build()
803                    .user_data(TOKEN_SEND);
804                    // SAFETY: same as the coalesced send above —
805                    // `send_buf` is owned and pinned by `send_in_flight`
806                    // until the TOKEN_SEND CQE clears it.
807                    unsafe { ring.submission().push(&sqe).expect("SQ full") };
808                    send_in_flight = true;
809                    send_offset = 0;
810                    *last_send = std::time::Instant::now();
811                    send_submit_ts = Some(*last_send);
812                }
813            }
814        }
815
816        // Periodic TCP_INFO dump — debug level. Amortized so the
817        // per-iteration cost is a single `AND` + predictable branch.
818        if info_log_timer
819            .tick(
820                std::time::Duration::from_secs(1),
821                busy_spin || idle_spins < 1000,
822            )
823            .is_some()
824        {
825            super::log_tcp_info(tcp_fd, "live_stream", slot_idx);
826        }
827
828        // --- Submit SQEs to kernel (non-blocking) ---
829        // Skip the syscall when no new SQEs were pushed — an empty
830        // io_uring_enter still costs ~200 ns of mode-switch overhead and
831        // showed up as 6 % of total CPU in profiles of the sender loop.
832        let pending = ring.submission().len();
833        if pending > 0 {
834            ring.submit()
835                .map_err(|e| io::Error::other(format!("io_uring submit: {e}")))?;
836        }
837
838        // --- Collect CQEs (must drain before pushing new SQEs) ---
839        // Collecting into a small stack array avoids the CQ borrow
840        // conflicting with SQ pushes during processing.
841        let mut cqes: [(u64, i32); 4] = [(0, 0); 4];
842        let mut cqe_count = 0;
843        for cqe in ring.completion() {
844            if cqe_count < cqes.len() {
845                cqes[cqe_count] = (cqe.user_data(), cqe.result());
846                cqe_count += 1;
847            }
848        }
849
850        let any_cqe = cqe_count > 0;
851        for &(token, result) in &cqes[..cqe_count] {
852            idle_spins = 0;
853            match token {
854                TOKEN_RECV => {
855                    if result <= 0 {
856                        return Err(io::Error::other(format!(
857                            "replica disconnected (recv returned {result})"
858                        )));
859                    }
860                    let n = result as usize;
861                    parse_buf.extend_from_slice(&recv_buf[..n]);
862
863                    // Extract complete ack frames from parse_buf.
864                    let mut cursor = 0;
865                    while cursor + 4 <= parse_buf.len() {
866                        let frame_len = u32::from_le_bytes(
867                            parse_buf[cursor..cursor + 4]
868                                .try_into()
869                                .expect("bounds checked: 4-byte slice"),
870                        ) as usize;
871                        if frame_len == 0 || frame_len > MAX_CONTROL_FRAME {
872                            return Err(io::Error::other(format!(
873                                "invalid ack frame length: {frame_len}"
874                            )));
875                        }
876                        if cursor + 4 + frame_len > parse_buf.len() {
877                            break; // Incomplete frame.
878                        }
879                        let payload = &parse_buf[cursor + 4..cursor + 4 + frame_len];
880                        if let Ok(ReplicaMessage::Ack(ack)) = decode_replica_message(payload) {
881                            // Eviction on violation: returning Err tears the
882                            // connection down; the accept loop's cleanup
883                            // disengages the cursors and frees the slot.
884                            cursors
885                                .record_ack(slot_idx, &ack, sent.get())
886                                .map_err(io::Error::other)?;
887                            metrics.ack_latency_us[slot_idx]
888                                .store(last_send.elapsed().as_micros() as u64, Ordering::Relaxed);
889                        }
890                        cursor += 4 + frame_len;
891                    }
892                    // Compact parse_buf.
893                    if cursor > 0 {
894                        let remaining = parse_buf.len() - cursor;
895                        parse_buf.copy_within(cursor.., 0);
896                        parse_buf.truncate(remaining);
897                    }
898
899                    // Resubmit RECV.
900                    let sqe = opcode::Recv::new(
901                        types::Fixed(0),
902                        recv_buf.as_mut_ptr(),
903                        recv_buf.len() as u32,
904                    )
905                    .build()
906                    .user_data(TOKEN_RECV);
907                    // SAFETY: `recv_buf` is owned by this task; only
908                    // one RECV is ever in flight (we just consumed the
909                    // previous CQE). Ring is single-threaded.
910                    unsafe { ring.submission().push(&sqe).expect("SQ full") };
911                }
912
913                TOKEN_SEND => {
914                    if result < 0 {
915                        return Err(io::Error::other(format!("send error (returned {result})")));
916                    }
917                    let sent = result as usize;
918                    send_offset += sent;
919                    if send_offset >= send_buf.len() {
920                        // Fully sent. Measure end-to-end SEND latency: a
921                        // healthy io_uring TCP SEND completes in
922                        // microseconds; > threshold implies the kernel
923                        // waited on cwnd / peer window / retransmit.
924                        if let Some(ts) = send_submit_ts.take() {
925                            let elapsed = ts.elapsed();
926                            if elapsed.as_millis() >= SLOW_SEND_THRESHOLD_MS {
927                                tracing::debug!(
928                                    slot = slot_idx,
929                                    elapsed_us = elapsed.as_micros() as u64,
930                                    bytes = send_buf.len(),
931                                    "slow SEND completion"
932                                );
933                                super::log_tcp_info(tcp_fd, "slow_send", slot_idx);
934                            }
935                        }
936                        metrics.bytes_sent[slot_idx]
937                            .fetch_add(send_buf.len() as u64, Ordering::Relaxed);
938                        send_buf.clear();
939                        send_offset = 0;
940                        send_in_flight = false;
941                    } else {
942                        // Partial send — resubmit remainder.
943                        let sqe = opcode::Send::new(
944                            types::Fixed(0),
945                            send_buf[send_offset..].as_ptr(),
946                            (send_buf.len() - send_offset) as u32,
947                        )
948                        .build()
949                        .user_data(TOKEN_SEND);
950                        // SAFETY: `send_buf` is owned and still pinned
951                        // (`send_in_flight` stays true across partial
952                        // sends); ring is single-threaded.
953                        unsafe { ring.submission().push(&sqe).expect("SQ full") };
954                    }
955                }
956
957                _ => {}
958            }
959        }
960
961        // --- Idle wait ---
962        if !any_cqe && send_buf.is_empty() {
963            if busy_spin || idle_spins < 1000 {
964                idle_spins = idle_spins.wrapping_add(1);
965                std::hint::spin_loop();
966            } else {
967                std::thread::yield_now();
968            }
969        }
970    }
971}