Skip to main content

melin_server_runtime/replication/
tcp_sender.rs

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