armdb 0.1.14

//! End-to-end integration tests for Fixed replication.
//!
//! These tests spin up a real `FixedReplicationServer` + `FixedReplicationClient`
//! talking over TCP, exercising the handshake, full-scan catch-up, and SPSC
//! streaming paths against real `FixedTree` engines.

#![cfg(feature = "replication")]

use std::net::TcpListener;
use std::sync::Arc;
use std::time::{Duration, Instant};

use armdb::FixedConfig;
use armdb::FixedTree;
use armdb::ShutdownSignal;
use armdb::fixed_replication::{
    FixedReplicationClient, FixedReplicationServer, FixedReplicationTarget,
};
use tempfile::TempDir;

// Keys are 8-byte big-endian u64s; values are 16 bytes.
const VLEN: usize = 16;
type K = [u8; 8];

fn pick_port() -> u16 {
    let l = TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port");
    l.local_addr().expect("local_addr").port()
}

fn test_config(shard_count: usize) -> FixedConfig {
    FixedConfig {
        shard_count,
        grow_step: 64,
        ..FixedConfig::test()
    }
}

fn kbe(i: u64) -> K {
    i.to_be_bytes()
}

fn vbuf(i: u64) -> [u8; VLEN] {
    let mut out = [0u8; VLEN];
    out[..8].copy_from_slice(&i.to_be_bytes());
    out[8..].copy_from_slice(&(i.wrapping_mul(31)).to_be_bytes());
    out
}

/// Poll `cond` until it returns true or `timeout` elapses. Returns `true`
/// on success, `false` otherwise.
fn wait_for(timeout: Duration, mut cond: impl FnMut() -> bool) -> bool {
    let start = Instant::now();
    while start.elapsed() < timeout {
        if cond() {
            return true;
        }
        std::thread::sleep(Duration::from_millis(20));
    }
    cond()
}

// ---------------------------------------------------------------------------
// test_initial_sync_empty_follower
// ---------------------------------------------------------------------------
#[test]
fn test_initial_sync_empty_follower() {
    // Leader: write 100 keys, delete 10. Fresh follower should converge to 90.
    let leader_dir = TempDir::new().unwrap();
    let leader = Arc::new(
        FixedTree::<K, VLEN>::open(leader_dir.path(), test_config(2)).expect("open leader"),
    );
    for i in 0..100u64 {
        leader.put(&kbe(i), &vbuf(i)).unwrap();
    }
    for i in 0..10u64 {
        leader.delete(&kbe(i)).unwrap();
    }
    leader.flush().unwrap();

    let port = pick_port();
    let bind_addr = format!("127.0.0.1:{port}").parse().unwrap();
    let stop_leader = ShutdownSignal::new();
    let server =
        FixedReplicationServer::start(bind_addr, leader.fixed_engine_access(), stop_leader.clone())
            .expect("start server");

    // Give the acceptor a moment to start listening.
    std::thread::sleep(Duration::from_millis(50));

    let follower_dir = TempDir::new().unwrap();
    let follower = Arc::new(
        FixedTree::<K, VLEN>::open(follower_dir.path(), test_config(2)).expect("open follower"),
    );
    let target: Arc<dyn FixedReplicationTarget> = follower.clone();

    let stop_follower = ShutdownSignal::new();
    let client = FixedReplicationClient::start(bind_addr, target, stop_follower.clone())
        .expect("start client");

    // Wait for convergence: keys 10..100 present.
    let f = follower.clone();
    assert!(
        wait_for(Duration::from_secs(15), || {
            (10..100u64).all(|i| f.get(&kbe(i)) == Some(vbuf(i)))
        }),
        "follower did not converge: len={}",
        follower.len()
    );

    // Keys 0..10 must be absent on follower.
    for i in 0..10u64 {
        assert!(
            follower.get(&kbe(i)).is_none(),
            "deleted key {i} leaked to follower"
        );
    }

    stop_follower.shutdown();
    drop(client);
    stop_leader.shutdown();
    drop(server);
}

// ---------------------------------------------------------------------------
// test_streaming_live_writes
// ---------------------------------------------------------------------------
#[test]
fn test_streaming_live_writes() {
    let leader_dir = TempDir::new().unwrap();
    let leader = Arc::new(
        FixedTree::<K, VLEN>::open(leader_dir.path(), test_config(2)).expect("open leader"),
    );

    let port = pick_port();
    let bind_addr = format!("127.0.0.1:{port}").parse().unwrap();
    let stop_leader = ShutdownSignal::new();
    let server =
        FixedReplicationServer::start(bind_addr, leader.fixed_engine_access(), stop_leader.clone())
            .expect("start server");

    std::thread::sleep(Duration::from_millis(50));

    let follower_dir = TempDir::new().unwrap();
    let follower = Arc::new(
        FixedTree::<K, VLEN>::open(follower_dir.path(), test_config(2)).expect("open follower"),
    );
    let target: Arc<dyn FixedReplicationTarget> = follower.clone();
    let stop_follower = ShutdownSignal::new();
    let client = FixedReplicationClient::start(bind_addr, target, stop_follower.clone())
        .expect("start client");

    // Let the handshake + empty phase-1 complete before writing.
    std::thread::sleep(Duration::from_millis(300));

    for i in 0..500u64 {
        leader.put(&kbe(i), &vbuf(i)).unwrap();
    }
    leader.flush().unwrap();

    let f = follower.clone();
    assert!(
        wait_for(Duration::from_secs(20), || {
            (0..500u64).all(|i| f.get(&kbe(i)) == Some(vbuf(i)))
        }),
        "streaming did not converge: follower len={}",
        follower.len()
    );

    stop_follower.shutdown();
    drop(client);
    stop_leader.shutdown();
    drop(server);
}

// ---------------------------------------------------------------------------
// test_reconnect_after_disconnect
// ---------------------------------------------------------------------------
#[test]
fn test_reconnect_after_disconnect() {
    let leader_dir = TempDir::new().unwrap();
    let leader = Arc::new(
        FixedTree::<K, VLEN>::open(leader_dir.path(), test_config(2)).expect("open leader"),
    );
    for i in 0..50u64 {
        leader.put(&kbe(i), &vbuf(i)).unwrap();
    }
    leader.flush().unwrap();

    let port = pick_port();
    let bind_addr = format!("127.0.0.1:{port}").parse().unwrap();
    let stop_leader = ShutdownSignal::new();
    let server =
        FixedReplicationServer::start(bind_addr, leader.fixed_engine_access(), stop_leader.clone())
            .expect("start server");

    std::thread::sleep(Duration::from_millis(50));

    let follower_dir = TempDir::new().unwrap();
    let follower = Arc::new(
        FixedTree::<K, VLEN>::open(follower_dir.path(), test_config(2)).expect("open follower"),
    );

    // First session: connect, catch up 50 keys.
    let stop_follower1 = ShutdownSignal::new();
    let target1: Arc<dyn FixedReplicationTarget> = follower.clone();
    let client1 = FixedReplicationClient::start(bind_addr, target1, stop_follower1.clone())
        .expect("start client1");

    let f = follower.clone();
    assert!(
        wait_for(Duration::from_secs(15), || {
            (0..50u64).all(|i| f.get(&kbe(i)) == Some(vbuf(i)))
        }),
        "first-phase catch-up failed"
    );

    // Disconnect the follower.
    stop_follower1.shutdown();
    drop(client1);

    // Leader writes 50 more keys while follower is offline.
    for i in 50..100u64 {
        leader.put(&kbe(i), &vbuf(i)).unwrap();
    }
    leader.flush().unwrap();

    // Reconnect a fresh client to the same follower engine.
    let stop_follower2 = ShutdownSignal::new();
    let target2: Arc<dyn FixedReplicationTarget> = follower.clone();
    let client2 = FixedReplicationClient::start(bind_addr, target2, stop_follower2.clone())
        .expect("start client2");

    let f = follower.clone();
    assert!(
        wait_for(Duration::from_secs(15), || {
            (0..100u64).all(|i| f.get(&kbe(i)) == Some(vbuf(i)))
        }),
        "reconnect did not converge: len={}",
        follower.len()
    );

    stop_follower2.shutdown();
    drop(client2);
    stop_leader.shutdown();
    drop(server);
}

// ---------------------------------------------------------------------------
// test_topology_mismatch_errors
// ---------------------------------------------------------------------------
#[test]
fn test_topology_mismatch_errors() {
    // Leader: 2 shards. Follower: 4 shards. Client must refuse to apply.
    let leader_dir = TempDir::new().unwrap();
    let leader = Arc::new(
        FixedTree::<K, VLEN>::open(leader_dir.path(), test_config(2)).expect("open leader"),
    );
    for i in 0..20u64 {
        leader.put(&kbe(i), &vbuf(i)).unwrap();
    }
    leader.flush().unwrap();

    let port = pick_port();
    let bind_addr = format!("127.0.0.1:{port}").parse().unwrap();
    let stop_leader = ShutdownSignal::new();
    let server =
        FixedReplicationServer::start(bind_addr, leader.fixed_engine_access(), stop_leader.clone())
            .expect("start server");

    std::thread::sleep(Duration::from_millis(50));

    let follower_dir = TempDir::new().unwrap();
    let follower = Arc::new(
        FixedTree::<K, VLEN>::open(follower_dir.path(), test_config(4))
            .expect("open follower with 4 shards"),
    );
    let target: Arc<dyn FixedReplicationTarget> = follower.clone();
    let stop_follower = ShutdownSignal::new();
    let client = FixedReplicationClient::start(bind_addr, target, stop_follower.clone())
        .expect("start client");

    // Give the client time to handshake and reject several times.
    std::thread::sleep(Duration::from_secs(2));

    // Follower state must remain empty — nothing got applied.
    assert_eq!(
        follower.len(),
        0,
        "topology-mismatched follower applied events"
    );

    stop_follower.shutdown();
    drop(client);
    stop_leader.shutdown();
    drop(server);
}

// ---------------------------------------------------------------------------
// test_follower_restart_clean_uses_sidecar
// ---------------------------------------------------------------------------
#[test]
fn test_follower_restart_clean_uses_sidecar() {
    // Sync a follower, shut it down cleanly (via close), confirm the sidecar
    // exists, reopen and verify data is still there.
    let leader_dir = TempDir::new().unwrap();
    let leader = Arc::new(
        FixedTree::<K, VLEN>::open(leader_dir.path(), test_config(2)).expect("open leader"),
    );
    for i in 0..80u64 {
        leader.put(&kbe(i), &vbuf(i)).unwrap();
    }
    leader.flush().unwrap();

    let port = pick_port();
    let bind_addr = format!("127.0.0.1:{port}").parse().unwrap();
    let stop_leader = ShutdownSignal::new();
    let server =
        FixedReplicationServer::start(bind_addr, leader.fixed_engine_access(), stop_leader.clone())
            .expect("start server");

    std::thread::sleep(Duration::from_millis(50));

    let follower_dir = TempDir::new().unwrap();
    let follower = Arc::new(
        FixedTree::<K, VLEN>::open(follower_dir.path(), test_config(2)).expect("open follower"),
    );
    let target: Arc<dyn FixedReplicationTarget> = follower.clone();
    let stop_follower = ShutdownSignal::new();
    let client = FixedReplicationClient::start(bind_addr, target, stop_follower.clone())
        .expect("start client");

    // Wait for convergence (use a borrow, no extra Arc clone).
    assert!(
        wait_for(Duration::from_secs(15), || {
            (0..80u64).all(|i| follower.get(&kbe(i)) == Some(vbuf(i)))
        }),
        "initial sync did not converge"
    );

    // Tear down networking first so the client releases its target Arc.
    stop_follower.shutdown();
    drop(client);
    stop_leader.shutdown();
    drop(server);

    // Give background threads a moment to release the target Arc.
    for _ in 0..50 {
        if Arc::strong_count(&follower) == 1 {
            break;
        }
        std::thread::sleep(Duration::from_millis(20));
    }

    // Close follower cleanly to write versions sidecar.
    let follower = Arc::try_unwrap(follower).unwrap_or_else(|arc| {
        panic!(
            "follower still has outstanding Arc refs: strong_count={}",
            Arc::strong_count(&arc)
        )
    });
    follower.close().expect("clean close of follower");

    // Verify sidecars exist.
    let follower_root = follower_dir.path();
    for shard_id in 0..2 {
        let shard_dir = follower_root.join(format!("shard_{:03}", shard_id));
        let sidecar = shard_dir.join("fixed.versions");
        assert!(
            sidecar.exists(),
            "missing versions sidecar for shard {shard_id}: {sidecar:?}"
        );
    }

    // Reopen follower — recovery should use the sidecar and data is intact.
    let reopened = FixedTree::<K, VLEN>::open(follower_root, test_config(2))
        .expect("reopen follower after clean close");
    for i in 0..80u64 {
        assert_eq!(
            reopened.get(&kbe(i)),
            Some(vbuf(i)),
            "key {i} missing after reopen via sidecar"
        );
    }
    assert_eq!(reopened.len(), 80);
}

// ---------------------------------------------------------------------------
// test_catchup_multibatch_does_not_stall
//
// Regression test for the Phase-1 catch-up stall: before the fix the
// follower only Acked every `ACK_INTERVAL = 1000` events, but the server
// sends batches of at most `BATCH_MAX_ENTRIES = 256`. That meant the
// follower never produced an Ack for the first 3 batches, so the server
// stalled on `read_frame` until the follower's 5s heartbeat fired.
// Catch-up of N batches took ~5s * ceil(3N/4).
//
// We verify the fix by loading a single shard with >3 batches (800 keys
// over one shard ≥ 4 batches) and asserting convergence completes in well
// under the pre-fix lower bound.
// ---------------------------------------------------------------------------
#[test]
fn test_catchup_multibatch_does_not_stall() {
    const KEYS: u64 = 800; // > 3 * BATCH_MAX_ENTRIES (=256)

    let leader_dir = TempDir::new().unwrap();
    // Single shard so all keys land on one shard → multi-batch Phase-1 scan.
    let leader = Arc::new(
        FixedTree::<K, VLEN>::open(leader_dir.path(), test_config(1)).expect("open leader"),
    );
    for i in 0..KEYS {
        leader.put(&kbe(i), &vbuf(i)).unwrap();
    }
    leader.flush().unwrap();

    let port = pick_port();
    let bind_addr = format!("127.0.0.1:{port}").parse().unwrap();
    let stop_leader = ShutdownSignal::new();
    let server =
        FixedReplicationServer::start(bind_addr, leader.fixed_engine_access(), stop_leader.clone())
            .expect("start server");

    std::thread::sleep(Duration::from_millis(50));

    let follower_dir = TempDir::new().unwrap();
    let follower = Arc::new(
        FixedTree::<K, VLEN>::open(follower_dir.path(), test_config(1)).expect("open follower"),
    );
    let target: Arc<dyn FixedReplicationTarget> = follower.clone();

    let stop_follower = ShutdownSignal::new();
    let start = Instant::now();
    let client = FixedReplicationClient::start(bind_addr, target, stop_follower.clone())
        .expect("start client");

    let f = follower.clone();
    assert!(
        wait_for(Duration::from_secs(5), || f.len() as u64 == KEYS),
        "catch-up did not converge within 5s: len={}, elapsed={:?}",
        follower.len(),
        start.elapsed()
    );
    let elapsed = start.elapsed();

    // With the fix catch-up completes in well under a second on any reasonable
    // machine; the pre-fix bug stalled ≥5s just for the heartbeat-driven Acks.
    // We give a generous 3s ceiling to keep the test robust on slow CI.
    assert!(
        elapsed < Duration::from_secs(3),
        "catch-up took {elapsed:?} — suggests Phase-1 is stalling on Acks"
    );

    stop_follower.shutdown();
    drop(client);
    stop_leader.shutdown();
    drop(server);
}

// ---------------------------------------------------------------------------
// test_protocol_version_mismatch_rejected
//
// The server must reject a SyncRequest whose `protocol_version` disagrees
// with its own `PROTOCOL_VERSION`, sending an Error frame before ShardInfo.
// This is a smoke test that a hand-crafted mismatched SyncRequest elicits
// an Error response on the wire.
// ---------------------------------------------------------------------------
#[test]
fn test_protocol_version_mismatch_rejected() {
    use armdb::fixed_replication::protocol::{
        FixedMessageType, SyncRequest, decode_error, read_frame, write_frame,
    };
    use std::io::BufReader;
    use std::net::TcpStream;

    let leader_dir = TempDir::new().unwrap();
    let leader = Arc::new(
        FixedTree::<K, VLEN>::open(leader_dir.path(), test_config(1)).expect("open leader"),
    );

    let port = pick_port();
    let bind_addr: std::net::SocketAddr = format!("127.0.0.1:{port}").parse().unwrap();
    let stop_leader = ShutdownSignal::new();
    let server =
        FixedReplicationServer::start(bind_addr, leader.fixed_engine_access(), stop_leader.clone())
            .expect("start server");

    std::thread::sleep(Duration::from_millis(50));

    let mut stream = TcpStream::connect(bind_addr).expect("connect to leader");
    stream.set_nodelay(true).ok();
    let mut reader = BufReader::new(stream.try_clone().expect("clone"));

    // Hand-craft a SyncRequest claiming an unsupported protocol_version.
    let bogus = SyncRequest {
        shard_id: 0,
        protocol_version: 99,
        flags: 0,
    };
    write_frame(&mut stream, &bogus.encode()).expect("write frame");

    // Server should respond with an Error frame before closing.
    let frame = read_frame(&mut reader).expect("read error frame");
    assert_eq!(frame.msg_type, FixedMessageType::Error);
    let msg = decode_error(&frame.payload);
    assert!(
        msg.contains("protocol version"),
        "unexpected error message: {msg}"
    );

    stop_leader.shutdown();
    drop(server);
}