armdb 0.7.0

//! Leader-side replication server for FixedStore.
//!
//! - Spawns a TCP acceptor thread.
//! - On the first accepted follower connection, installs SPSC producers
//!   into every shard (lazy install — Mitigation B of spec §5/§7).
//! - Each accepted connection spawns a per-shard serve thread that does
//!   Phase 1 (full scan catch-up, honoring FLAG_EMPTY_STATE to skip DELETED
//!   slots) then Phase 2 (SPSC streaming) if the per-shard consumer is
//!   still available.
//!
//! Single-follower-streaming-per-shard constraint: the SPSC consumer for
//! each shard is held by at most one connection at a time.  If a concurrent
//! follower tries to claim a consumer that is already taken the server sends
//! an Error frame and returns, letting the client reconnect later.  When a
//! connection finishes (cleanly or on error), `ShardConsumerGuard::drop`
//! returns the consumer to the pending slot so the next connection can
//! claim it and enter Phase-2 streaming.

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

use rtrb::{Consumer, Producer, RingBuffer};

use crate::error::{DbError, DbResult};
use crate::shutdown::ShutdownSignal;

use super::engine_access::ArcEngine;
use super::event::FixedReplicationEvent;
use super::protocol::*;

pub const SPSC_CAPACITY: usize = 8192;
const SCAN_CHUNK_BYTES: usize = 64 * 1024;
/// Phase-2 read timeout. The socket stays **blocking** (so writes never see a
/// spurious `WouldBlock` from a full send buffer — see F-05); only reads time
/// out, giving the poll loop a responsive cadence for shutdown checks, the
/// dropped-events probe, and disconnect detection without busy-spinning.
const PHASE2_READ_TIMEOUT_MS: u64 = 20;

/// P2-2: bound on a single blocking `write_frame`. A follower that stops reading
/// fills the socket send buffer; without a write timeout the blocking write hangs
/// forever, stalling shutdown observation and `FixedReplicationServer::drop()`.
/// Set generously (a slow-but-alive follower must not trip it) but bounded so a
/// dead peer cannot pin a handler thread. Symmetric with the 2×-heartbeat read
/// side that detects a silent disconnect.
const WRITE_TIMEOUT_SECS: u64 = HEARTBEAT_INTERVAL_SECS * 2;

// Per-shard slot: holds either (Some producer, Some consumer) before first install,
// (None, Some consumer) between install and first follower accept,
// (None, None) after consumer handed off, or (None, None) forever after.
type PendingSlot = crate::sync::Mutex<(
    Option<Producer<FixedReplicationEvent>>,
    Option<Consumer<FixedReplicationEvent>>,
)>;

pub struct FixedReplicationServer {
    stop: ShutdownSignal,
    acceptor_handle: Option<JoinHandle<()>>,
    handler_handles: Arc<crate::sync::Mutex<Vec<JoinHandle<()>>>>,
    #[allow(dead_code)]
    producers_installed: Arc<AtomicBool>,
}

impl FixedReplicationServer {
    pub fn start(
        bind_addr: SocketAddr,
        engine: ArcEngine,
        signal: ShutdownSignal,
    ) -> DbResult<Self> {
        let shard_count = engine.shard_count();
        let mut pending: Vec<PendingSlot> = Vec::with_capacity(shard_count);
        for _ in 0..shard_count {
            let (p, c) = RingBuffer::new(SPSC_CAPACITY);
            pending.push(crate::sync::Mutex::new((Some(p), Some(c))));
        }
        let pending: Arc<Vec<PendingSlot>> = Arc::new(pending);
        let producers_installed = Arc::new(AtomicBool::new(false));
        let handler_handles = Arc::new(crate::sync::Mutex::new(Vec::new()));

        let listener = TcpListener::bind(bind_addr).map_err(DbError::from)?;
        listener.set_nonblocking(true).ok();

        let acceptor_handle = {
            let engine = engine.clone();
            let pending = pending.clone();
            let producers_installed = producers_installed.clone();
            let stop = signal.clone();
            let hh = handler_handles.clone();
            thread::spawn(move || {
                acceptor_loop(listener, engine, pending, producers_installed, hh, stop);
            })
        };

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

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

impl Drop for FixedReplicationServer {
    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();
        }
    }
}

fn acceptor_loop(
    listener: TcpListener,
    engine: ArcEngine,
    pending: Arc<Vec<PendingSlot>>,
    producers_installed: Arc<AtomicBool>,
    handler_handles: Arc<crate::sync::Mutex<Vec<JoinHandle<()>>>>,
    stop: ShutdownSignal,
) {
    while !stop.is_shutdown() {
        match listener.accept() {
            Ok((stream, addr)) => {
                tracing::info!(%addr, "fixed follower connected");
                stream.set_nodelay(true).ok();
                // The listener is non-blocking for the acceptor loop; on
                // BSD/macOS (and in some Linux configs) accepted streams
                // inherit that flag.  We rely on blocking reads during
                // Phase-1 Ack round-trips, so force it back.
                stream.set_nonblocking(false).ok();
                // Cap Phase-1 blocking reads so a hung follower can't pin
                // a handler thread forever.  A healthy follower sends
                // either an Ack immediately after each batch or a
                // Heartbeat every HEARTBEAT_INTERVAL_SECS; 2× that gives
                // a comfortable ceiling before we conclude the peer is
                // gone.  (Phase 2 switches to non-blocking polling and
                // overrides this setting locally.)
                stream
                    .set_read_timeout(Some(Duration::from_secs(HEARTBEAT_INTERVAL_SECS * 2)))
                    .ok();

                if !producers_installed.swap(true, Ordering::SeqCst) {
                    tracing::info!("first fixed follower — installing SPSC producers");
                    let mut producers = Vec::with_capacity(pending.len());
                    for slot in pending.iter() {
                        let mut guard = crate::sync::lock(slot);
                        producers.push(guard.0.take().expect("producer present on first install"));
                    }
                    engine.install_replication_producers(producers);
                }

                let engine = engine.clone();
                let pending = pending.clone();
                let stop = stop.clone();
                let hh = handler_handles.clone();
                let handle = thread::spawn(move || {
                    if let Err(e) = serve_connection(stream, engine, pending, stop) {
                        tracing::error!(error = %e, "fixed replication connection error");
                    }
                });
                let mut handles = crate::sync::lock(&hh);
                handles.retain(|h| !h.is_finished());
                handles.push(handle);
            }
            Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                stop.wait_timeout(Duration::from_millis(50));
            }
            Err(e) => {
                tracing::error!(error = %e, "fixed accept error");
                stop.wait_timeout(Duration::from_millis(100));
            }
        }
    }
    tracing::info!("fixed replication acceptor stopped");
}

/// RAII guard that holds the SPSC consumer while a connection is active.
///
/// On drop (whether clean or due to an error) the consumer is returned to the
/// `pending` slot so the next connecting follower can claim it and enter
/// Phase-2 streaming.
struct ShardConsumerGuard {
    pending: Arc<Vec<PendingSlot>>,
    shard_id: usize,
    consumer: Option<Consumer<FixedReplicationEvent>>,
}

impl ShardConsumerGuard {
    fn new(
        pending: Arc<Vec<PendingSlot>>,
        shard_id: usize,
        consumer: Consumer<FixedReplicationEvent>,
    ) -> Self {
        Self {
            pending,
            shard_id,
            consumer: Some(consumer),
        }
    }

    fn take(&mut self) -> Consumer<FixedReplicationEvent> {
        self.consumer.take().expect("consumer present")
    }

    fn put_back(&mut self, consumer: Consumer<FixedReplicationEvent>) {
        self.consumer = Some(consumer);
    }
}

impl Drop for ShardConsumerGuard {
    fn drop(&mut self) {
        if let Some(consumer) = self.consumer.take() {
            let mut guard = crate::sync::lock(&self.pending[self.shard_id]);
            if guard.1.is_none() {
                guard.1 = Some(consumer);
            }
        }
    }
}

fn serve_connection(
    stream: TcpStream,
    engine: ArcEngine,
    pending: Arc<Vec<PendingSlot>>,
    stop: ShutdownSignal,
) -> DbResult<()> {
    let mut reader = stream.try_clone().map_err(DbError::from)?;
    // P2-2: bound writes so a follower that stops reading cannot block a
    // `write_frame` on the blocking socket indefinitely. `reader` is a separate
    // dup'd fd (its own read timeout is set per phase), so this affects only the
    // write side.
    let _ = stream.set_write_timeout(Some(Duration::from_secs(WRITE_TIMEOUT_SECS)));
    let mut writer = BufWriter::new(stream);

    // Read initial SyncRequest.
    let frame = read_frame(&mut reader).map_err(DbError::from)?;
    if frame.msg_type != FixedMessageType::SyncRequest {
        return Err(DbError::Replication(format!(
            "expected SyncRequest, got {:?}",
            frame.msg_type
        )));
    }
    let req = SyncRequest::decode(&frame.payload).map_err(DbError::from)?;
    if req.protocol_version != PROTOCOL_VERSION {
        let msg = format!(
            "protocol version mismatch: leader {PROTOCOL_VERSION}, follower {}",
            req.protocol_version
        );
        write_frame(&mut writer, &encode_error(&msg)).map_err(DbError::from)?;
        return Err(DbError::Replication(msg));
    }
    let shard_id = req.shard_id as usize;
    if shard_id >= engine.shard_count() {
        write_frame(&mut writer, &encode_error("invalid shard_id")).map_err(DbError::from)?;
        return Err(DbError::Replication(format!("invalid shard_id {shard_id}")));
    }

    // Send ShardInfo.
    let info = ShardInfo {
        shard_count: engine.shard_count() as u8,
        key_len: engine.key_len() as u16,
        value_len: engine.value_len() as u16,
        slot_size: engine.slot_size(),
        current_slot_count: engine.current_slot_count(shard_id),
        shard_prefix_bits: engine.shard_prefix_bits(),
    };
    write_frame(&mut writer, &info.encode()).map_err(DbError::from)?;

    let skip_deleted = (req.flags & FLAG_EMPTY_STATE) != 0;
    tracing::info!(
        shard_id,
        skip_deleted,
        protocol_version = req.protocol_version,
        "fixed follower handshake complete"
    );

    // Take the SPSC consumer for this shard.
    // If another connection already holds it, reject with an error frame.
    let mut consumer = {
        let mut guard = crate::sync::lock(&pending[shard_id]);
        match guard.1.take() {
            Some(c) => c,
            None => {
                let msg = "shard already streaming";
                write_frame(&mut writer, &encode_error(msg)).map_err(DbError::from)?;
                return Err(DbError::Replication(msg.to_string()));
            }
        }
    };

    // F-02 / P1-6: drain the SPSC ring, then clear any stale "events dropped"
    // flag on claim. The upcoming Phase-1 full scan reconciles all committed
    // state from disk, so buffered ring events are redundant and draining them
    // loses nothing. Draining also frees ring capacity so writes during Phase 1
    // are captured instead of overflowing a still-full ring and immediately
    // re-arming the rescan flag. Without the drain a hot shard livelocks:
    // reconnect -> clear flag -> Phase-1 (never touches the still-full ring) ->
    // the next write overflows and re-sets the flag -> Phase-2 sees it and bails,
    // returning the same full consumer -> reconnect, forever. Draining converts
    // that "re-arm on every single write" pathology into a normal overflow that
    // only re-triggers if Phase-1-window writes actually exceed ring capacity.
    // Only a drop that occurs *after* this point (during Phase 1/2) trips the
    // rescan signal caught by the Phase-2 poll.
    while consumer.pop().is_ok() {}
    let _ = engine.take_dropped(shard_id);

    // Wrap the consumer in a guard so it is returned to the pending slot on
    // any exit path (clean shutdown, network error, or panic).
    let mut consumer_guard = ShardConsumerGuard::new(pending.clone(), shard_id, consumer);

    // Phase 1: full scan — consumer stays inside the guard so that if Phase 1
    // fails the guard's Drop returns the consumer to the pending slot.
    let total = phase1_full_scan(
        &engine,
        shard_id,
        &mut writer,
        &mut reader,
        skip_deleted,
        &stop,
    )?;
    write_frame(
        &mut writer,
        &CaughtUp {
            shard_id: shard_id as u8,
            total_scanned: total,
        }
        .encode(),
    )
    .map_err(DbError::from)?;
    tracing::info!(shard_id, total, "fixed catch-up complete");

    // Phase 2: streaming — take consumer only now (phase2 always returns it).
    let consumer = consumer_guard.take();
    let consumer = phase2_streaming(&engine, shard_id, consumer, &mut writer, &mut reader, &stop)?;
    consumer_guard.put_back(consumer);
    Ok(())
}

fn phase1_full_scan(
    engine: &ArcEngine,
    shard_id: usize,
    writer: &mut BufWriter<TcpStream>,
    reader: &mut TcpStream,
    skip_deleted: bool,
    stop: &ShutdownSignal,
) -> DbResult<u64> {
    use crate::fixed::slot::{
        SLOT_HEADER_SIZE, STATUS_DELETED, STATUS_FREE, STATUS_OCCUPIED, meta_of, pack_meta,
        read_slot, status_of, version_of,
    };

    let slot_size = engine.slot_size() as usize;
    let key_len = engine.key_len();
    let value_len = engine.value_len();
    let slot_count = engine.current_slot_count(shard_id);
    let slots_per_chunk = (SCAN_CHUNK_BYTES / slot_size).max(1);

    let mut total_scanned = 0u64;
    let mut batch = SlotBatchEncoder::new(shard_id as u8, key_len, value_len);

    let mut slot_id = 0u32;
    while slot_id < slot_count {
        if stop.is_shutdown() {
            return Ok(total_scanned);
        }
        let remaining = slot_count - slot_id;
        let this_chunk = remaining.min(slots_per_chunk as u32) as usize;
        let chunk = engine.read_shard_chunk(shard_id, slot_id, this_chunk)?;

        for i in 0..this_chunk {
            let off = i * slot_size;
            let slot_buf = &chunk[off..off + slot_size];
            let meta = meta_of(slot_buf);
            let status = status_of(meta);
            let current_slot = slot_id + i as u32;

            match status {
                STATUS_OCCUPIED => {
                    if let Some((_meta, key, value)) = read_slot(slot_buf, key_len, value_len) {
                        batch.add_occupied(current_slot, meta, key, value);
                        total_scanned += 1;
                    } else {
                        let reset_meta = pack_meta(STATUS_FREE, version_of(meta));
                        batch.add_reset(current_slot, reset_meta);
                        total_scanned += 1;
                    }
                }
                STATUS_DELETED => {
                    if !skip_deleted {
                        let key = &slot_buf[SLOT_HEADER_SIZE..SLOT_HEADER_SIZE + key_len];
                        batch.add_deleted(current_slot, meta, key);
                        total_scanned += 1;
                    }
                }
                STATUS_FREE => {
                    if version_of(meta) != 0 {
                        let reset_meta = pack_meta(STATUS_FREE, version_of(meta));
                        batch.add_reset(current_slot, reset_meta);
                        total_scanned += 1;
                    }
                }
                other => {
                    return Err(DbError::Replication(format!(
                        "fixed Phase-1 scan found unknown slot status {other} at shard {shard_id} slot {current_slot}"
                    )));
                }
            }

            if !batch.is_empty()
                && (batch.len() as usize >= BATCH_MAX_ENTRIES || batch.bytes() >= BATCH_MAX_BYTES)
            {
                flush_and_wait_ack(writer, reader, batch, engine, shard_id)?;
                batch = SlotBatchEncoder::new(shard_id as u8, key_len, value_len);
            }
        }

        slot_id += this_chunk as u32;
    }

    if !batch.is_empty() {
        flush_and_wait_ack(writer, reader, batch, engine, shard_id)?;
    }

    metrics::counter!(
        "armdb.fixed.catchup_slots_scanned",
        "shard" => shard_id.to_string()
    )
    .increment(total_scanned);

    Ok(total_scanned)
}

fn flush_and_wait_ack(
    writer: &mut BufWriter<TcpStream>,
    reader: &mut TcpStream,
    batch: SlotBatchEncoder,
    engine: &ArcEngine,
    shard_id: usize,
) -> DbResult<()> {
    // A well-behaved follower emits at most one Heartbeat per
    // HEARTBEAT_INTERVAL_SECS (=5s) while idle.  The per-stream read
    // timeout (set in `acceptor_loop`) bounds wall-clock time, but a
    // peer that floods heartbeats at sub-timeout intervals could still
    // keep the loop hot forever, so cap the count as defence-in-depth.
    const MAX_HEARTBEATS: u32 = 8;

    let frame = batch.finish();
    write_frame(writer, &frame).map_err(DbError::from)?;
    writer.flush().map_err(DbError::from)?;
    // Loop until we consume an Ack. The follower may have queued a
    // Heartbeat before our batch arrived (sent during its idle
    // read-timeout branch); skip those and keep reading until Ack.
    let mut heartbeats_skipped: u32 = 0;
    loop {
        let ack_frame = read_frame(reader).map_err(DbError::from)?;
        match ack_frame.msg_type {
            FixedMessageType::Ack => {
                let ack = Ack::decode(&ack_frame.payload).map_err(DbError::from)?;
                engine.update_min_replicated_version(shard_id, ack.max_version_seen);
                return Ok(());
            }
            FixedMessageType::Heartbeat => {
                // Follower-side keepalive while waiting for our batch; ignore.
                heartbeats_skipped += 1;
                if heartbeats_skipped > MAX_HEARTBEATS {
                    return Err(DbError::Replication(format!(
                        "too many consecutive heartbeats ({heartbeats_skipped}) \
                         while waiting for Phase-1 Ack"
                    )));
                }
                continue;
            }
            other => {
                return Err(DbError::Replication(format!(
                    "expected Ack during Phase-1 catch-up, got {other:?}"
                )));
            }
        }
    }
}

fn phase2_streaming(
    engine: &ArcEngine,
    shard_id: usize,
    mut consumer: Consumer<FixedReplicationEvent>,
    writer: &mut BufWriter<TcpStream>,
    reader: &mut TcpStream,
    stop: &ShutdownSignal,
) -> DbResult<Consumer<FixedReplicationEvent>> {
    use crate::fixed::slot::{SLOT_HEADER_SIZE, meta_of};

    let key_len = engine.key_len();
    let value_len = engine.value_len();
    let slot_size = engine.slot_size() as usize;
    let mut last_heartbeat = Instant::now();
    let hb_interval = Duration::from_secs(HEARTBEAT_INTERVAL_SECS);

    // F-05: keep the socket blocking; bound only reads with a short timeout so
    // a full send buffer during a write cannot masquerade as a peer disconnect.
    reader
        .set_read_timeout(Some(Duration::from_millis(PHASE2_READ_TIMEOUT_MS)))
        .ok();
    // F-04: resumable reader so a read timeout mid-frame does not lose bytes.
    let mut frame_reader = FrameReader::new();

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

        // F-02: the leader's write path dropped SPSC events for this shard
        // (ring overflow while Phase-1 blocked, or a write burst), so the
        // stream is now missing mutations. A silent gap would let a healthy,
        // heartbeating follower sit on stale data forever, so signal it to
        // reconnect and re-run Phase-1 (full-scan reconcile). Preserve the
        // consumer for the next connection by returning it, not erroring.
        if engine.take_dropped(shard_id) {
            metrics::counter!(
                "armdb.fixed.dropped_rescan",
                "shard" => shard_id.to_string()
            )
            .increment(1);
            let _ = write_frame(writer, &encode_error("events dropped, rescan required"))
                .and_then(|_| writer.flush());
            return Ok(consumer);
        }

        let mut batch = SlotBatchEncoder::new(shard_id as u8, key_len, value_len);
        while (batch.len() as usize) < BATCH_MAX_ENTRIES && batch.bytes() < BATCH_MAX_BYTES {
            match consumer.pop() {
                Ok(FixedReplicationEvent::Write { slot_id, payload }) => {
                    debug_assert_eq!(payload.len(), slot_size);
                    let meta = meta_of(&payload);
                    let key = &payload[SLOT_HEADER_SIZE..SLOT_HEADER_SIZE + key_len];
                    let value = &payload
                        [SLOT_HEADER_SIZE + key_len..SLOT_HEADER_SIZE + key_len + value_len];
                    batch.add_occupied(slot_id, meta, key, value);
                }
                Ok(FixedReplicationEvent::Delete { slot_id, meta, key }) => {
                    batch.add_deleted(slot_id, meta, &key);
                }
                Err(_) => break,
            }
        }

        if !batch.is_empty() {
            let frame_events = batch.len() as u64;
            let frame = batch.finish();
            if write_frame(writer, &frame)
                .and_then(|_| writer.flush())
                .is_err()
            {
                // Peer disconnected during batch write.
                return Ok(consumer);
            }
            metrics::counter!(
                "armdb.fixed.streaming_events_sent",
                "shard" => shard_id.to_string()
            )
            .increment(frame_events);

            // Best-effort Ack check (bounded by the short read timeout).
            match frame_reader.read_frame(reader) {
                Ok(f) if f.msg_type == FixedMessageType::Ack => {
                    if let Ok(ack) = Ack::decode(&f.payload) {
                        engine.update_min_replicated_version(shard_id, ack.max_version_seen);
                    }
                }
                Ok(_) => {}
                Err(ref e)
                    if e.kind() == std::io::ErrorKind::WouldBlock
                        || e.kind() == std::io::ErrorKind::TimedOut => {}
                Err(e) => {
                    // Peer disconnected or connection reset — treat as a clean
                    // exit so the consumer can be returned to the pending slot
                    // for the next connecting follower.
                    tracing::debug!(shard_id, error = %e, "fixed streaming: peer disconnected");
                    return Ok(consumer);
                }
            }
        } else {
            // Idle: no events in consumer ring buffer.
            // Read-timeout probe to detect peer disconnect before the
            // heartbeat timeout fires (so the consumer is returned promptly
            // to the pending slot for the next reconnecting follower).
            match frame_reader.read_frame(reader) {
                Err(ref e)
                    if e.kind() == std::io::ErrorKind::WouldBlock
                        || e.kind() == std::io::ErrorKind::TimedOut => {}
                Ok(f) if f.msg_type == FixedMessageType::Heartbeat => {
                    // Follower keepalive; ignore in Phase 2 idle.
                }
                Ok(_) => {}
                Err(e) => {
                    // Peer disconnected — return consumer for next follower.
                    tracing::debug!(shard_id, error = %e, "fixed streaming idle: peer disconnected");
                    return Ok(consumer);
                }
            }
            if last_heartbeat.elapsed() >= hb_interval {
                if write_frame(writer, &encode_heartbeat())
                    .and_then(|_| writer.flush())
                    .is_err()
                {
                    // Peer disconnected during heartbeat write.
                    return Ok(consumer);
                }
                last_heartbeat = Instant::now();
            }
            thread::sleep(Duration::from_millis(TAIL_POLL_MS));
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::TcpListener;
    use std::sync::atomic::AtomicBool;

    /// Minimal engine that reports one dropped-events flag then clears it.
    struct MockEngine {
        dropped: AtomicBool,
    }

    impl crate::fixed_replication::FixedEngineAccess for MockEngine {
        fn shard_count(&self) -> usize {
            1
        }
        fn key_len(&self) -> usize {
            8
        }
        fn value_len(&self) -> usize {
            16
        }
        fn slot_size(&self) -> u16 {
            32
        }
        fn current_slot_count(&self, _shard_id: usize) -> u32 {
            0
        }
        fn shard_prefix_bits(&self) -> u8 {
            0
        }
        fn read_shard_chunk(
            &self,
            _shard_id: usize,
            _start_slot: u32,
            count: usize,
        ) -> DbResult<Vec<u8>> {
            Ok(vec![0u8; count * 32])
        }
        fn install_replication_producers(&self, _producers: Vec<Producer<FixedReplicationEvent>>) {}
        fn update_min_replicated_version(&self, _shard_id: usize, _version: u32) {}
        fn take_dropped(&self, _shard_id: usize) -> bool {
            self.dropped.swap(false, Ordering::Relaxed)
        }
    }

    /// F-02: when the engine reports dropped events, Phase-2 must send an
    /// `Error` frame telling the follower to rescan, and must return the
    /// consumer (not lose it) so the next connection can stream.
    #[test]
    fn test_f02_phase2_dropped_sends_error_and_returns_consumer() {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let mut client = TcpStream::connect(addr).unwrap();
        let (server_stream, _) = listener.accept().unwrap();

        let engine: ArcEngine = Arc::new(MockEngine {
            dropped: AtomicBool::new(true),
        });
        let (_producer, consumer) = RingBuffer::<FixedReplicationEvent>::new(8);
        let mut writer = BufWriter::new(server_stream.try_clone().unwrap());
        let mut reader = server_stream;
        let stop = ShutdownSignal::new();

        // First loop iteration observes take_dropped == true and returns.
        let returned =
            phase2_streaming(&engine, 0, consumer, &mut writer, &mut reader, &stop).unwrap();
        // Consumer preserved for the next follower.
        drop(returned);

        // The follower peer must receive an Error frame requesting a rescan.
        client
            .set_read_timeout(Some(Duration::from_secs(5)))
            .unwrap();
        let mut fr = FrameReader::new();
        let frame = loop {
            match fr.read_frame(&mut client) {
                Ok(f) => break f,
                Err(ref e)
                    if e.kind() == std::io::ErrorKind::WouldBlock
                        || e.kind() == std::io::ErrorKind::TimedOut =>
                {
                    continue;
                }
                Err(e) => panic!("unexpected error reading Error frame: {e}"),
            }
        };
        assert_eq!(frame.msg_type, FixedMessageType::Error);
        assert!(
            decode_error(&frame.payload).contains("events dropped"),
            "unexpected error payload: {:?}",
            decode_error(&frame.payload)
        );
    }

    /// P1-6: a claim must DRAIN the SPSC ring, not merely clear the dropped flag.
    /// An overflowed (still-full) ring plus a single write arriving during the
    /// Phase-1 scan re-arms the rescan flag, so Phase-2 bails — and because the
    /// ring is never drained, every reconnect repeats it forever (livelock). With
    /// the drain, that Phase-1 write lands in the freed ring and streams normally.
    struct LivelockMock {
        dropped: AtomicBool,
        /// (producer for the claimed ring, already-injected latch)
        injector: crate::sync::Mutex<(Producer<FixedReplicationEvent>, bool)>,
    }

    impl crate::fixed_replication::FixedEngineAccess for LivelockMock {
        fn shard_count(&self) -> usize {
            1
        }
        fn key_len(&self) -> usize {
            8
        }
        fn value_len(&self) -> usize {
            16
        }
        fn slot_size(&self) -> u16 {
            32
        }
        fn current_slot_count(&self, _shard_id: usize) -> u32 {
            4
        }
        fn shard_prefix_bits(&self) -> u8 {
            0
        }
        fn read_shard_chunk(
            &self,
            _shard_id: usize,
            _start_slot: u32,
            count: usize,
        ) -> DbResult<Vec<u8>> {
            // Simulate a leader write that arrives DURING the follower's Phase-1
            // scan: push once into the SPSC ring. If claim did not drain it the
            // ring is still full, the push fails, and the dropped flag is
            // re-armed — the exact condition that makes Phase-2 bail.
            let mut g = crate::sync::lock(&self.injector);
            if !g.1 {
                g.1 = true;
                let ev = FixedReplicationEvent::Delete {
                    slot_id: 0,
                    meta: 1,
                    key: vec![0u8; 8],
                };
                if g.0.push(ev).is_err() {
                    self.dropped.store(true, Ordering::Relaxed);
                }
            }
            Ok(vec![0u8; count * 32])
        }
        fn install_replication_producers(&self, _producers: Vec<Producer<FixedReplicationEvent>>) {}
        fn update_min_replicated_version(&self, _shard_id: usize, _version: u32) {}
        fn take_dropped(&self, _shard_id: usize) -> bool {
            self.dropped.swap(false, Ordering::Relaxed)
        }
    }

    #[test]
    fn test_p1_6_claim_drains_overflowed_ring_no_livelock() {
        // Ring pre-filled to capacity — a prior overflow left it full.
        let (mut prod, cons) = RingBuffer::<FixedReplicationEvent>::new(4);
        while prod
            .push(FixedReplicationEvent::Delete {
                slot_id: 0,
                meta: 1,
                key: vec![0u8; 8],
            })
            .is_ok()
        {}

        let pending: Arc<Vec<PendingSlot>> =
            Arc::new(vec![crate::sync::Mutex::new((None, Some(cons)))]);
        let engine: ArcEngine = Arc::new(LivelockMock {
            dropped: AtomicBool::new(false),
            injector: crate::sync::Mutex::new((prod, false)),
        });

        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let addr = listener.local_addr().unwrap();
        let mut client = TcpStream::connect(addr).unwrap();
        let (server_stream, _) = listener.accept().unwrap();

        let stop = ShutdownSignal::new();
        let stop_srv = stop.clone();
        let handle = thread::spawn(move || {
            let _ = serve_connection(server_stream, engine, pending, stop_srv);
        });

        write_frame(
            &mut client,
            &SyncRequest {
                shard_id: 0,
                protocol_version: PROTOCOL_VERSION,
                flags: 0,
            }
            .encode(),
        )
        .unwrap();

        client
            .set_read_timeout(Some(Duration::from_secs(2)))
            .unwrap();
        let mut fr = FrameReader::new();
        let mut caught_up = false;
        // The frame right after CaughtUp is the Phase-2 verdict: a streaming
        // SlotBatch (drained → fixed) or an Error("events dropped") (livelock).
        let mut verdict: Option<FixedMessageType> = None;
        for _ in 0..16 {
            match fr.read_frame(&mut client) {
                Ok(f) => {
                    if caught_up {
                        verdict = Some(f.msg_type);
                        break;
                    }
                    match f.msg_type {
                        FixedMessageType::CaughtUp => caught_up = true,
                        FixedMessageType::Error => {
                            verdict = Some(FixedMessageType::Error);
                            break;
                        }
                        _ => {}
                    }
                }
                Err(ref e)
                    if e.kind() == std::io::ErrorKind::WouldBlock
                        || e.kind() == std::io::ErrorKind::TimedOut => {}
                Err(e) => panic!("unexpected read error: {e}"),
            }
        }

        stop.shutdown();
        drop(client);
        let _ = handle.join();

        assert!(caught_up, "Phase-1 never completed (no CaughtUp frame)");
        assert_ne!(
            verdict,
            Some(FixedMessageType::Error),
            "claim did not drain the overflowed ring: a Phase-1 write re-armed the \
             rescan flag and Phase-2 bailed — the livelock this fix prevents"
        );
    }
}