armdb 0.14.2

mod apply;
pub mod client;
mod log_reader;
pub mod protocol;
pub mod server;
pub(crate) mod snapshot;

pub use apply::{ApplyOutcome, ReplicationTarget};
pub use client::{ReplicationClient, ReplicationClientOptions};
pub use server::{ReplicationServer, ReplicationServerOptions};

use std::sync::atomic::{AtomicU64, Ordering};

use crate::entry::{EntryHeader, entry_size};
use crate::error::{DbError, DbResult};
use crate::shard::Shard;

/// Entry captured from the write path via SPSC channel.
pub struct ReplicationEntry {
    /// Raw serialized entry bytes (header + key + value + padding).
    pub data: Vec<u8>,
    /// Key length of the tree that produced this entry.
    pub key_len: u16,
}

/// Registry of the single replication target (tree) on a follower.
/// Each `Engine` belongs to exactly one `Tree`, so only one target exists.
pub struct ReplicationRegistry {
    target: Box<dyn ReplicationTarget>,
}

impl ReplicationRegistry {
    pub fn new(target: Box<dyn ReplicationTarget>) -> Self {
        target.mark_replication_follower();
        Self { target }
    }

    /// Apply a replicated entry in streaming mode (key_len known from SPSC).
    /// Writes raw bytes to the follower's shard, then updates the tree index.
    ///
    /// `last_applied_gsn` tracks the highest GSN successfully applied by the
    /// client. Entries with `gsn <= last_applied_gsn` are silently skipped
    /// (C12 defensive guard against stale resends from a buggy/old leader).
    pub fn apply_streaming(
        &self,
        shard: &Shard,
        entry: &ReplicationEntry,
        last_applied_gsn: &AtomicU64,
    ) -> DbResult<()> {
        // Pre-flight: header + key + value must fit in data. (C8)
        if entry.data.len() < HEADER_SIZE {
            return Err(DbError::Replication("entry header truncated".into()));
        }
        let header = parse_header(&entry.data)?;

        // C12 defensive: reject stale GSN before acquiring the shard lock.
        let cursor_gsn = last_applied_gsn.load(Ordering::Relaxed);
        if header.sequence() <= cursor_gsn {
            tracing::warn!(
                gsn = header.sequence(),
                cursor_gsn,
                "stale entry rejected (streaming)"
            );
            return Ok(());
        }

        let k = entry.key_len as usize;
        let required = HEADER_SIZE + k + header.value_len as usize;
        if entry.data.len() < required {
            return Err(DbError::Replication("entry data truncated".into()));
        }

        let (seq, needs_sync) = {
            let mut inner = shard.lock();
            let (file_id, entry_offset) =
                inner.append_raw_entry(shard.id, entry.key_len, &entry.data)?;

            let key = &entry.data[HEADER_SIZE..HEADER_SIZE + k];
            let value = &entry.data[HEADER_SIZE + k..HEADER_SIZE + k + header.value_len as usize];

            let outcome = self.target.apply_entry(
                &mut inner,
                shard.id,
                file_id,
                entry_offset,
                &header,
                key,
                value,
            )?;

            match &outcome {
                ApplyOutcome::Replaced(old) | ApplyOutcome::TombstoneRemoved(old) => {
                    let dead = entry_size(self.target.key_len(), old.len);
                    inner.add_dead_bytes(old.file_id, dead);
                }
                _ => {}
            }

            let seq = header.sequence();
            shard
                .gsn()
                .fetch_max(seq + 1, std::sync::atomic::Ordering::Relaxed);
            (seq, inner.should_sync())
        };

        if needs_sync {
            shard.flush()?;
        }

        // Advance after successful apply.
        last_applied_gsn.fetch_max(seq, Ordering::Relaxed);
        Ok(())
    }

    /// Apply a catch-up entry after validating its wire key length and CRC.
    ///
    /// `last_applied_gsn` tracks the highest GSN successfully applied by the
    /// client. Entries with `gsn <= last_applied_gsn` are silently skipped
    /// (C12 defensive guard against stale resends from a buggy/old leader).
    pub(crate) fn apply_catchup(
        &self,
        shard: &Shard,
        entry: &ReplicationEntry,
        last_applied_gsn: &AtomicU64,
    ) -> DbResult<()> {
        // Pre-flight: header + key + value must fit in data. (C8)
        if entry.data.len() < HEADER_SIZE {
            return Err(DbError::Replication("entry header truncated".into()));
        }
        let header = parse_header(&entry.data)?;

        // C12 defensive: reject stale GSN before acquiring the shard lock.
        let cursor_gsn = last_applied_gsn.load(Ordering::Relaxed);
        if header.sequence() <= cursor_gsn {
            tracing::warn!(
                gsn = header.sequence(),
                cursor_gsn,
                "stale entry rejected (catch-up)"
            );
            return Ok(());
        }

        let k = entry.key_len as usize;
        if k != self.target.key_len() {
            return Err(DbError::Replication(format!(
                "catch-up key length mismatch: expected {}, got {k}",
                self.target.key_len()
            )));
        }
        let required = HEADER_SIZE + k + header.value_len as usize;
        if entry.data.len() < required {
            return Err(DbError::Replication("entry data truncated".into()));
        }
        let key = &entry.data[HEADER_SIZE..HEADER_SIZE + k];
        let value = &entry.data[HEADER_SIZE + k..required];
        let actual_crc = crate::entry::compute_crc32(header.gsn, header.value_len, key, value);
        if actual_crc != header.crc32 {
            return Err(DbError::CrcMismatch {
                expected: header.crc32,
                actual: actual_crc,
            });
        }

        let (seq, needs_sync) = {
            let mut inner = shard.lock();
            let (file_id, entry_offset) =
                inner.append_raw_entry(shard.id, entry.key_len, &entry.data)?;

            let outcome = self.target.apply_entry(
                &mut inner,
                shard.id,
                file_id,
                entry_offset,
                &header,
                key,
                value,
            )?;

            match &outcome {
                ApplyOutcome::Replaced(old) | ApplyOutcome::TombstoneRemoved(old) => {
                    let dead = entry_size(self.target.key_len(), old.len);
                    inner.add_dead_bytes(old.file_id, dead);
                }
                _ => {}
            }

            let seq = header.sequence();
            shard
                .gsn()
                .fetch_max(seq + 1, std::sync::atomic::Ordering::Relaxed);
            (seq, inner.should_sync())
        };

        if needs_sync {
            shard.flush()?;
        }

        // Advance after successful apply.
        last_applied_gsn.fetch_max(seq, Ordering::Relaxed);
        Ok(())
    }
}

const HEADER_SIZE: usize = 16;

fn parse_header(data: &[u8]) -> DbResult<EntryHeader> {
    if data.len() < HEADER_SIZE {
        return Err(DbError::Replication("entry header truncated".into()));
    }
    use zerocopy::FromBytes;
    EntryHeader::read_from_bytes(&data[..HEADER_SIZE])
        .map_err(|_| DbError::Replication("invalid entry header".into()))
}

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

    #[test]
    fn parse_header_too_short_does_not_panic() {
        let short = vec![0u8; 5];
        let err = parse_header(&short).unwrap_err();
        assert!(matches!(err, DbError::Replication(_)));
    }
}

#[cfg(test)]
mod write_sync_tests {
    use std::collections::BTreeMap;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicU64, Ordering};

    use super::*;
    use crate::config::IoBackend;
    use crate::disk_loc::DiskLoc;
    use crate::entry::{compute_crc32, serialize_entry};
    use crate::flush_test;
    use crate::shard::ShardInner;

    #[derive(Debug, Clone, PartialEq, Eq)]
    struct AppliedValue {
        location: DiskLoc,
        value: [u8; 8],
        sequence: u64,
    }

    #[derive(Default)]
    struct TargetState {
        values: BTreeMap<[u8; 8], AppliedValue>,
    }

    #[derive(Default)]
    struct StatefulTarget {
        state: crate::sync::Mutex<TargetState>,
    }

    impl StatefulTarget {
        fn value(&self, key: [u8; 8]) -> Option<AppliedValue> {
            crate::sync::lock(&self.state).values.get(&key).cloned()
        }
    }

    impl ReplicationTarget for StatefulTarget {
        fn apply_entry(
            &self,
            _shard_inner: &mut ShardInner,
            _shard_id: u8,
            file_id: u32,
            entry_offset: u64,
            header: &EntryHeader,
            key: &[u8],
            value: &[u8],
        ) -> DbResult<ApplyOutcome> {
            let key: [u8; 8] = key
                .try_into()
                .map_err(|_| DbError::Replication("test target key length mismatch".into()))?;
            let value: [u8; 8] = value
                .try_into()
                .map_err(|_| DbError::Replication("test target value length mismatch".into()))?;
            let location = DiskLoc::new(
                file_id,
                (entry_offset + HEADER_SIZE as u64 + key.len() as u64) as u32,
                header.value_len,
            );
            let applied = AppliedValue {
                location,
                value,
                sequence: header.sequence(),
            };
            let old = crate::sync::lock(&self.state).values.insert(key, applied);
            Ok(old.map_or(ApplyOutcome::Inserted, |old| {
                ApplyOutcome::Replaced(old.location)
            }))
        }

        fn try_apply_entry(
            &self,
            shard_inner: &mut ShardInner,
            shard_id: u8,
            file_id: u32,
            entry_offset: u64,
            header: &EntryHeader,
            raw_after_header: &[u8],
        ) -> DbResult<ApplyOutcome> {
            let key_len = self.key_len();
            if raw_after_header.len() < key_len + header.value_len as usize {
                return Ok(ApplyOutcome::NotMatched);
            }
            let key = &raw_after_header[..key_len];
            let value = &raw_after_header[key_len..key_len + header.value_len as usize];
            if compute_crc32(header.gsn, header.value_len, key, value) != header.crc32 {
                return Ok(ApplyOutcome::NotMatched);
            }
            self.apply_entry(
                shard_inner,
                shard_id,
                file_id,
                entry_offset,
                header,
                key,
                value,
            )
        }

        fn key_len(&self) -> usize {
            size_of::<u64>()
        }
    }

    fn open_follower_shard(dir: &std::path::Path) -> Shard {
        Shard::open_with_sync_thresholds(
            0,
            dir,
            1 << 20,
            64 * 1024,
            false,
            false,
            false,
            IoBackend::default(),
            4,
            None,
            Arc::new(AtomicU64::new(0)),
        )
        .expect("open one-shard follower")
    }

    fn entry(sequence: u64) -> ReplicationEntry {
        ReplicationEntry {
            data: serialize_entry(
                sequence,
                &sequence.to_be_bytes(),
                &sequence.wrapping_mul(10).to_be_bytes(),
                false,
            ),
            key_len: size_of::<u64>() as u16,
        }
    }

    fn key(sequence: u64) -> [u8; 8] {
        sequence.to_be_bytes()
    }

    type Apply = fn(&ReplicationRegistry, &Shard, &ReplicationEntry, &AtomicU64) -> DbResult<()>;

    fn assert_failure_retry_keeps_watermark_and_retires_dead_bytes(apply: Apply) {
        flush_test::reset();
        let dir = tempfile::tempdir().expect("temporary follower directory");
        let shard = open_follower_shard(dir.path());
        let target = Arc::new(StatefulTarget::default());
        let registry = ReplicationRegistry::new(Box::new(Arc::clone(&target)));
        let last_applied_gsn = AtomicU64::new(0);

        for sequence in 1..=3 {
            apply(&registry, &shard, &entry(sequence), &last_applied_gsn)
                .expect("pre-threshold apply");
        }
        assert_eq!(
            flush_test::data_successes(),
            0,
            "three follower writes stay below the threshold"
        );
        assert_eq!(last_applied_gsn.load(Ordering::Relaxed), 3);

        let fourth = entry(4);
        let (first_seq4_file_id, first_seq4_offset) = {
            let inner = shard.lock();
            (inner.active.file_id, inner.active.write_offset)
        };
        let first_seq4_location = DiskLoc::new(
            first_seq4_file_id,
            (first_seq4_offset + HEADER_SIZE as u64 + size_of::<u64>() as u64) as u32,
            8,
        );
        flush_test::FAIL_NEXT_DATA_SYNC.store(true, Ordering::SeqCst);
        let error = apply(&registry, &shard, &fourth, &last_applied_gsn)
            .expect_err("the threshold barrier must surface the injected failure");
        assert!(
            format!("{error}").contains("test: injected data sync failure"),
            "apply must preserve the barrier error: {error}"
        );
        assert_eq!(last_applied_gsn.load(Ordering::Relaxed), 3);
        let failed = target
            .value(key(4))
            .expect("the index update precedes the failed durability barrier");
        assert_eq!(
            failed,
            AppliedValue {
                location: first_seq4_location,
                value: 40_u64.to_be_bytes(),
                sequence: 4,
            }
        );

        apply(&registry, &shard, &fourth, &last_applied_gsn)
            .expect("the retained watermark must permit a resend");
        assert_eq!(
            flush_test::data_successes(),
            1,
            "the resend completes exactly one follower data barrier"
        );
        assert_eq!(last_applied_gsn.load(Ordering::Relaxed), 4);
        assert_eq!(
            target.value(key(4)),
            Some(AppliedValue {
                location: DiskLoc::new(
                    first_seq4_file_id,
                    (first_seq4_offset
                        + entry_size(size_of::<u64>(), size_of::<u64>() as u32)
                        + HEADER_SIZE as u64
                        + size_of::<u64>() as u64) as u32,
                    8,
                ),
                value: 40_u64.to_be_bytes(),
                sequence: 4,
            }),
            "the resend replaces the first applied seq4 location"
        );
        assert_eq!(
            shard.lock().dead_bytes.get(&failed.location.file_id),
            Some(&entry_size(8, 8)),
            "the first seq4 location is retired after the resend"
        );
    }

    #[test]
    fn write_sync_repl_streaming_threshold() {
        assert_failure_retry_keeps_watermark_and_retires_dead_bytes(
            ReplicationRegistry::apply_streaming,
        );
    }

    #[test]
    fn write_sync_repl_catchup_threshold() {
        assert_failure_retry_keeps_watermark_and_retires_dead_bytes(
            ReplicationRegistry::apply_catchup,
        );
    }

    #[test]
    fn write_sync_repl_catchup_invalid_crc_rejects_before_append() {
        flush_test::reset();
        let dir = tempfile::tempdir().expect("temporary follower directory");
        let shard = open_follower_shard(dir.path());
        let target = Arc::new(StatefulTarget::default());
        let registry = ReplicationRegistry::new(Box::new(Arc::clone(&target)));
        let last_applied_gsn = AtomicU64::new(0);
        let gsn_before = shard.gsn().load(Ordering::Relaxed);
        let write_offset_before = shard.lock().active.write_offset;
        let mut invalid = entry(1);
        invalid.data[HEADER_SIZE + size_of::<u64>()] ^= 1;
        let error = registry
            .apply_catchup(&shard, &invalid, &last_applied_gsn)
            .expect_err("a bad catch-up CRC must be rejected before append");
        assert!(matches!(error, DbError::CrcMismatch { .. }));
        assert_eq!(
            target.value(key(1)),
            None,
            "a bad catch-up CRC must not reach the target"
        );
        assert_eq!(last_applied_gsn.load(Ordering::Relaxed), 0);
        assert_eq!(shard.gsn().load(Ordering::Relaxed), gsn_before);
        assert_eq!(shard.lock().active.write_offset, write_offset_before);
        assert_eq!(flush_test::data_successes(), 0);

        registry
            .apply_catchup(&shard, &entry(1), &last_applied_gsn)
            .expect("a valid resend of the rejected GSN must apply");
        assert_eq!(last_applied_gsn.load(Ordering::Relaxed), 1);
        assert_eq!(
            target.value(key(1)).map(|applied| applied.sequence),
            Some(1)
        );
    }
}

/// F1: SPSC overflow observability + recoverability.
///
/// The write path pushes each entry into a bounded SPSC ring for streaming; a
/// full ring drops the push. The fix (a) bumps a per-shard counter on every
/// dropped push so the server can *notice*, and (b) guarantees the dropped
/// entry's bytes are on disk so an overflow-recovery catch-up round
/// (`ShardLogReader`) reconstructs them in GSN order. This test drives the write
/// path directly against a deliberately tiny ring and asserts both properties —
/// the deterministic core of the server-side recovery loop, which the network
/// convergence path (timing-dependent to force overflow) builds on.
#[cfg(test)]
mod overflow_recovery_tests {
    use std::sync::Arc;
    use std::sync::atomic::AtomicU64;

    use super::ReplicationEntry;
    use super::log_reader::ShardLogReader;
    use crate::shard::Shard;

    #[test]
    fn ring_overflow_bumps_counter_and_keeps_entries_recoverable() {
        let dir = tempfile::tempdir().unwrap();
        let gsn = Arc::new(AtomicU64::new(0));
        let shard = Shard::open(
            0,
            dir.path(),
            1 << 24,    // 16 MiB — one file for the whole test
            256 * 1024, // write buffer
            false,
            false,
            false,
            crate::config::IoBackend::default(),
            gsn,
        )
        .expect("open shard");

        // Deliberately tiny ring, and we never pop from `_consumer`, so every
        // push past its capacity is dropped.
        let (producer, _consumer) = rtrb::RingBuffer::<ReplicationEntry>::new(4);
        shard.set_replication_producer(producer);

        let key_len: u16 = 8;
        let n: u64 = 20;
        {
            let mut inner = shard.lock();
            for i in 1..=n {
                let key = i.to_be_bytes();
                let value = [i as u8; 8];
                inner.append_entry(0, &key, &value, false).expect("append");
            }
        }

        // (a) Dropped pushes are observable via the per-shard counter.
        let overflow = shard.replication_overflow();
        assert!(
            overflow > 0,
            "a full ring must bump the overflow counter (got {overflow})"
        );

        // (b) Every entry — including those dropped from the ring — is on disk
        // and recoverable in GSN order by the overflow-recovery data path.
        shard.flush().expect("flush");
        let snapshot = shard.catchup_snapshot().expect("snapshot");
        let mut reader = ShardLogReader::new(snapshot, 0, key_len, usize::MAX, &mut || Ok(()))
            .expect("open log reader");

        let mut gsns = Vec::new();
        while let Some(e) = reader.next_entry().expect("read entry") {
            gsns.push(e.gsn);
        }

        // GSNs are assigned from the shared counter (seeded at 0), so the `n`
        // entries carry GSN 0..n-1.
        assert_eq!(
            gsns,
            (0..n).collect::<Vec<_>>(),
            "every ring-dropped entry must remain recoverable from disk in GSN order"
        );
    }
}