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}