armdb/fixed_replication/

server.rs

//! 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;

// 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)?;
    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 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()));
            }
        }
    };

    // 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);

    reader.set_nonblocking(true).ok();

    loop {
        if stop.is_shutdown() {
            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);

            // Non-blocking Ack check.
            match 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.
            // Non-blocking 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 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));
        }
    }
}