1use std::io::{self, Write};
10use std::net::{SocketAddr, TcpStream};
11use std::os::unix::io::RawFd;
12use std::sync::atomic::{AtomicBool, Ordering};
13
14use tracing::{info, warn};
15
16use melin_app::Application;
17use melin_journal::JournalWrite;
18use melin_transport_core::pipeline::{InputSlot, JournalStage, JournalStageRun};
19
20use super::auth::authenticate_with_primary;
21use super::receiver_transport::{ReceiverTransport, SessionExit, streaming_loop};
22use super::{
23 ReplicaPipelineHandles, build_replica_pipeline_with_threads, sleep_checking_flags,
24 teardown_replica_pipeline,
25};
26use melin_transport_core::replication::protocol::{
27 Ack, Handshake, MAX_CONTROL_FRAME, MAX_DATA_FRAME, PrimaryMessage, decode_primary_message,
28 encode_ack, encode_handshake, read_frame,
29};
30
31#[inline]
36fn arm_tcp_quickack(fd: RawFd) {
37 let on: libc::c_int = 1;
38 let rc = unsafe {
42 libc::setsockopt(
43 fd,
44 libc::IPPROTO_TCP,
45 libc::TCP_QUICKACK,
46 &on as *const libc::c_int as *const libc::c_void,
47 std::mem::size_of::<libc::c_int>() as libc::socklen_t,
48 )
49 };
50 if rc != 0 {
51 let _ = rc;
52 }
53}
54
55const TOKEN_RECV: u64 = 0;
60const TOKEN_SEND: u64 = 1;
61const TOKEN_PROVIDE: u64 = 2;
62
63const NUM_RECV_BUFFERS: u16 = 16;
65const RECV_BUF_SIZE: usize = MAX_DATA_FRAME;
66const RECV_BUF_GROUP_ID: u16 = 0;
67
68const IORING_CQE_F_BUFFER: u32 = 1 << 0;
70const IORING_CQE_F_MORE: u32 = 1 << 1;
71const IORING_CQE_BUFFER_SHIFT: u32 = 16;
72
73struct UringTransport {
79 ring: io_uring::IoUring,
80 tcp_fd: RawFd,
81 _recv_pool: Vec<u8>,
85 pool_ptr: *mut u8,
86 multishot_active: bool,
87 connected: bool,
88 ack_buf: Vec<u8>,
89 ack_offset: usize,
90 ack_in_flight: bool,
91}
92
93impl UringTransport {
94 fn new(tcp_stream: &TcpStream) -> io::Result<Self> {
95 use io_uring::{IoUring, opcode, types};
96 use std::os::unix::io::AsRawFd;
97
98 let tcp_fd = tcp_stream.as_raw_fd();
99
100 let mut ring: IoUring = IoUring::builder()
101 .setup_single_issuer()
102 .build(64)
103 .map_err(|e| io::Error::other(format!("io_uring init failed: {e}")))?;
104
105 ring.submitter()
106 .register_files(&[tcp_fd])
107 .map_err(|e| io::Error::other(format!("io_uring register_files failed: {e}")))?;
108
109 arm_tcp_quickack(tcp_fd);
110
111 {
113 let mut cpuset: libc::cpu_set_t = unsafe { std::mem::zeroed() };
114 unsafe { libc::CPU_SET(0, &mut cpuset) };
115 let _ = ring.submitter().register_iowq_aff(&cpuset);
116 }
117
118 let mut recv_pool: Vec<u8> = vec![0u8; NUM_RECV_BUFFERS as usize * RECV_BUF_SIZE];
119 let pool_ptr = recv_pool.as_mut_ptr();
120
121 {
123 let sqe = opcode::ProvideBuffers::new(
124 pool_ptr,
125 RECV_BUF_SIZE as i32,
126 NUM_RECV_BUFFERS,
127 RECV_BUF_GROUP_ID,
128 0,
129 )
130 .build()
131 .user_data(TOKEN_PROVIDE);
132 unsafe { ring.submission().push(&sqe).expect("SQ full") };
133 ring.submit_and_wait(1)?;
134 let cqe = ring
135 .completion()
136 .next()
137 .ok_or_else(|| io::Error::other("no cqe after ProvideBuffers"))?;
138 if cqe.result() < 0 {
139 return Err(io::Error::other(format!(
140 "ProvideBuffers failed: {}",
141 cqe.result()
142 )));
143 }
144 }
145
146 let sqe = opcode::RecvMulti::new(types::Fd(tcp_fd), RECV_BUF_GROUP_ID)
148 .build()
149 .user_data(TOKEN_RECV);
150 unsafe { ring.submission().push(&sqe).expect("SQ full") };
151
152 Ok(UringTransport {
153 ring,
154 tcp_fd,
155 _recv_pool: recv_pool,
156 pool_ptr,
157 multishot_active: true,
158 connected: true,
159 ack_buf: Vec::with_capacity(64),
160 ack_offset: 0,
161 ack_in_flight: false,
162 })
163 }
164}
165
166impl ReceiverTransport for UringTransport {
167 fn poll_recv(&mut self, recv_buf: &mut Vec<u8>) -> io::Result<bool> {
168 use io_uring::{opcode, types};
169
170 if !self.connected {
171 return Err(io::Error::other("not connected"));
172 }
173
174 let pending = self.ring.submission().len();
177 if pending > 0 {
178 self.ring.submit()?;
179 }
180
181 let mut cqes: [(u64, i32, u32); 16] = [(0, 0, 0); 16];
183 let mut cqe_count = 0;
184 {
185 let mut cq = self.ring.completion();
186 while cqe_count < cqes.len() {
187 match cq.next() {
188 Some(cqe) => {
189 cqes[cqe_count] = (cqe.user_data(), cqe.result(), cqe.flags());
190 cqe_count += 1;
191 }
192 None => break,
193 }
194 }
195 }
196
197 let mut any_recv = false;
198
199 for &(token, result, flags) in &cqes[..cqe_count] {
200 match token {
201 TOKEN_RECV => {
202 if (flags & IORING_CQE_F_MORE) == 0 {
203 self.multishot_active = false;
204 }
205 if result < 0 {
206 if result == -libc::ENOBUFS {
207 tracing::debug!("recv multishot ENOBUFS — pool exhausted");
208 continue;
209 }
210 self.connected = false;
211 return Err(io::Error::other(format!(
212 "primary disconnected (recv returned {result})"
213 )));
214 }
215 if result == 0 {
216 self.connected = false;
217 return Err(io::Error::other("primary disconnected (recv returned 0)"));
218 }
219 if (flags & IORING_CQE_F_BUFFER) == 0 {
220 self.connected = false;
221 return Err(io::Error::other("recv cqe missing F_BUFFER flag"));
222 }
223
224 arm_tcp_quickack(self.tcp_fd);
225 let n = result as usize;
226 let buf_id = (flags >> IORING_CQE_BUFFER_SHIFT) as usize;
227 let buf_ptr = unsafe { self.pool_ptr.add(buf_id * RECV_BUF_SIZE) };
230 let slice = unsafe { std::slice::from_raw_parts(buf_ptr, n) };
231 recv_buf.extend_from_slice(slice);
232 any_recv = true;
233
234 let provide_sqe = opcode::ProvideBuffers::new(
236 buf_ptr,
237 RECV_BUF_SIZE as i32,
238 1,
239 RECV_BUF_GROUP_ID,
240 buf_id as u16,
241 )
242 .build()
243 .user_data(TOKEN_PROVIDE);
244 unsafe { self.ring.submission().push(&provide_sqe).expect("SQ full") };
245 }
246
247 TOKEN_PROVIDE if result < 0 => {
248 tracing::debug!("ProvideBuffers re-provision failed: {result}");
249 }
250
251 TOKEN_SEND => {
252 if result < 0 {
253 self.connected = false;
254 return Err(io::Error::other(format!(
255 "ack send error (returned {result})"
256 )));
257 }
258 let sent = result as usize;
259 self.ack_offset += sent;
260 if self.ack_offset >= self.ack_buf.len() {
261 self.ack_buf.clear();
262 self.ack_offset = 0;
263 self.ack_in_flight = false;
264 } else {
265 let sqe = opcode::Send::new(
267 types::Fixed(0),
268 self.ack_buf[self.ack_offset..].as_ptr(),
269 (self.ack_buf.len() - self.ack_offset) as u32,
270 )
271 .build()
272 .user_data(TOKEN_SEND);
273 unsafe { self.ring.submission().push(&sqe).expect("SQ full") };
274 }
275 }
276
277 _ => {}
278 }
279 }
280
281 if !self.multishot_active {
283 let sqe = opcode::RecvMulti::new(types::Fd(self.tcp_fd), RECV_BUF_GROUP_ID)
284 .build()
285 .user_data(TOKEN_RECV);
286 unsafe { self.ring.submission().push(&sqe).expect("SQ full") };
287 self.multishot_active = true;
288 }
289
290 Ok(any_recv)
291 }
292
293 fn send_ack(&mut self, ack: &Ack) -> io::Result<bool> {
294 use io_uring::{opcode, types};
295
296 if self.ack_in_flight {
297 return Ok(false);
298 }
299
300 self.ack_buf.clear();
301 encode_ack(ack, &mut self.ack_buf);
302 let sqe = opcode::Send::new(
303 types::Fixed(0),
304 self.ack_buf.as_ptr(),
305 self.ack_buf.len() as u32,
306 )
307 .build()
308 .user_data(TOKEN_SEND);
309 unsafe { self.ring.submission().push(&sqe).expect("SQ full") };
310 self.ack_in_flight = true;
311 self.ack_offset = 0;
312 Ok(true)
313 }
314
315 fn ack_in_flight(&self) -> bool {
316 self.ack_in_flight
317 }
318
319 fn is_connected(&mut self) -> bool {
320 self.connected
321 }
322}
323
324pub type ReceiverResult<A, W> = Result<Option<(A, W)>, Box<dyn std::error::Error>>;
330
331#[allow(clippy::too_many_arguments)]
332pub fn run_receiver<A, W>(
333 primary_addr: SocketAddr,
334 journal_path: &std::path::Path,
335 signing_key: &ed25519_dalek::SigningKey,
336 shutdown: &AtomicBool,
337 promote: &AtomicBool,
338 snapshot_interval_ms: u64,
339 snapshot_path: std::path::PathBuf,
340 cores: crate::server::PipelineCores,
341 receiver_core: usize,
342 group_commit_delay: std::time::Duration,
343 pipeline_depth: usize,
344 busy_spin: bool,
345 rotation: Option<(u64, std::sync::Arc<AtomicBool>)>,
346 factory: std::sync::Arc<dyn melin_app::app_factory::AppFactory<App = A>>,
347) -> ReceiverResult<A, W>
348where
349 A: Application + Send + 'static,
350 A::Event: Send + Sync + 'static,
351 A::Report: Send + 'static,
352 A::QueryResponse: Send + 'static,
353 W: JournalWrite<A::Event> + Send + 'static,
354 JournalStage<A::Event, W>: JournalStageRun<A::Event, Writer = W>,
355{
356 let (mut exchange, mut journal_writer, mut last_sequence, mut chain_hash) =
357 if journal_path.exists() {
358 let engine = if snapshot_path.exists() {
359 info!("recovering replica from snapshot + journal");
360 melin_transport_core::JournaledApp::<A, W>::recover_from_snapshot(
361 &snapshot_path,
362 journal_path,
363 )?
364 } else {
365 melin_transport_core::JournaledApp::<A, W>::recover(factory.empty(), journal_path)?
366 };
367 let next = engine.next_sequence();
368 let last = next.saturating_sub(1);
369 let hash = engine.chain_hash().unwrap_or([0u8; 32]);
370 let (mut exchange, writer) = engine.into_parts();
371 factory.apply_operator_policy(&mut exchange);
372 (Some(exchange), Some(writer), last, hash)
373 } else {
374 (None, None, 0u64, [0u8; 32])
375 };
376
377 let mut backoff = std::time::Duration::from_secs(1);
378 const MAX_BACKOFF: std::time::Duration = std::time::Duration::from_secs(30);
379
380 let mut send_buf = Vec::with_capacity(64);
381 let mut pipeline: Option<ReplicaPipelineHandles<A, W>> = None;
382
383 loop {
385 if let Some(p) = pipeline.as_ref() {
386 last_sequence = p.last_seq.load(Ordering::Acquire);
387 if let Some(ref lock) = p.chain_hash_lock {
388 chain_hash = lock.load().chain_hash;
389 }
390 }
391
392 if shutdown.load(Ordering::Relaxed) {
393 if let Some(mut p) = pipeline.take() {
394 p.input_producer
395 .publish(InputSlot::<A::Event>::shutdown_sentinel());
396 let _ = teardown_replica_pipeline::<A, W>(p);
397 }
398 return Ok(None);
399 }
400 if promote.load(Ordering::Acquire) {
401 info!("promotion triggered while disconnected");
402 if let Some(mut p) = pipeline.take() {
403 p.input_producer
404 .publish(InputSlot::<A::Event>::shutdown_sentinel());
405 if let Some((e, w)) = teardown_replica_pipeline::<A, W>(p) {
406 exchange = Some(e);
407 journal_writer = Some(w);
408 }
409 }
410 return match (exchange, journal_writer) {
411 (Some(e), Some(w)) => Ok(Some((e, w))),
412 _ => Err("promotion requested but no local state available".into()),
413 };
414 }
415
416 info!(primary = %primary_addr, "connecting to primary as replica");
418
419 let stream = match TcpStream::connect(primary_addr) {
420 Ok(s) => s,
421 Err(e) => {
422 warn!(
423 error = %e,
424 backoff_secs = backoff.as_secs(),
425 "failed to connect to primary — retrying"
426 );
427 sleep_checking_flags(backoff, shutdown, promote);
428 if shutdown.load(Ordering::Relaxed) {
429 return Ok(None);
430 }
431 if promote.load(Ordering::Acquire) {
432 info!("promotion triggered during reconnect backoff");
433 if let Some(mut p) = pipeline.take() {
434 p.input_producer
435 .publish(InputSlot::<A::Event>::shutdown_sentinel());
436 if let Some((e, w)) = teardown_replica_pipeline::<A, W>(p) {
437 exchange = Some(e);
438 journal_writer = Some(w);
439 }
440 }
441 return match (exchange, journal_writer) {
442 (Some(e), Some(w)) => Ok(Some((e, w))),
443 _ => Err("promotion requested but no local state available".into()),
444 };
445 }
446 backoff = (backoff * 2).min(MAX_BACKOFF);
447 continue;
448 }
449 };
450 if let Err(e) = stream.set_nodelay(true) {
451 warn!(error = %e, "failed to set TCP_NODELAY on replica receive socket");
452 }
453 {
454 use std::os::unix::io::AsRawFd;
455 let busy_poll: libc::c_int = 50;
456 let rc = unsafe {
457 libc::setsockopt(
458 stream.as_raw_fd(),
459 libc::SOL_SOCKET,
460 libc::SO_BUSY_POLL,
461 &busy_poll as *const libc::c_int as *const libc::c_void,
462 std::mem::size_of::<libc::c_int>() as libc::socklen_t,
463 )
464 };
465 if rc != 0 {
466 let err = std::io::Error::last_os_error();
467 warn!(error = %err, "failed to set SO_BUSY_POLL on replica receive socket");
468 }
469 }
470 stream.set_read_timeout(Some(std::time::Duration::from_secs(5)))?;
471
472 let mut reader = stream.try_clone()?;
473 let mut tcp_writer = stream;
474
475 if let Err(e) = authenticate_with_primary(&mut reader, &mut tcp_writer, signing_key) {
476 warn!(error = %e, "authentication failed — retrying");
477 backoff = (backoff * 2).min(MAX_BACKOFF);
478 continue;
479 }
480 info!("authenticated with primary");
481
482 let handshake = Handshake {
484 last_sequence,
485 chain_hash,
486 };
487 send_buf.clear();
488 encode_handshake(&handshake, &mut send_buf);
489 tcp_writer.write_all(&send_buf)?;
490 tcp_writer.flush()?;
491 send_buf.clear();
492
493 let response_frame = read_frame(&mut reader, MAX_CONTROL_FRAME)?;
495 let response = decode_primary_message(&response_frame)?;
496 let primary_genesis_entry = match response {
497 PrimaryMessage::StreamStart {
498 start_sequence,
499 genesis_entry,
500 } => {
501 info!(start_sequence, "streaming started");
502 genesis_entry
503 }
504 PrimaryMessage::NeedSnapshot => {
505 info!("primary requires snapshot transfer — receiving snapshot");
506
507 if let Some(mut p) = pipeline.take() {
508 p.input_producer
509 .publish(InputSlot::<A::Event>::shutdown_sentinel());
510 let _ = teardown_replica_pipeline::<A, W>(p);
511 }
512
513 let _ = std::fs::remove_file(journal_path);
514 let _ = std::fs::remove_file(&snapshot_path);
515
516 let begin_frame = read_frame(&mut reader, MAX_CONTROL_FRAME)?;
517 let (snap_len, snap_sequence, snap_chain_hash) =
518 match decode_primary_message(&begin_frame)? {
519 PrimaryMessage::SnapshotBegin {
520 snapshot_len,
521 snap_sequence,
522 snap_chain_hash,
523 } => (snapshot_len, snap_sequence, snap_chain_hash),
524 other => {
525 return Err(format!("expected SnapshotBegin, got {other:?}").into());
526 }
527 };
528
529 info!(snap_sequence, snap_len, "receiving snapshot");
530
531 let tmp_path = snapshot_path.with_extension("snapshot.tmp");
532 {
533 let mut tmp_file = std::fs::File::create(&tmp_path)?;
534 let mut received: u64 = 0;
535 let mut running_crc: u32 = 0;
536 loop {
537 let chunk_frame = read_frame(&mut reader, MAX_DATA_FRAME)?;
538 match decode_primary_message(&chunk_frame)? {
539 PrimaryMessage::SnapshotChunk(data) => {
540 std::io::Write::write_all(&mut tmp_file, &data)?;
541 received += data.len() as u64;
542 running_crc = crc32c::crc32c_append(running_crc, &data);
543 }
544 PrimaryMessage::SnapshotEnd {
545 crc32c: expected_crc,
546 } => {
547 tmp_file.sync_all()?;
548 drop(tmp_file);
549
550 if received != snap_len {
551 let _ = std::fs::remove_file(&tmp_path);
552 return Err(format!(
553 "snapshot length mismatch: expected {snap_len} bytes, got {received}"
554 )
555 .into());
556 }
557
558 if running_crc != expected_crc {
559 let _ = std::fs::remove_file(&tmp_path);
560 return Err(format!(
561 "snapshot CRC mismatch: expected {expected_crc:#x}, got {running_crc:#x}"
562 )
563 .into());
564 }
565
566 std::fs::rename(&tmp_path, &snapshot_path)?;
567 info!(snap_sequence, received, "snapshot received and verified");
568 break;
569 }
570 other => {
571 let _ = std::fs::remove_file(&tmp_path);
572 return Err(
573 format!("expected SnapshotChunk/End, got {other:?}").into()
574 );
575 }
576 }
577 }
578 }
579
580 let (snap_exchange, _snap_seq, snap_hash) =
581 melin_transport_core::snapshot::load::<A>(&snapshot_path)?;
582 if snap_hash != snap_chain_hash {
583 return Err(format!(
584 "snapshot chain hash mismatch: primary sent {snap_chain_hash:02x?}, \
585 loaded snapshot has {snap_hash:02x?}"
586 )
587 .into());
588 }
589 exchange = Some(snap_exchange);
590
591 let writer = W::create_continuing(journal_path, snap_sequence + 1, snap_hash)?;
592 journal_writer = Some(writer);
593
594 let ss_frame = read_frame(&mut reader, MAX_CONTROL_FRAME)?;
595 match decode_primary_message(&ss_frame)? {
596 PrimaryMessage::StreamStart {
597 start_sequence,
598 genesis_entry,
599 } => {
600 info!(start_sequence, "streaming resumed after snapshot transfer");
601 genesis_entry
602 }
603 other => {
604 return Err(
605 format!("expected StreamStart after snapshot, got {other:?}").into(),
606 );
607 }
608 }
609 }
610 PrimaryMessage::HashMismatch => {
611 return Err("chain hash mismatch — replica has divergent history".into());
612 }
613 _ => {
614 return Err(format!("unexpected response: {response:?}").into());
615 }
616 };
617
618 if pipeline.is_none() && journal_writer.is_none() {
620 let writer =
621 melin_journal::create_fresh_replica::<_, W>(journal_path, &primary_genesis_entry)?;
622 let mut fresh = factory.empty();
623 factory.apply_operator_policy(&mut fresh);
624 exchange = Some(fresh);
625 journal_writer = Some(writer);
626 }
627
628 if pipeline.is_none() {
630 let cur_exchange = exchange.take().expect("exchange initialized");
631 let cur_writer = journal_writer.take().expect("journal_writer initialized");
632 pipeline = Some(build_replica_pipeline_with_threads::<A, W>(
633 cur_exchange,
634 cur_writer,
635 cores,
636 snapshot_interval_ms,
637 snapshot_path.clone(),
638 group_commit_delay,
639 busy_spin,
640 rotation.clone(),
641 )?);
642 }
643
644 let result = {
646 let p = pipeline.as_mut().expect("pipeline must exist by here");
647 let input_producer = &mut p.input_producer;
648 let journal_cursor = p.journal_cursor.as_ref();
649 std::thread::scope(|s| {
650 let handle = std::thread::Builder::new()
651 .name("replica-receiver".into())
652 .spawn_scoped(s, || {
653 melin_app::affinity::pin_thread("replica-receiver", receiver_core);
654 let mut transport = match UringTransport::new(&tcp_writer) {
655 Ok(t) => t,
656 Err(e) => {
657 tracing::error!(error = %e, "UringTransport init failed");
658 return super::receiver_transport::StreamingResult {
659 exit: SessionExit::Disconnected,
660 received_data: false,
661 };
662 }
663 };
664 streaming_loop::<UringTransport, A::Event>(
665 &mut transport,
666 input_producer,
667 journal_cursor,
668 shutdown,
669 promote,
670 pipeline_depth,
671 busy_spin,
672 Vec::with_capacity(MAX_DATA_FRAME + 4),
673 None,
674 )
675 })
676 .expect("spawn replica-receiver thread");
677 handle.join().expect("replica-receiver thread panicked")
678 })
679 };
680
681 if !matches!(result.exit, SessionExit::Disconnected)
683 && let Some(p) = pipeline.as_mut()
684 {
685 p.input_producer
686 .publish(InputSlot::<A::Event>::shutdown_sentinel());
687 }
688
689 match result.exit {
690 SessionExit::Shutdown => {
691 if let Some(p) = pipeline.take() {
692 let _ = teardown_replica_pipeline::<A, W>(p);
693 }
694 return Ok(None);
695 }
696
697 SessionExit::Promote => {
698 return match pipeline.take() {
699 Some(p) => match teardown_replica_pipeline::<A, W>(p) {
700 Some((e, w)) => Ok(Some((e, w))),
701 None => Err("pipeline failed during promotion".into()),
702 },
703 None => Err("pipeline missing on promote".into()),
704 };
705 }
706
707 SessionExit::Fatal(e) => {
708 if let Some(p) = pipeline.take() {
709 let _ = teardown_replica_pipeline::<A, W>(p);
710 }
711 return Err(e);
712 }
713
714 SessionExit::Disconnected => {
715 if result.received_data {
716 backoff = std::time::Duration::from_secs(1);
717 }
718
719 warn!(
720 last_sequence,
721 backoff_secs = backoff.as_secs(),
722 "reconnecting to primary"
723 );
724 sleep_checking_flags(backoff, shutdown, promote);
725 backoff = (backoff * 2).min(MAX_BACKOFF);
726 }
727 }
728 }
729}