armdb 0.7.0

use std::io::{BufWriter, Write};
use std::net::{Shutdown, SocketAddr, TcpListener, TcpStream};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::thread::{self, JoinHandle};
use std::time::Duration;

use crate::error::{DbError, DbResult};
use crate::frame_reader::FrameReader;
use crate::shard::{Shard, Shards};
use crate::shutdown::ShutdownSignal;

use super::ReplicationEntry;
use super::log_reader::ShardLogReader;
use super::protocol::*;
use super::snapshot::ShardCatchupSnapshot;

const BATCH_MAX_ENTRIES: usize = 256;
const BATCH_MAX_BYTES: usize = 64 * 1024;
const TAIL_POLL_MS: u64 = 1;
/// Heartbeat interval used by the independent connection-liveness worker.
pub const HEARTBEAT_INTERVAL_SECS: u64 = 5;

type HandlerHandles = Arc<crate::sync::Mutex<Vec<JoinHandle<()>>>>;
type ConsumerSlots = Arc<Vec<crate::sync::Mutex<Option<rtrb::Consumer<ReplicationEntry>>>>>;

struct ConnectionWriter {
    writer: Arc<crate::sync::Mutex<BufWriter<FrameDeadlineWriter>>>,
    state: crate::sync::Mutex<ConnectionWriterState>,
}

struct FrameDeadlineWriter {
    stream: TcpStream,
    frame_timeout: Duration,
    deadline: Option<std::time::Instant>,
}

impl FrameDeadlineWriter {
    fn new(stream: TcpStream, frame_timeout: Duration) -> Self {
        Self {
            stream,
            frame_timeout,
            deadline: None,
        }
    }

    fn arm(&mut self) -> std::io::Result<()> {
        self.deadline = Some(
            std::time::Instant::now()
                .checked_add(self.frame_timeout)
                .ok_or_else(|| {
                    std::io::Error::new(
                        std::io::ErrorKind::InvalidInput,
                        "replication frame timeout is too large",
                    )
                })?,
        );
        Ok(())
    }

    fn clear(&mut self) {
        self.deadline = None;
    }

    fn remaining(&self) -> std::io::Result<Duration> {
        let deadline = self
            .deadline
            .ok_or_else(|| std::io::Error::other("replication frame deadline is not armed"))?;
        deadline
            .checked_duration_since(std::time::Instant::now())
            .filter(|remaining| !remaining.is_zero())
            .ok_or_else(|| {
                std::io::Error::new(
                    std::io::ErrorKind::TimedOut,
                    "replication frame deadline elapsed",
                )
            })
    }
}

impl Write for FrameDeadlineWriter {
    fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
        configure_write_timeout(&self.stream, self.remaining()?)?;
        self.stream.write(buffer)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        configure_write_timeout(&self.stream, self.remaining()?)?;
        self.stream.flush()
    }
}

struct ConnectionWriterState {
    authorized_gsn: u64,
    failed: bool,
}

type SharedWriter = Arc<ConnectionWriter>;

#[cfg(test)]
fn write_locked_frame<W: std::io::Write>(
    writer: &Arc<crate::sync::Mutex<W>>,
    frame: &Frame,
) -> DbResult<()> {
    let mut writer = crate::sync::lock(writer);
    write_frame(&mut *writer, frame)?;
    Ok(())
}

fn new_shared_writer(
    stream: TcpStream,
    frame_timeout: Duration,
    authorized_gsn: u64,
) -> SharedWriter {
    Arc::new(ConnectionWriter {
        writer: Arc::new(crate::sync::Mutex::new(BufWriter::new(
            FrameDeadlineWriter::new(stream, frame_timeout),
        ))),
        state: crate::sync::Mutex::new(ConnectionWriterState {
            authorized_gsn,
            failed: false,
        }),
    })
}

fn authorize_bootstrap(writer: &SharedWriter, authorized_gsn: u64) {
    crate::sync::lock(&writer.state).authorized_gsn = authorized_gsn;
}

fn bootstrap_authorized_gsn(from_gsn: u64, leader_known_gsn: u64) -> u64 {
    from_gsn.saturating_sub(1).min(leader_known_gsn)
}

fn write_shared_frame(
    writer: &SharedWriter,
    frame: &Frame,
    offered_gsn: Option<u64>,
) -> DbResult<()> {
    let mut state = crate::sync::lock(&writer.state);
    if state.failed {
        return Err(DbError::Io(std::io::Error::new(
            std::io::ErrorKind::BrokenPipe,
            "replication connection writer already failed",
        )));
    }

    let frame_result = {
        let mut frame_writer = crate::sync::lock(&writer.writer);
        match frame_writer.get_mut().arm() {
            Ok(()) => {
                let result = write_frame(&mut *frame_writer, frame);
                frame_writer.get_mut().clear();
                result
            }
            Err(error) => Err(error),
        }
    };
    if let Err(error) = frame_result {
        state.failed = true;
        return Err(DbError::Io(error));
    }
    if let Some(gsn) = offered_gsn {
        state.authorized_gsn = state.authorized_gsn.max(gsn);
    }
    Ok(())
}

fn authorized_gsn(writer: &SharedWriter) -> Option<u64> {
    let state = crate::sync::lock(&writer.state);
    (!state.failed).then_some(state.authorized_gsn)
}

fn writer_failed(writer: &SharedWriter) -> bool {
    crate::sync::lock(&writer.state).failed
}

fn configure_write_timeout(stream: &TcpStream, interval: Duration) -> std::io::Result<()> {
    stream.set_write_timeout(Some(interval))?;
    Ok(())
}

struct ConsumerSlotGuard {
    slots: ConsumerSlots,
    shard_id: usize,
    consumer: Option<rtrb::Consumer<ReplicationEntry>>,
}

impl ConsumerSlotGuard {
    fn take(slots: ConsumerSlots, shard_id: usize) -> Option<Self> {
        let consumer = crate::sync::lock(&slots[shard_id]).take()?;
        Some(Self {
            slots,
            shard_id,
            consumer: Some(consumer),
        })
    }

    fn consumer_mut(&mut self) -> &mut rtrb::Consumer<ReplicationEntry> {
        self.consumer
            .as_mut()
            .expect("consumer slot guard always owns a consumer")
    }
}

impl Drop for ConsumerSlotGuard {
    fn drop(&mut self) {
        if let Some(consumer) = self.consumer.take() {
            *crate::sync::lock(&self.slots[self.shard_id]) = Some(consumer);
        }
    }
}

struct ConnectionLiveness {
    peer_dead: Arc<AtomicBool>,
    reader: TcpStream,
    handle: Option<JoinHandle<()>>,
}

impl ConnectionLiveness {
    fn peer_dead(&self) -> bool {
        self.peer_dead.load(Ordering::Relaxed)
    }
}

impl Drop for ConnectionLiveness {
    fn drop(&mut self) {
        let _ = self.reader.shutdown(Shutdown::Read);
        if let Some(handle) = self.handle.take() {
            let _ = handle.join();
        }
    }
}

fn spawn_liveness(
    reader: TcpStream,
    writer: SharedWriter,
    shard_id: usize,
    min_gsn: Arc<Vec<AtomicU64>>,
    interval: Duration,
    stop: ShutdownSignal,
) -> DbResult<ConnectionLiveness> {
    reader.set_read_timeout(Some(interval))?;
    let shutdown_reader = reader.try_clone()?;
    let peer_dead = Arc::new(AtomicBool::new(false));
    let peer_dead_worker = peer_dead.clone();
    let handle = thread::spawn(move || {
        let mut reader = reader;
        // Resumable reader: an ack torn by an SO_RCVTIMEO timeout mid-frame
        // is reassembled on the next read instead of being destroyed.
        let mut frames = FrameReader::<MessageType>::new();
        let mut awaiting_reply = false;
        // Bytes buffered at the moment the HB was sent / progress was last
        // checked: growth between timeouts proves the peer is alive even if
        // a complete frame has not been assembled yet ("progress = alive").
        let mut buffered_baseline = 0usize;
        while !stop.is_shutdown() {
            match frames.read_frame(&mut reader) {
                Ok(frame) if frame.msg_type == MessageType::Ack => {
                    let Ok(ack) = AckMessage::decode(&frame.payload) else {
                        peer_dead_worker.store(true, Ordering::Relaxed);
                        break;
                    };
                    if ack.shard_id != shard_id as u8 {
                        peer_dead_worker.store(true, Ordering::Relaxed);
                        break;
                    }
                    let Some(authorized_gsn) = authorized_gsn(&writer) else {
                        peer_dead_worker.store(true, Ordering::Relaxed);
                        break;
                    };
                    if ack.last_gsn > authorized_gsn {
                        peer_dead_worker.store(true, Ordering::Relaxed);
                        break;
                    }
                    min_gsn[shard_id].fetch_max(ack.last_gsn, Ordering::Relaxed);
                    awaiting_reply = false;
                }
                Err(error)
                    if matches!(
                        error.kind(),
                        std::io::ErrorKind::TimedOut | std::io::ErrorKind::WouldBlock
                    ) =>
                {
                    if awaiting_reply {
                        if frames.buffered() > buffered_baseline {
                            // Partial ack bytes arrived within the window:
                            // the peer is alive, the frame will complete on
                            // a later read. No HB re-send — we are still
                            // waiting for the reply to the previous one.
                            buffered_baseline = frames.buffered();
                            continue;
                        }
                        peer_dead_worker.store(true, Ordering::Relaxed);
                        break;
                    }
                    if write_shared_frame(&writer, &encode_heartbeat(), None).is_err() {
                        peer_dead_worker.store(true, Ordering::Relaxed);
                        break;
                    }
                    awaiting_reply = true;
                    buffered_baseline = frames.buffered();
                }
                Ok(_) | Err(_) => {
                    peer_dead_worker.store(true, Ordering::Relaxed);
                    break;
                }
            }
        }
    });

    Ok(ConnectionLiveness {
        peer_dead,
        reader: shutdown_reader,
        handle: Some(handle),
    })
}

/// Options for tuning the replication server (e.g. in tests).
pub struct ReplicationServerOptions {
    pub heartbeat_interval_secs: u64,
    pub max_catchup_runs: usize,
}

impl Default for ReplicationServerOptions {
    fn default() -> Self {
        Self {
            heartbeat_interval_secs: HEARTBEAT_INTERVAL_SECS,
            max_catchup_runs: 65_536,
        }
    }
}

/// Replication server running on the leader node.
/// Accepts follower connections and streams entries per-shard.
///
/// # Trust model
///
/// The protocol assumes a **trusted network**: there is no handshake, authentication,
/// or authorization. Any TCP peer that reaches the port receives the shard's full
/// history and takes the single per-shard consumer slot (single-follower-per-shard),
/// so it can also evict a legitimate follower (DoS). Adding auth is wire-breaking, so
/// it is not done by default — bind to a private interface and front untrusted
/// networks with an mTLS tunnel. See the "Network trust model" section in
/// `docs/design.md`.
pub struct ReplicationServer {
    stop: ShutdownSignal,
    acceptor_handle: Option<JoinHandle<()>>,
    handler_handles: HandlerHandles,
    /// Per-shard last GSN ack'd by the at-most-one streaming follower. The
    /// subsystem enforces single-follower-per-shard via an Error frame on a
    /// second concurrent connect; reconnects after disconnect work normally.
    pub min_replicated_gsn: Arc<Vec<AtomicU64>>,
}

impl ReplicationServer {
    /// Start the replication server.
    ///
    /// `consumers`: one SPSC consumer per shard (taken from the ring buffers
    /// installed via `Shard::set_replication_producer`).
    pub fn start(
        bind_addr: SocketAddr,
        shards: Arc<Shards>,
        consumers: Vec<rtrb::Consumer<ReplicationEntry>>,
        max_file_size: u64,
        signal: ShutdownSignal,
    ) -> DbResult<Self> {
        Self::start_with_options(
            bind_addr,
            shards,
            consumers,
            max_file_size,
            signal,
            ReplicationServerOptions::default(),
        )
    }

    pub fn start_with_options(
        bind_addr: SocketAddr,
        shards: Arc<Shards>,
        consumers: Vec<rtrb::Consumer<ReplicationEntry>>,
        max_file_size: u64,
        signal: ShutdownSignal,
        options: ReplicationServerOptions,
    ) -> DbResult<Self> {
        let shard_count = shards.len();
        let heartbeat_secs = options.heartbeat_interval_secs;
        let max_catchup_runs = options.max_catchup_runs;
        if heartbeat_secs == 0 {
            return Err(DbError::Config(
                "heartbeat_interval_secs must be greater than zero",
            ));
        }
        if max_catchup_runs == 0 {
            return Err(DbError::Config(
                "max_catchup_runs must be greater than zero",
            ));
        }

        let min_replicated_gsn: Arc<Vec<AtomicU64>> =
            Arc::new((0..shard_count).map(|_| AtomicU64::new(0)).collect());

        // Wrap consumers in Arc<Mutex> so the acceptor thread can hand them out
        let consumers: ConsumerSlots = Arc::new(
            consumers
                .into_iter()
                .map(|consumer| crate::sync::Mutex::new(Some(consumer)))
                .collect(),
        );

        let listener = TcpListener::bind(bind_addr)?;
        listener.set_nonblocking(true)?;

        let stop2 = signal.clone();
        let shards2 = shards.clone();
        let min_gsn2 = min_replicated_gsn.clone();
        let consumers2 = consumers.clone();
        let handler_handles: HandlerHandles = Arc::new(crate::sync::Mutex::new(Vec::new()));
        let hh2 = handler_handles.clone();

        let acceptor = thread::spawn(move || {
            tracing::info!(%bind_addr, "replication server started");
            while !stop2.is_shutdown() {
                match listener.accept() {
                    Ok((stream, addr)) => {
                        tracing::info!(%addr, "follower connected");
                        // The listener is non-blocking (to poll for connections),
                        // and on some platforms the accepted stream inherits that
                        // flag. Force it back to blocking so the handler's reads —
                        // in particular the liveness reader — actually *block* up to the
                        // read timeout below. Without this the liveness reader's
                        // `read_frame` returns `WouldBlock` immediately on an idle
                        // (but healthy) peer, which the liveness check would
                        // misread as a dead peer.
                        let _ = stream.set_nonblocking(false);
                        let _ = stream.set_nodelay(true);
                        // Bound the SyncRequest read so a peer that connects and
                        // then never sends anything cannot stall the handler
                        // thread indefinitely and block clean shutdown.  The
                        // connection-liveness reader sets its own timeout on
                        // its own cloned half later.
                        let _ =
                            stream.set_read_timeout(Some(Duration::from_secs(2 * heartbeat_secs)));
                        let shards = shards2.clone();
                        let consumers = consumers2.clone();
                        let stop_handler = stop2.clone();
                        let min_gsn = min_gsn2.clone();
                        let hh = hh2.clone();
                        let handle = thread::spawn(move || {
                            if let Err(e) = handle_connection_in_thread(
                                stream,
                                &shards,
                                &consumers,
                                max_file_size,
                                &stop_handler,
                                &min_gsn,
                                heartbeat_secs,
                                max_catchup_runs,
                            ) {
                                tracing::error!(%addr, error = %e, "handler thread error");
                            }
                        });
                        crate::sync::lock(&hh).push(handle);
                    }
                    Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                        stop2.wait_timeout(Duration::from_millis(50));
                    }
                    Err(e) => {
                        tracing::error!(error = %e, "accept error");
                        stop2.wait_timeout(Duration::from_millis(100));
                    }
                }
            }
            tracing::info!("replication server stopped");
        });

        Ok(Self {
            stop: signal,
            acceptor_handle: Some(acceptor),
            handler_handles,
            min_replicated_gsn,
        })
    }

    pub fn stop(&self) {
        self.stop.shutdown();
    }
}

impl crate::compaction::CompactionGuard for ReplicationServer {
    fn min_replicated_gsn(&self, shard_id: u8) -> u64 {
        self.min_replicated_gsn
            .get(shard_id as usize)
            .map(|v| v.load(Ordering::Relaxed))
            .unwrap_or(u64::MAX)
    }
}

impl Drop for ReplicationServer {
    fn drop(&mut self) {
        self.stop.shutdown();
        if let Some(h) = self.acceptor_handle.take() {
            let _ = h.join();
        }
        let mut handles = crate::sync::lock(&self.handler_handles);
        for h in handles.drain(..) {
            let _ = h.join();
        }
    }
}

/// Per-connection handler that runs entirely on the spawned thread.
///
/// Reads the SyncRequest, enforces the single-follower contract (A2, C4, C5),
/// runs catch-up + streaming, and returns the consumer to its slot on every
/// exit path (clean, error, or shutdown).
#[allow(clippy::too_many_arguments)]
fn handle_connection_in_thread(
    stream: TcpStream,
    shards: &Arc<Shards>,
    consumers: &ConsumerSlots,
    max_file_size: u64,
    stop: &ShutdownSignal,
    min_gsn: &Arc<Vec<AtomicU64>>,
    heartbeat_secs: u64,
    max_catchup_runs: usize,
) -> DbResult<()> {
    let mut reader = stream.try_clone().map_err(crate::error::DbError::Io)?;
    let writer = new_shared_writer(stream, Duration::from_secs(heartbeat_secs), 0);

    // Read initial SyncRequest to determine which shard
    let frame = read_frame(&mut reader)?;

    if frame.msg_type != MessageType::SyncRequest {
        write_shared_frame(
            &writer,
            &encode_error(
                ReplicationErrorCode::InvalidRequest,
                &format!("expected SyncRequest, got {:?}", frame.msg_type),
            ),
            None,
        )?;
        return Ok(());
    }

    let req = match SyncRequest::decode(&frame.payload) {
        Ok(req) => req,
        Err(error) => {
            write_shared_frame(
                &writer,
                &encode_error(
                    ReplicationErrorCode::InvalidRequest,
                    &format!("invalid SyncRequest: {error}"),
                ),
                None,
            )?;
            return Ok(());
        }
    };

    if req.protocol_version != VAR_PROTOCOL_VERSION {
        write_shared_frame(
            &writer,
            &encode_error(
                ReplicationErrorCode::ProtocolMismatch,
                &format!(
                    "variable replication protocol mismatch: leader {}, follower {}",
                    VAR_PROTOCOL_VERSION, req.protocol_version
                ),
            ),
            None,
        )?;
        return Ok(());
    }

    let shard_id = req.shard_id as usize;
    if shard_id >= shards.len() {
        write_shared_frame(
            &writer,
            &encode_error(
                ReplicationErrorCode::InvalidRequest,
                &format!("invalid shard_id {shard_id}"),
            ),
            None,
        )?;
        return Ok(());
    }

    let leader_known_gsn = shards[shard_id]
        .gsn()
        .load(Ordering::Relaxed)
        .saturating_sub(1);
    authorize_bootstrap(
        &writer,
        bootstrap_authorized_gsn(req.from_gsn, leader_known_gsn),
    );

    // Send ShardInfo
    let info = ShardInfo {
        protocol_version: VAR_PROTOCOL_VERSION,
        shard_count: shards.len() as u8,
        max_file_size,
    };
    write_shared_frame(&writer, &info.encode(), None)?;

    // Single-follower contract (A2, C4, C5): take the SPSC consumer atomically.
    // If it's absent, a stream is already active for this shard — send Error and exit.
    let Some(mut consumer_guard) = ConsumerSlotGuard::take(consumers.clone(), shard_id) else {
        tracing::warn!(
            shard_id,
            "shard already streaming, rejecting second connection"
        );
        let _ = write_shared_frame(
            &writer,
            &encode_error(
                ReplicationErrorCode::InvalidRequest,
                "shard already streaming",
            ),
            None,
        );
        return Ok(());
    };

    let key_len = req.key_len;
    let from_gsn = req.from_gsn;

    let liveness = spawn_liveness(
        reader,
        writer.clone(),
        shard_id,
        min_gsn.clone(),
        Duration::from_secs(heartbeat_secs),
        stop.clone(),
    )?;

    let outcome = serve_shard_inner(
        &writer,
        shards,
        shard_id,
        from_gsn,
        consumer_guard.consumer_mut(),
        key_len,
        stop,
        &liveness,
        max_catchup_runs,
    );

    if let Err(error) = &outcome
        && !writer_failed(&writer)
    {
        let _ = write_shared_frame(
            &writer,
            &encode_error(replication_error_code(error), &error.to_string()),
            None,
        );
    }
    outcome
}

#[allow(clippy::too_many_arguments)]
fn serve_shard_inner(
    writer: &SharedWriter,
    shards: &[Shard],
    shard_id: usize,
    from_gsn: u64,
    consumer: &mut rtrb::Consumer<ReplicationEntry>,
    key_len: u16,
    stop: &ShutdownSignal,
    liveness: &ConnectionLiveness,
    max_catchup_runs: usize,
) -> DbResult<()> {
    let shard = &shards[shard_id];

    // F1: snapshot the overflow counter BEFORE any streaming, so a ring drop
    // that happens *during* catch-up (when the consumer is idle) is not masked
    // — the first streaming iteration observes the advance and recovers.
    let initial_snapshot = shard.catchup_snapshot()?;
    let snapshot_max_gsn = initial_snapshot.max_gsn;
    let overflow_baseline = initial_snapshot.overflow_baseline;

    // `last_streamed` = the highest GSN the follower is known to hold
    // contiguously. It asked for `from_gsn`, so it already has `from_gsn - 1`;
    // catch-up and every streamed batch advance it.
    let mut last_streamed = from_gsn.saturating_sub(1);

    // Phase 1: Catch-up via ShardLogReader (globally GSN-sorted, F2).
    if from_gsn <= snapshot_max_gsn {
        tracing::info!(shard_id, from_gsn, snapshot_max_gsn, "starting catch-up");

        let last = stream_snapshot(
            writer,
            initial_snapshot,
            shard_id,
            from_gsn,
            key_len,
            max_catchup_runs,
            stop,
            liveness,
        )?;
        last_streamed = last.max(last_streamed);

        if stop.is_shutdown() {
            return Ok(());
        }

        tracing::info!(shard_id, last_streamed, "catch-up complete");
    } else {
        drop(initial_snapshot);
    }

    // Phase 2: Streaming via SPSC.
    {
        tracing::info!(shard_id, "entering streaming mode");

        // F1: the ring drop count we have already reconciled by a catch-up
        // round. When the live counter exceeds it, the ring lost an entry whose
        // GSN may sit *below* a still-queued entry; streaming further would
        // advance the follower's watermark past a permanent hole. Instead we
        // drain the ring and re-stream from disk from `last_streamed + 1`.
        let mut accounted_overflow = overflow_baseline;

        // One overflow-recovery round. `observed` is captured *before* draining
        // so any drop that occurs during the round advances the counter beyond
        // it and forces another round on the next iteration (no drop is masked).
        // Every drained/dropped entry's bytes are on disk, so the disk re-read
        // below reconstructs them; the follower dedups by GSN.
        macro_rules! recover_overflow {
            () => {{
                while consumer.pop().is_ok() {}
                let snapshot = match shard.catchup_snapshot() {
                    Ok(snapshot) => snapshot,
                    Err(e) => return Err(e.into()),
                };
                let observed = snapshot.overflow_baseline;
                let last = match stream_snapshot(
                    writer,
                    snapshot,
                    shard_id,
                    last_streamed.saturating_add(1),
                    key_len,
                    max_catchup_runs,
                    stop,
                    liveness,
                ) {
                    Ok(v) => v,
                    Err(e) => return Err(e.into()),
                };
                last_streamed = last.max(last_streamed);
                accounted_overflow = observed;
                tracing::warn!(
                    shard_id,
                    last_streamed,
                    "recovered replication ring overflow via catch-up round"
                );
            }};
        }

        loop {
            if stop.is_shutdown() {
                break;
            }

            // The liveness worker detected the peer is gone (an unanswered
            // heartbeat, a closed stream, or an invalid response). Stop streaming and
            // cleanup so the SPSC consumer is released — otherwise this loop
            // keeps writing to a half-dead socket until the send buffer fills
            // (minutes of TCP retransmission) while a real reconnect is bounced.
            if liveness.peer_dead() {
                tracing::info!(shard_id, "peer dead — ending streaming, releasing consumer");
                break;
            }

            let mut batch = Vec::new();
            let mut batch_bytes = 0;

            while batch.len() < BATCH_MAX_ENTRIES && batch_bytes < BATCH_MAX_BYTES {
                match consumer.pop() {
                    Ok(entry) => {
                        batch_bytes += entry.data.len();
                        let gsn = extract_gsn(&entry.data);
                        batch.push(WireEntry {
                            entry_len: entry.data.len() as u32,
                            key_len: entry.key_len,
                            gsn,
                            data: entry.data,
                        });
                    }
                    Err(_) => break, // Empty
                }
            }

            if batch.is_empty() {
                // Idle: an overflow while quiescent still needs recovery before
                // we settle into heartbeating.
                if shard.replication_overflow() != accounted_overflow {
                    recover_overflow!();
                    continue;
                }
                thread::sleep(Duration::from_millis(TAIL_POLL_MS));
                continue;
            }

            // F1: before streaming this batch, verify no un-accounted ring drop
            // happened at or before it. If one did, the batch may straddle a
            // hole — discard it (its bytes are on disk) and recover instead.
            if shard.replication_overflow() != accounted_overflow {
                recover_overflow!();
                continue;
            }

            // Highest GSN in this batch — advances `last_streamed` once sent.
            let batch_max_gsn = batch.iter().map(|e| e.gsn).max().unwrap_or(last_streamed);

            let msg = EntryBatch {
                shard_id: shard_id as u8,
                entries: batch,
            };
            // ACKs are handled exclusively by the liveness worker.
            write_shared_frame(writer, &msg.encode(), Some(batch_max_gsn))?;
            last_streamed = batch_max_gsn.max(last_streamed);
        }
    }
    Ok(())
}

/// Stream every on-disk entry with `gsn >= from_gsn` to the follower in global
/// GSN order, in `EntryBatch` frames. Flushes the shard first so the
/// `ShardLogReader` sees all buffered entries. Returns the highest GSN streamed,
/// or `from_gsn - 1` if nothing was found (so the caller's `last_streamed`
/// bookkeeping does not regress).
///
/// Used for both the initial catch-up (Phase 1) and F1 overflow recovery
/// rounds: because it re-reads from disk and the follower dedups by GSN, calling
/// it repeatedly is idempotent.
#[allow(clippy::too_many_arguments)]
fn stream_snapshot(
    writer: &SharedWriter,
    snapshot: ShardCatchupSnapshot,
    shard_id: usize,
    from_gsn: u64,
    key_len: u16,
    max_runs: usize,
    stop: &ShutdownSignal,
    liveness: &ConnectionLiveness,
) -> DbResult<u64> {
    let snapshot_max_gsn = snapshot.max_gsn;
    let mut cancel = || {
        if stop.is_shutdown() {
            Err(DbError::Replication(
                "replication server shutting down".into(),
            ))
        } else if liveness.peer_dead() {
            Err(DbError::Replication("replication peer disconnected".into()))
        } else {
            Ok(())
        }
    };
    let mut log_reader = ShardLogReader::new(snapshot, from_gsn, key_len, max_runs, &mut cancel)?;

    let mut last_gsn = from_gsn.saturating_sub(1);

    loop {
        if stop.is_shutdown() || liveness.peer_dead() {
            return Ok(last_gsn);
        }

        let mut batch = Vec::new();
        let mut batch_bytes = 0;

        loop {
            if batch.len() >= BATCH_MAX_ENTRIES || batch_bytes >= BATCH_MAX_BYTES {
                break;
            }
            match log_reader.next_entry()? {
                Some(entry) => {
                    batch_bytes += entry.data.len();
                    batch.push(WireEntry {
                        entry_len: entry.data.len() as u32,
                        key_len: entry.key_len,
                        gsn: entry.gsn,
                        data: entry.data,
                    });
                }
                None => break,
            }
        }

        if batch.is_empty() {
            break;
        }

        let msg = EntryBatch {
            shard_id: shard_id as u8,
            entries: batch,
        };
        let batch_max_gsn = msg
            .entries
            .iter()
            .map(|entry| entry.gsn)
            .max()
            .unwrap_or(last_gsn);
        write_shared_frame(writer, &msg.encode(), Some(batch_max_gsn))?;
        last_gsn = batch_max_gsn;
    }

    write_shared_frame(
        writer,
        &CaughtUp {
            shard_id: shard_id as u8,
            leader_gsn: snapshot_max_gsn,
        }
        .encode(),
        Some(snapshot_max_gsn),
    )?;
    Ok(last_gsn)
}

fn replication_error_code(error: &DbError) -> ReplicationErrorCode {
    match error {
        DbError::CatchupResourceLimit { .. } => ReplicationErrorCode::ResourceLimit,
        DbError::CorruptedEntry { .. } | DbError::CrcMismatch { .. } => {
            ReplicationErrorCode::CorruptedLog
        }
        #[cfg(feature = "encryption")]
        DbError::EncryptionError(_) => ReplicationErrorCode::CorruptedLog,
        DbError::Io(_) => ReplicationErrorCode::RetryableIo,
        _ => ReplicationErrorCode::RetryableIo,
    }
}

/// Extract GSN (sequence only, no tombstone bit) from raw entry bytes.
fn extract_gsn(data: &[u8]) -> u64 {
    if data.len() < 8 {
        return 0;
    }
    let gsn = u64::from_ne_bytes(data[..8].try_into().expect("impossible"));
    gsn & crate::entry::SEQUENCE_MASK
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn shared_writer_keeps_concurrent_frames_atomic() {
        let bytes = Arc::new(crate::sync::Mutex::new(Vec::<u8>::new()));
        let barrier = Arc::new(std::sync::Barrier::new(3));
        let handles: Vec<_> = [
            Frame {
                msg_type: MessageType::Heartbeat,
                payload: Vec::new(),
            },
            Frame {
                msg_type: MessageType::EntryBatch,
                payload: vec![7; 128 * 1024],
            },
        ]
        .into_iter()
        .map(|frame| {
            let bytes = bytes.clone();
            let barrier = barrier.clone();
            std::thread::spawn(move || {
                barrier.wait();
                write_locked_frame(&bytes, &frame).unwrap();
            })
        })
        .collect();
        barrier.wait();
        for handle in handles {
            handle.join().unwrap();
        }
        let encoded = crate::sync::lock(&bytes).clone();
        let mut cursor = std::io::Cursor::new(encoded);
        let first = read_frame(&mut cursor).unwrap();
        let second = read_frame(&mut cursor).unwrap();
        assert_ne!(first.msg_type, second.msg_type);
        assert_eq!(cursor.position() as usize, cursor.get_ref().len());
    }

    #[test]
    fn consumer_slot_guard_returns_consumer_on_drop_and_unwind() {
        let (_producer, consumer) = rtrb::RingBuffer::<ReplicationEntry>::new(4);
        let slots = Arc::new(vec![crate::sync::Mutex::new(Some(consumer))]);
        {
            let guard = ConsumerSlotGuard::take(slots.clone(), 0).unwrap();
            assert!(crate::sync::lock(&slots[0]).is_none());
            drop(guard);
        }
        assert!(crate::sync::lock(&slots[0]).is_some());
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe({
            let slots = slots.clone();
            move || {
                let _guard = ConsumerSlotGuard::take(slots, 0).unwrap();
                panic!("injected handler panic");
            }
        }));
        assert!(result.is_err());
        assert!(crate::sync::lock(&slots[0]).is_some());
    }

    fn tcp_pair() -> (TcpStream, TcpStream) {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let client = TcpStream::connect(listener.local_addr().unwrap()).unwrap();
        let (server, _) = listener.accept().unwrap();
        // Mirror the production connection setup (see `set_nodelay` on the
        // accept path): without TCP_NODELAY, Nagle delays the second of the
        // three small writes in `write_frame`, tearing the frame across the
        // liveness read-timeout boundary.
        server.set_nodelay(true).unwrap();
        client.set_nodelay(true).unwrap();
        (server, client)
    }

    fn wait_until(timeout: Duration, mut predicate: impl FnMut() -> bool) -> bool {
        let started = std::time::Instant::now();
        while started.elapsed() < timeout {
            if predicate() {
                return true;
            }
            std::thread::sleep(Duration::from_millis(5));
        }
        predicate()
    }

    fn offer_gsn(writer: &SharedWriter, follower: &mut TcpStream, gsn: u64) {
        let frame = EntryBatch {
            shard_id: 0,
            entries: vec![WireEntry {
                entry_len: 8,
                key_len: 8,
                gsn,
                data: vec![0; 8],
            }],
        }
        .encode();
        write_shared_frame(writer, &frame, Some(gsn)).unwrap();
        assert_eq!(
            read_frame(follower).unwrap().msg_type,
            MessageType::EntryBatch
        );
    }

    #[test]
    fn zero_heartbeat_interval_is_rejected() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let address = listener.local_addr().unwrap();
        drop(listener);

        // Acquire a real lease on a tempdir so `Shards::new` gets a valid lock
        // file. The test validates heartbeat rejection, not locking semantics.
        let dir = tempfile::tempdir().unwrap();
        let lock = crate::lock::acquire(dir.path()).unwrap();

        let result = ReplicationServer::start_with_options(
            address,
            Arc::new(Shards::new(Vec::new(), lock)),
            Vec::new(),
            1024,
            ShutdownSignal::new(),
            ReplicationServerOptions {
                heartbeat_interval_secs: 0,
                ..Default::default()
            },
        );

        assert!(matches!(
            result,
            Err(DbError::Config(
                "heartbeat_interval_secs must be greater than zero"
            ))
        ));
    }

    #[test]
    fn ack_bound_clamps_bootstrap_to_leader_known_gsn() {
        assert_eq!(bootstrap_authorized_gsn(u64::MAX, 42), 42);
        assert_eq!(bootstrap_authorized_gsn(8, 42), 7);
        assert_eq!(bootstrap_authorized_gsn(0, 42), 0);
    }

    #[test]
    fn heartbeat_ack_advances_only_the_offered_watermark() {
        let (server, mut follower) = tcp_pair();
        follower
            .set_read_timeout(Some(Duration::from_secs(1)))
            .unwrap();
        let writer = new_shared_writer(server.try_clone().unwrap(), Duration::from_secs(1), 0);
        offer_gsn(&writer, &mut follower, 7);
        let min_gsn = Arc::new(vec![AtomicU64::new(0)]);
        let stop = ShutdownSignal::new();
        let liveness = spawn_liveness(
            server,
            writer,
            0,
            min_gsn.clone(),
            Duration::from_millis(20),
            stop,
        )
        .unwrap();
        let frame = read_frame(&mut follower).unwrap();
        assert_eq!(frame.msg_type, MessageType::Heartbeat);
        write_frame(
            &mut follower,
            &AckMessage {
                shard_id: 0,
                last_gsn: 7,
            }
            .encode(),
        )
        .unwrap();
        assert!(wait_until(Duration::from_secs(1), || {
            min_gsn[0].load(Ordering::Relaxed) == 7
        }));

        assert_eq!(
            read_frame(&mut follower).unwrap().msg_type,
            MessageType::Heartbeat
        );
        write_frame(
            &mut follower,
            &AckMessage {
                shard_id: 0,
                last_gsn: 6,
            }
            .encode(),
        )
        .unwrap();
        assert_eq!(
            read_frame(&mut follower).unwrap().msg_type,
            MessageType::Heartbeat
        );
        assert_eq!(min_gsn[0].load(Ordering::Relaxed), 7);
        assert!(!liveness.peer_dead());
    }

    #[test]
    fn liveness_reassembles_ack_split_across_timeout_windows() {
        let (server, mut follower) = tcp_pair();
        follower
            .set_read_timeout(Some(Duration::from_secs(2)))
            .unwrap();
        let writer = new_shared_writer(server.try_clone().unwrap(), Duration::from_secs(1), 0);
        offer_gsn(&writer, &mut follower, 7);
        let min_gsn = Arc::new(vec![AtomicU64::new(0)]);
        let interval = Duration::from_millis(200);
        let liveness = spawn_liveness(
            server,
            writer,
            0,
            min_gsn.clone(),
            interval,
            ShutdownSignal::new(),
        )
        .unwrap();

        // Wait for the heartbeat, then answer with an ack torn in two: the
        // first fragment right away (inside the current awaiting window),
        // the remainder ~1.5 intervals later — strictly AFTER the point
        // where the legacy reader died (one full window with an incomplete
        // frame), but strictly INSIDE the resumable reader's next window
        // (R1).
        let frame = read_frame(&mut follower).unwrap();
        assert_eq!(frame.msg_type, MessageType::Heartbeat);

        let mut ack_bytes = Vec::new();
        write_frame(
            &mut ack_bytes,
            &AckMessage {
                shard_id: 0,
                last_gsn: 7,
            }
            .encode(),
        )
        .unwrap();
        // 14 bytes: type(1) + len(4) + payload(9). Tear inside the header.
        let (first, rest) = ack_bytes.split_at(3);

        use std::io::Write as _;
        follower.write_all(first).unwrap();
        follower.flush().unwrap();

        std::thread::sleep(interval + interval / 2); // ~1.5 intervals

        // The worker has been through at least one awaiting timeout with
        // partial bytes buffered — by the progress rule it must be alive,
        // and the ack is not assembled yet.
        assert!(!liveness.peer_dead());
        assert_eq!(min_gsn[0].load(Ordering::Relaxed), 0);

        follower.write_all(rest).unwrap();
        follower.flush().unwrap();

        assert!(wait_until(Duration::from_secs(2), || {
            min_gsn[0].load(Ordering::Relaxed) == 7
        }));
        assert!(!liveness.peer_dead());
    }

    #[test]
    fn liveness_declares_peer_dead_after_silent_awaiting_window() {
        // The "empty window = death" contract is preserved: HB sent, not a
        // single byte in reply for a full window -> peer dead.
        let (server, mut follower) = tcp_pair();
        follower
            .set_read_timeout(Some(Duration::from_secs(2)))
            .unwrap();
        let writer = new_shared_writer(server.try_clone().unwrap(), Duration::from_secs(1), 0);
        offer_gsn(&writer, &mut follower, 7);
        let liveness = spawn_liveness(
            server,
            writer,
            0,
            Arc::new(vec![AtomicU64::new(0)]),
            Duration::from_millis(50),
            ShutdownSignal::new(),
        )
        .unwrap();
        assert_eq!(
            read_frame(&mut follower).unwrap().msg_type,
            MessageType::Heartbeat
        );
        assert!(wait_until(Duration::from_secs(2), || liveness.peer_dead()));
    }

    #[test]
    fn future_ack_above_offered_bound_terminates_liveness() {
        let (server, mut follower) = tcp_pair();
        follower
            .set_read_timeout(Some(Duration::from_secs(1)))
            .unwrap();
        let writer = new_shared_writer(server.try_clone().unwrap(), Duration::from_secs(1), 0);
        offer_gsn(&writer, &mut follower, 7);
        let min_gsn = Arc::new(vec![AtomicU64::new(0)]);
        let liveness = spawn_liveness(
            server,
            writer,
            0,
            min_gsn.clone(),
            Duration::from_millis(20),
            ShutdownSignal::new(),
        )
        .unwrap();

        assert_eq!(
            read_frame(&mut follower).unwrap().msg_type,
            MessageType::Heartbeat
        );
        write_frame(
            &mut follower,
            &AckMessage {
                shard_id: 0,
                last_gsn: u64::MAX,
            }
            .encode(),
        )
        .unwrap();

        assert!(wait_until(Duration::from_secs(1), || liveness.peer_dead()));
        assert_eq!(min_gsn[0].load(Ordering::Relaxed), 0);
    }

    #[cfg(unix)]
    #[test]
    fn backpressure_write_timeout_releases_liveness_and_consumer() {
        use std::os::fd::AsRawFd;
        use std::sync::mpsc;

        let (server, _follower) = tcp_pair();
        let send_buffer_bytes: libc::c_int = 4096;
        let set_buffer_result = unsafe {
            libc::setsockopt(
                server.as_raw_fd(),
                libc::SOL_SOCKET,
                libc::SO_SNDBUF,
                (&send_buffer_bytes as *const libc::c_int).cast(),
                std::mem::size_of_val(&send_buffer_bytes) as libc::socklen_t,
            )
        };
        assert_eq!(set_buffer_result, 0);
        let writer = new_shared_writer(server.try_clone().unwrap(), Duration::from_millis(50), 0);
        let liveness = spawn_liveness(
            server,
            writer.clone(),
            0,
            Arc::new(vec![AtomicU64::new(0)]),
            Duration::from_millis(10),
            ShutdownSignal::new(),
        )
        .unwrap();
        let (_producer, consumer) = rtrb::RingBuffer::<ReplicationEntry>::new(4);
        let slots = Arc::new(vec![crate::sync::Mutex::new(Some(consumer))]);
        let guard = ConsumerSlotGuard::take(slots.clone(), 0).unwrap();

        let (result_tx, result_rx) = mpsc::channel();
        let writer_for_data = writer.clone();
        let data_sender = std::thread::spawn(move || {
            let started = std::time::Instant::now();
            let result = write_shared_frame(
                &writer_for_data,
                &Frame {
                    msg_type: MessageType::EntryBatch,
                    payload: vec![7; MAX_FRAME_SIZE],
                },
                Some(7),
            );
            result_tx.send((result, started.elapsed())).unwrap();
        });

        let (result, elapsed) = result_rx
            .recv_timeout(Duration::from_secs(2))
            .expect("backpressured write must respect its configured timeout");
        assert!(matches!(result, Err(DbError::Io(_))));
        assert!(elapsed < Duration::from_secs(1));
        assert!(writer_failed(&writer));
        assert_eq!(crate::sync::lock(&writer.state).authorized_gsn, 0);
        let retry = write_shared_frame(&writer, &encode_heartbeat(), None);
        assert!(matches!(
            retry,
            Err(DbError::Io(error)) if error.kind() == std::io::ErrorKind::BrokenPipe
        ));
        assert!(wait_until(Duration::from_secs(1), || liveness.peer_dead()));
        data_sender.join().unwrap();

        let cleanup_started = std::time::Instant::now();
        drop(liveness);
        drop(guard);
        assert!(cleanup_started.elapsed() < Duration::from_secs(1));
        assert!(crate::sync::lock(&slots[0]).is_some());
    }

    #[cfg(unix)]
    #[test]
    fn trickle_reader_cannot_extend_absolute_frame_deadline() {
        use std::io::Read;
        use std::os::fd::AsRawFd;
        use std::sync::mpsc;

        let (server, mut follower) = tcp_pair();
        let follower_shutdown = follower.try_clone().unwrap();
        let send_buffer_bytes: libc::c_int = 4096;
        let set_buffer_result = unsafe {
            libc::setsockopt(
                server.as_raw_fd(),
                libc::SOL_SOCKET,
                libc::SO_SNDBUF,
                (&send_buffer_bytes as *const libc::c_int).cast(),
                std::mem::size_of_val(&send_buffer_bytes) as libc::socklen_t,
            )
        };
        assert_eq!(set_buffer_result, 0);
        let frame_deadline = Duration::from_millis(100);
        let writer = new_shared_writer(server.try_clone().unwrap(), frame_deadline, 0);
        let liveness = spawn_liveness(
            server,
            writer.clone(),
            0,
            Arc::new(vec![AtomicU64::new(0)]),
            Duration::from_millis(10),
            ShutdownSignal::new(),
        )
        .unwrap();
        let (_producer, consumer) = rtrb::RingBuffer::<ReplicationEntry>::new(4);
        let slots = Arc::new(vec![crate::sync::Mutex::new(Some(consumer))]);
        let guard = ConsumerSlotGuard::take(slots.clone(), 0).unwrap();

        let stop_trickle = Arc::new(AtomicBool::new(false));
        let stop_trickle_worker = stop_trickle.clone();
        let trickle_reader = std::thread::spawn(move || {
            let mut buffer = [0u8; 4096];
            while !stop_trickle_worker.load(Ordering::Relaxed) {
                std::thread::sleep(Duration::from_millis(30));
                if follower.read(&mut buffer).is_err() {
                    break;
                }
            }
        });

        let (result_tx, result_rx) = mpsc::channel();
        let writer_for_data = writer.clone();
        let data_sender = std::thread::spawn(move || {
            let started = std::time::Instant::now();
            let result = write_shared_frame(
                &writer_for_data,
                &Frame {
                    msg_type: MessageType::EntryBatch,
                    payload: vec![7; MAX_FRAME_SIZE],
                },
                Some(7),
            );
            result_tx.send((result, started.elapsed())).unwrap();
        });

        let bounded_result = result_rx.recv_timeout(Duration::from_millis(500));
        stop_trickle.store(true, Ordering::Relaxed);
        let _ = follower_shutdown.shutdown(Shutdown::Both);
        trickle_reader.join().unwrap();
        data_sender.join().unwrap();

        let (result, elapsed) =
            bounded_result.expect("trickle reads must not extend one absolute frame deadline");
        assert!(matches!(result, Err(DbError::Io(_))));
        assert!(elapsed < Duration::from_millis(500));
        assert!(writer_failed(&writer));
        assert_eq!(crate::sync::lock(&writer.state).authorized_gsn, 0);
        let retry = write_shared_frame(&writer, &encode_heartbeat(), None);
        assert!(matches!(
            retry,
            Err(DbError::Io(error)) if error.kind() == std::io::ErrorKind::BrokenPipe
        ));
        assert!(wait_until(Duration::from_secs(1), || liveness.peer_dead()));

        let cleanup_started = std::time::Instant::now();
        drop(liveness);
        drop(guard);
        assert!(cleanup_started.elapsed() < Duration::from_secs(1));
        assert!(crate::sync::lock(&slots[0]).is_some());
    }

    #[cfg(feature = "encryption")]
    #[test]
    fn replication_error_code_maps_encryption_error_to_corrupted_log() {
        let error = DbError::EncryptionError("decryption/authentication failed".into());
        assert_eq!(
            replication_error_code(&error),
            ReplicationErrorCode::CorruptedLog
        );
    }

    #[cfg(feature = "encryption")]
    #[test]
    fn replication_error_code_maps_encrypted_corrupt_tag_snapshot_to_corrupted_log() {
        use crate::crypto::PageCipher;
        use crate::entry::serialize_entry;

        let dir = tempfile::tempdir().unwrap();
        let shard = Shard::open_encrypted(
            0,
            dir.path(),
            1 << 20,
            64 * 1024,
            false,
            false,
            false,
            crate::config::IoBackend::default(),
            Some(Arc::new(PageCipher::new(&[0x42; 32]).unwrap())),
            Arc::new(AtomicU64::new(2)),
        )
        .unwrap();
        let entry = serialize_entry(1, &1u64.to_be_bytes(), b"encrypted", false);
        shard.lock().append_raw_entry(0, 8, &entry).unwrap();
        shard.rotate_active_for_test(8).unwrap();
        let snapshot = shard.catchup_snapshot().unwrap();

        crate::test_faults::corrupt_tags(dir.path());
        let error = ShardLogReader::new(snapshot, 0, 8, 8, &mut || Ok(())).unwrap_err();
        assert!(matches!(error, DbError::EncryptionError(_)));
        assert_eq!(
            replication_error_code(&error),
            ReplicationErrorCode::CorruptedLog
        );
    }
}