armdb 0.7.0

//! Integration tests for armdb/src/replication/ (Bitcask variable replication).
//!
//! Each test spins up a real `VarTree` leader + server and a follower + client
//! talking over TCP, exercises catch-up or streaming, then asserts convergence.

#![cfg(all(feature = "replication", feature = "var-collections"))]

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

use armdb::Config;
use armdb::ShutdownSignal;
use armdb::VarTree;
use armdb::replication::{ReplicationClientOptions, ReplicationRegistry, ReplicationServerOptions};
use armdb::serialize_entry;

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

/// Bind to an ephemeral port and immediately release the listener so the
/// OS keeps the port reserved until we bind again. This is racy by design —
/// use only in tests where the window is tiny.
fn next_bind_addr() -> SocketAddr {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap();
    drop(listener);
    addr
}

/// Poll `f` every 20 ms until it returns `true` or `timeout` elapses.
fn wait_until(timeout: Duration, mut f: impl FnMut() -> bool) -> bool {
    let start = Instant::now();
    loop {
        if f() {
            return true;
        }
        if start.elapsed() >= timeout {
            return f(); // one last check
        }
        std::thread::sleep(Duration::from_millis(20));
    }
}

/// Config with a small shard count to keep fd usage low in tests.
fn test_cfg() -> Config {
    Config::test()
}

// ---------------------------------------------------------------------------
// Test 1: catch_up_smoke  (C1, C10)
// ---------------------------------------------------------------------------
/// Leader writes N > BATCH_MAX_ENTRIES (256) entries before the follower
/// connects. Follower must catch up via the log-reader path and end up with
/// all 512 entries.
#[test]
fn catch_up_smoke() {
    let leader_dir = tempfile::tempdir().unwrap();
    let follower_dir = tempfile::tempdir().unwrap();
    let addr = next_bind_addr();

    let leader = VarTree::<[u8; 8]>::open(leader_dir.path(), test_cfg()).unwrap();
    for i in 0u64..512 {
        leader
            .put(&i.to_be_bytes(), format!("v{i}").as_bytes())
            .unwrap();
    }

    let leader_signal = ShutdownSignal::new();
    let _server = leader
        .start_replication_server(addr, leader_signal.clone())
        .unwrap();

    let follower = Arc::new(VarTree::<[u8; 8]>::open(follower_dir.path(), test_cfg()).unwrap());
    let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));
    let follower_signal = ShutdownSignal::new();
    let _client = follower
        .start_replication_client(addr, registry, follower_signal.clone())
        .unwrap();

    assert!(
        wait_until(Duration::from_secs(15), || {
            (0u64..512).all(|i| follower.contains(&i.to_be_bytes()))
        }),
        "follower did not catch up within 15s (len={})",
        follower.len()
    );

    follower_signal.shutdown();
    leader_signal.shutdown();
}

// ---------------------------------------------------------------------------
// Test 2: streaming_after_catchup
// ---------------------------------------------------------------------------
/// After initial catch-up, the leader writes more entries. Follower must
/// receive them via the streaming path.
#[test]
fn streaming_after_catchup() {
    let leader_dir = tempfile::tempdir().unwrap();
    let follower_dir = tempfile::tempdir().unwrap();
    let addr = next_bind_addr();

    let leader = VarTree::<[u8; 8]>::open(leader_dir.path(), test_cfg()).unwrap();
    // Write the first batch before the follower connects (catch-up phase).
    for i in 0u64..64 {
        leader.put(&i.to_be_bytes(), b"initial").unwrap();
    }

    let leader_signal = ShutdownSignal::new();
    let _server = leader
        .start_replication_server(addr, leader_signal.clone())
        .unwrap();

    let follower = Arc::new(VarTree::<[u8; 8]>::open(follower_dir.path(), test_cfg()).unwrap());
    let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));
    let follower_signal = ShutdownSignal::new();
    let _client = follower
        .start_replication_client(addr, registry, follower_signal.clone())
        .unwrap();

    // Wait for catch-up to complete.
    assert!(
        wait_until(Duration::from_secs(10), || { follower.len() >= 64 }),
        "catch-up did not complete within 10s"
    );

    // Now write streaming entries (512..768).
    for i in 512u64..768 {
        leader.put(&i.to_be_bytes(), b"streamed").unwrap();
    }

    assert!(
        wait_until(Duration::from_secs(5), || {
            (512u64..768).all(|i| follower.contains(&i.to_be_bytes()))
        }),
        "follower did not receive streamed entries within 5s (len={})",
        follower.len()
    );

    follower_signal.shutdown();
    leader_signal.shutdown();
}

// ---------------------------------------------------------------------------
// Test 3: reconnect_no_loss_no_duplicates  (C11, C12)
// ---------------------------------------------------------------------------
/// Drop the client after catch-up, write more on leader, reconnect — follower
/// must end up with all entries and no logical duplicates (index len == 200).
#[test]
fn reconnect_no_loss_no_duplicates() {
    let leader_dir = tempfile::tempdir().unwrap();
    let follower_dir = tempfile::tempdir().unwrap();
    let addr = next_bind_addr();

    let leader = VarTree::<[u8; 8]>::open(leader_dir.path(), test_cfg()).unwrap();

    let leader_signal = ShutdownSignal::new();
    let _server = leader
        .start_replication_server_with_options(
            addr,
            leader_signal.clone(),
            ReplicationServerOptions {
                heartbeat_interval_secs: 1,
                ..Default::default()
            },
        )
        .unwrap();

    // --- First connection ---
    let follower = Arc::new(VarTree::<[u8; 8]>::open(follower_dir.path(), test_cfg()).unwrap());
    let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));
    let signal1 = ShutdownSignal::new();
    let client1 = follower
        .start_replication_client(addr, registry.clone(), signal1.clone())
        .unwrap();

    // Write first 100 entries.
    for i in 0u64..100 {
        leader.put(&i.to_be_bytes(), b"batch1").unwrap();
    }

    // Wait for first batch.
    assert!(
        wait_until(Duration::from_secs(10), || follower.len() >= 100),
        "follower did not receive first 100 entries"
    );

    // Drop the client (triggers ShutdownSignal via Drop).
    drop(client1);
    signal1.shutdown();

    // Give the server time to notice the disconnect.
    std::thread::sleep(Duration::from_millis(50));

    // Write second 100 entries while disconnected.
    for i in 100u64..200 {
        leader.put(&i.to_be_bytes(), b"batch2").unwrap();
    }

    // --- Second connection (same follower dir = cursor is loaded) ---
    let signal2 = ShutdownSignal::new();
    let _client2 = follower
        .start_replication_client_with_options(
            addr,
            registry,
            signal2.clone(),
            ReplicationClientOptions {
                reconnect_base_ms: 50,
                reconnect_max_ms: 200,
                ..Default::default()
            },
        )
        .unwrap();

    // Follower must end up with exactly 200 unique keys.
    assert!(
        wait_until(Duration::from_secs(15), || follower.len() >= 200),
        "follower did not receive all 200 entries after reconnect (len={})",
        follower.len()
    );

    // No logical duplicates: index reports exactly 200.
    assert_eq!(
        follower.len(),
        200,
        "expected 200 unique entries, got {}",
        follower.len()
    );

    signal2.shutdown();
    leader_signal.shutdown();
}

#[test]
fn reconnect_uses_in_memory_last_applied() {
    use armdb::replication::protocol::{
        EntryBatch, ShardInfo, SyncRequest, VAR_PROTOCOL_VERSION, WireEntry, read_frame,
        write_frame,
    };
    use std::io::BufReader;

    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let addr = listener.local_addr().unwrap();
    let mock = std::thread::spawn(move || {
        for (round, expected_from) in [(0, 1), (1, 11)] {
            let (stream, _) = listener.accept().unwrap();
            let mut reader = BufReader::new(stream.try_clone().unwrap());
            let mut writer = stream;
            let request = read_frame(&mut reader).unwrap();
            let request = SyncRequest::decode(&request.payload).unwrap();
            assert_eq!(request.from_gsn, expected_from);
            write_frame(
                &mut writer,
                &ShardInfo {
                    protocol_version: VAR_PROTOCOL_VERSION,
                    shard_count: 1,
                    max_file_size: 256 * 1024 * 1024,
                }
                .encode(),
            )
            .unwrap();
            let entries = ((round * 10 + 1)..=(round * 10 + 10))
                .map(|gsn| {
                    let key = (gsn as u64).to_be_bytes();
                    let data = serialize_entry(gsn as u64, &key, b"v", false);
                    WireEntry {
                        entry_len: data.len() as u32,
                        key_len: 8,
                        gsn: gsn as u64,
                        data,
                    }
                })
                .collect();
            write_frame(
                &mut writer,
                &EntryBatch {
                    shard_id: 0,
                    entries,
                }
                .encode(),
            )
            .unwrap();
        }
    });

    let dir = tempfile::tempdir().unwrap();
    let follower = Arc::new(
        VarTree::<[u8; 8]>::open(
            dir.path(),
            Config::balanced().shard_count(1).hints(true).build(),
        )
        .unwrap(),
    );
    let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));
    let signal = ShutdownSignal::new();
    let _client = follower
        .start_replication_client_with_options(
            addr,
            registry,
            signal.clone(),
            ReplicationClientOptions {
                reconnect_base_ms: 20,
                reconnect_max_ms: 20,
                heartbeat_interval_secs: 1,
            },
        )
        .unwrap();
    assert!(wait_until(Duration::from_secs(5), || follower.len() == 20));
    signal.shutdown();
    mock.join().unwrap();
}

// ---------------------------------------------------------------------------
// Test 4: rotation_during_streaming  (C2, C17)
// ---------------------------------------------------------------------------
/// Small max_file_size forces many log rotations while the follower streams.
/// Follower's index must resolve all entries correctly via updated file_id.
#[test]
fn rotation_during_streaming() {
    let leader_dir = tempfile::tempdir().unwrap();
    let follower_dir = tempfile::tempdir().unwrap();
    let addr = next_bind_addr();

    // 32 KiB per file → many rotations with 1000 × ~40-byte entries.
    let cfg = Config {
        shard_count: 2,
        max_file_size: 32 * 1024,
        write_buffer_size: 16 * 1024,
        ..Config::balanced().build()
    };

    let leader = VarTree::<[u8; 8]>::open(leader_dir.path(), cfg.clone()).unwrap();

    let leader_signal = ShutdownSignal::new();
    let _server = leader
        .start_replication_server(addr, leader_signal.clone())
        .unwrap();

    let follower = Arc::new(VarTree::<[u8; 8]>::open(follower_dir.path(), cfg).unwrap());
    let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));
    let follower_signal = ShutdownSignal::new();
    let _client = follower
        .start_replication_client(addr, registry, follower_signal.clone())
        .unwrap();

    // Write ~1000 entries with 32-byte values — enough to trigger many rotations.
    let value = b"rotation_test_value_32bytes_xxxxx";
    for i in 0u64..1000 {
        leader.put(&i.to_be_bytes(), value).unwrap();
    }

    assert!(
        wait_until(Duration::from_secs(15), || { follower.len() >= 1000 }),
        "follower did not catch up after rotations (len={})",
        follower.len()
    );

    // Spot-check: values must be readable (not stale DiskLoc).
    for i in [0u64, 499, 999] {
        assert!(
            follower.get(&i.to_be_bytes()).is_some(),
            "key {i} missing from follower after rotation"
        );
    }

    follower_signal.shutdown();
    leader_signal.shutdown();
}

// ---------------------------------------------------------------------------
// Test 5: multi_shard
// ---------------------------------------------------------------------------
/// shard_count=4 with shard_prefix_bits=0 (default hash routing). Each shard
/// streams independently via one thread per shard in the client.
#[test]
fn multi_shard() {
    let leader_dir = tempfile::tempdir().unwrap();
    let follower_dir = tempfile::tempdir().unwrap();
    let addr = next_bind_addr();

    let cfg = Config {
        shard_count: 4,
        ..Config::balanced().build()
    };

    let leader = VarTree::<[u8; 8]>::open(leader_dir.path(), cfg.clone()).unwrap();
    // Write entries that will spread across shards via hash routing.
    for i in 0u64..200 {
        leader.put(&i.to_be_bytes(), b"multi").unwrap();
    }

    let leader_signal = ShutdownSignal::new();
    let _server = leader
        .start_replication_server(addr, leader_signal.clone())
        .unwrap();

    let follower = Arc::new(VarTree::<[u8; 8]>::open(follower_dir.path(), cfg).unwrap());
    let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));
    let follower_signal = ShutdownSignal::new();
    let _client = follower
        .start_replication_client(addr, registry, follower_signal.clone())
        .unwrap();

    assert!(
        wait_until(Duration::from_secs(15), || {
            (0u64..200).all(|i| follower.contains(&i.to_be_bytes()))
        }),
        "multi-shard follower did not receive all entries (len={})",
        follower.len()
    );

    follower_signal.shutdown();
    leader_signal.shutdown();
}

// ---------------------------------------------------------------------------
// Test 6: encrypted_catchup  (C3, A6)
// ---------------------------------------------------------------------------
/// Encryption enabled; entries written without explicit flush. The server's
/// A6 flush rule pads the trailing encrypted page before the log reader runs,
/// so catch-up must deliver every entry.
#[cfg(feature = "encryption")]
#[test]
fn encrypted_catchup() {
    let leader_dir = tempfile::tempdir().unwrap();
    let follower_dir = tempfile::tempdir().unwrap();
    let addr = next_bind_addr();

    let cfg = Config {
        shard_count: 2,
        write_buffer_size: 8192,
        #[cfg(feature = "encryption")]
        encryption_key: Some([0x42u8; 32]),
        ..Config::balanced().build()
    };

    let leader = VarTree::<[u8; 8]>::open(leader_dir.path(), cfg.clone()).unwrap();
    // Write entries WITHOUT a manual flush — A6 flush must handle this.
    for i in 0u64..100 {
        leader.put(&i.to_be_bytes(), b"encrypted").unwrap();
    }
    // Intentionally no flush here — tests the A6 automatic flush path.

    let leader_signal = ShutdownSignal::new();
    let _server = leader
        .start_replication_server(addr, leader_signal.clone())
        .unwrap();

    // Follower must also open with the same encryption key.
    let follower = Arc::new(VarTree::<[u8; 8]>::open(follower_dir.path(), cfg).unwrap());
    let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));
    let follower_signal = ShutdownSignal::new();
    let _client = follower
        .start_replication_client(addr, registry, follower_signal.clone())
        .unwrap();

    assert!(
        wait_until(Duration::from_secs(15), || {
            (0u64..100).all(|i| follower.contains(&i.to_be_bytes()))
        }),
        "encrypted follower did not catch up (len={})",
        follower.len()
    );

    follower_signal.shutdown();
    leader_signal.shutdown();
}

// ---------------------------------------------------------------------------
// Test 7: reject_second_follower  (A2)
// ---------------------------------------------------------------------------
/// Two clients connect to the same shard. The second receives an Error frame
/// (its `start_replication_client` call succeeds — the error is delivered
/// over the wire and logged, then the client's thread backs off). The second
/// follower must remain empty while the first converges.
#[test]
fn reject_second_follower() {
    let leader_dir = tempfile::tempdir().unwrap();
    let follower1_dir = tempfile::tempdir().unwrap();
    let follower2_dir = tempfile::tempdir().unwrap();
    let addr = next_bind_addr();

    let leader = VarTree::<[u8; 8]>::open(leader_dir.path(), test_cfg()).unwrap();

    let leader_signal = ShutdownSignal::new();
    let _server = leader
        .start_replication_server(addr, leader_signal.clone())
        .unwrap();

    // First follower connects and claims the consumer.
    let follower1 = Arc::new(VarTree::<[u8; 8]>::open(follower1_dir.path(), test_cfg()).unwrap());
    let registry1 = Arc::new(ReplicationRegistry::new(follower1.as_replication_target()));
    let signal1 = ShutdownSignal::new();
    let _client1 = follower1
        .start_replication_client(addr, registry1, signal1.clone())
        .unwrap();

    // Give client1 time to connect and claim the shard consumer.
    std::thread::sleep(Duration::from_millis(300));

    // Second follower attempts to connect — server will send Error frame.
    let follower2 = Arc::new(VarTree::<[u8; 8]>::open(follower2_dir.path(), test_cfg()).unwrap());
    let registry2 = Arc::new(ReplicationRegistry::new(follower2.as_replication_target()));
    let signal2 = ShutdownSignal::new();
    let _client2 = follower2
        .start_replication_client(addr, registry2, signal2.clone())
        .unwrap();

    // Write entries on leader.
    for i in 0u64..50 {
        leader
            .put(&i.to_be_bytes(), format!("v{i}").as_bytes())
            .unwrap();
    }

    // follower1 must receive all entries.
    assert!(
        wait_until(Duration::from_secs(10), || { follower1.len() >= 50 }),
        "follower1 did not receive entries (len={})",
        follower1.len()
    );

    // follower2 must remain empty for the duration. Give it 3s to confirm
    // it has NOT received any entries (the Error frame keeps it from applying
    // anything during the backoff window).
    std::thread::sleep(Duration::from_secs(3));
    assert_eq!(
        follower2.len(),
        0,
        "follower2 should be empty but has {} entries",
        follower2.len()
    );

    signal1.shutdown();
    signal2.shutdown();
    leader_signal.shutdown();
}

#[test]
fn low_rate_stream_advances_min_replicated_gsn() {
    use std::sync::atomic::Ordering;

    let leader_dir = tempfile::tempdir().unwrap();
    let follower_dir = tempfile::tempdir().unwrap();
    let addr = next_bind_addr();
    let config = Config::balanced().shard_count(1).hints(true).build();
    let leader = VarTree::<[u8; 8]>::open(leader_dir.path(), config.clone()).unwrap();
    leader.put(&0u64.to_be_bytes(), b"seed").unwrap();
    let leader_signal = ShutdownSignal::new();
    let server = leader
        .start_replication_server_with_options(
            addr,
            leader_signal.clone(),
            ReplicationServerOptions {
                heartbeat_interval_secs: 1,
                ..Default::default()
            },
        )
        .unwrap();
    let follower = Arc::new(VarTree::<[u8; 8]>::open(follower_dir.path(), config).unwrap());
    let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));
    let follower_signal = ShutdownSignal::new();
    let _client = follower
        .start_replication_client(addr, registry, follower_signal.clone())
        .unwrap();
    assert!(wait_until(Duration::from_secs(5), || follower.len() == 1));
    assert!(wait_until(Duration::from_secs(5), || {
        server.min_replicated_gsn[0].load(Ordering::Relaxed) >= 1
    }));
    // The seed proves the connection completed catch-up; the following three
    // writes are therefore below the streaming ACK interval.
    std::thread::sleep(Duration::from_millis(100));
    for key in 1u64..=3 {
        leader.put(&key.to_be_bytes(), b"slow").unwrap();
    }
    assert!(wait_until(Duration::from_secs(5), || follower.len() == 4));
    assert!(wait_until(Duration::from_secs(5), || {
        server.min_replicated_gsn[0].load(Ordering::Relaxed) >= 4
    }));
    follower_signal.shutdown();
    leader_signal.shutdown();
}

#[test]
fn variable_protocol_mismatch_returns_typed_error() {
    use armdb::replication::protocol::{
        MessageType, ReplicationErrorCode, SyncRequest, VAR_PROTOCOL_VERSION, decode_error,
        read_frame, write_frame,
    };
    use std::io::BufReader;
    use std::net::TcpStream;

    let leader_dir = tempfile::tempdir().unwrap();
    let leader = VarTree::<[u8; 8]>::open(
        leader_dir.path(),
        Config::balanced().shard_count(1).hints(true).build(),
    )
    .unwrap();
    let addr = next_bind_addr();
    let signal = ShutdownSignal::new();
    let _server = leader
        .start_replication_server(addr, signal.clone())
        .unwrap();

    let stream = TcpStream::connect(addr).unwrap();
    stream
        .set_read_timeout(Some(Duration::from_secs(5)))
        .unwrap();
    let mut reader = BufReader::new(stream.try_clone().unwrap());
    let mut writer = stream;
    let req = SyncRequest {
        protocol_version: VAR_PROTOCOL_VERSION + 1,
        shard_id: 0,
        from_gsn: 1,
        key_len: 8,
    };
    write_frame(&mut writer, &req.encode()).unwrap();

    let frame = read_frame(&mut reader).unwrap();
    assert_eq!(frame.msg_type, MessageType::Error);
    assert_eq!(
        decode_error(&frame.payload).unwrap().code,
        ReplicationErrorCode::ProtocolMismatch
    );

    signal.shutdown();
}

// ---------------------------------------------------------------------------
// Test 8: shard_id_validation_terminates_client  (C15)
// ---------------------------------------------------------------------------
/// Fabricate a TCP server that sends a malformed EntryBatch (shard_id=99
/// on a stream that was opened for shard 0). The client must return a
/// terminal DbError::Replication and exit that shard worker.
///
/// A raw TCP listener acts as a rogue server; the behavioral assertion verifies
/// that the malformed batch is rejected without applying any entries.
#[test]
fn shard_id_validation_terminates_client() {
    use armdb::replication::protocol::{
        EntryBatch, ShardInfo, SyncRequest, VAR_PROTOCOL_VERSION, WireEntry, read_frame,
        write_frame,
    };
    use std::io::BufReader;

    // Build a fake server that accepts one connection, does the handshake,
    // then sends an EntryBatch with shard_id=99 (mismatched).
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let server_addr: SocketAddr = listener.local_addr().unwrap();

    // Spawn a thread acting as the rogue server.
    let rogue_handle = std::thread::spawn(move || {
        let (stream, _) = listener.accept().expect("accept");
        let _ = stream.set_nodelay(true);

        let mut reader = BufReader::new(stream.try_clone().unwrap());
        let mut writer = stream;

        // Read SyncRequest.
        let frame = read_frame(&mut reader).expect("read SyncRequest");
        let req = SyncRequest::decode(&frame.payload).expect("decode SyncRequest");

        // Send a valid ShardInfo (shard_count must match follower's shard_count).
        // Config::test() uses shard_count=2.
        let info = ShardInfo {
            protocol_version: VAR_PROTOCOL_VERSION,
            shard_count: 2,
            max_file_size: 256 * 1024 * 1024,
        };
        write_frame(&mut writer, &info.encode()).expect("write ShardInfo");

        // Now send an EntryBatch with a WRONG shard_id (not req.shard_id).
        let wrong_shard_id: u8 = if req.shard_id == 0 { 1 } else { 0 };
        // Actually use a value that will DEFINITELY mismatch on shard 0's thread:
        // shard_id=99 is always wrong for a 2-shard engine.
        let bad_batch = EntryBatch {
            shard_id: 99, // mismatch — client must reject
            entries: vec![WireEntry {
                entry_len: 0,
                key_len: 8,
                gsn: 1,
                data: vec![],
            }],
        };
        let _ = write_frame(&mut writer, &bad_batch.encode());
        // Server is done; close connection.
        let _ = wrong_shard_id;
    });

    // Set up a follower that connects to the rogue server.
    let follower_dir = tempfile::tempdir().unwrap();
    let follower = Arc::new(VarTree::<[u8; 8]>::open(follower_dir.path(), test_cfg()).unwrap());
    let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));

    // The client will connect, receive the bad batch, detect shard_id mismatch
    // (C15), log a terminal error, and exit the shard worker.
    let signal = ShutdownSignal::new();
    let _client = follower
        .start_replication_client(server_addr, registry, signal.clone())
        .unwrap();

    // Wait for the rogue server to finish sending its bad frame.
    rogue_handle.join().expect("rogue server thread panicked");

    // Give the client a moment to process the bad frame and observe the error.
    std::thread::sleep(Duration::from_millis(500));

    // The follower must have received zero entries (bad batch was rejected).
    assert_eq!(
        follower.len(),
        0,
        "follower must not apply entries from a mismatched shard_id"
    );

    // Shut down the client handle after the shard worker has terminated.
    signal.shutdown();
}

// ---------------------------------------------------------------------------
// Test: restart_resumes_from_durable_recovery_without_cursor  (F3)
// ---------------------------------------------------------------------------
/// A restarted follower derives its resume point from durable recovery and
/// catches up without creating a separate variable-replication cursor file.
#[test]
fn restart_resumes_from_durable_recovery_without_cursor() {
    let one_shard = || Config::balanced().shard_count(1).hints(true).build();

    let leader_dir = tempfile::tempdir().unwrap();
    let follower_dir = tempfile::tempdir().unwrap();
    let addr = next_bind_addr();

    let leader = VarTree::<[u8; 8]>::open(leader_dir.path(), one_shard()).unwrap();
    let leader_signal = ShutdownSignal::new();
    let _server = leader
        .start_replication_server_with_options(
            addr,
            leader_signal.clone(),
            ReplicationServerOptions {
                heartbeat_interval_secs: 1,
                ..Default::default()
            },
        )
        .unwrap();

    for i in 0u64..100 {
        leader.put(&i.to_be_bytes(), b"a").unwrap();
    }
    {
        let follower =
            Arc::new(VarTree::<[u8; 8]>::open(follower_dir.path(), one_shard()).unwrap());
        let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));
        let sig = ShutdownSignal::new();
        let client = follower
            .start_replication_client(addr, registry, sig.clone())
            .unwrap();
        assert!(
            wait_until(Duration::from_secs(15), || follower.len() >= 100),
            "phase A catch-up failed (len={})",
            follower.len()
        );
        sig.shutdown();
        drop(client);
        follower.clean_shutdown().unwrap();
    }

    for i in 100u64..200 {
        leader.put(&i.to_be_bytes(), b"b").unwrap();
    }

    let follower = Arc::new(VarTree::<[u8; 8]>::open(follower_dir.path(), one_shard()).unwrap());
    let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));
    let sig = ShutdownSignal::new();
    let _client = follower
        .start_replication_client(addr, registry, sig.clone())
        .unwrap();

    assert!(
        wait_until(Duration::from_secs(15), || {
            (0u64..200).all(|i| follower.contains(&i.to_be_bytes()))
        }),
        "follower did not resume from durable recovery (len={})",
        follower.len()
    );
    assert_eq!(follower.len(), 200);
    assert!(!follower_dir.path().join("shard_000/repl.cursor").exists());

    sig.shutdown();
    leader_signal.shutdown();
}

// ---------------------------------------------------------------------------
// Test: silent_leader_read_timeout_triggers_reconnect  (F4)
// ---------------------------------------------------------------------------
/// A leader that dies with no TCP RST (power-off, cable pull) leaves the
/// follower blocked in `read_frame` forever. With the F4 read timeout the
/// follower instead observes silence for `2×heartbeat_interval`, errors out,
/// and reconnects. We prove it by pointing the follower at a mock leader that
/// completes the handshake and then stays silent (holding the connection open,
/// so the follower sees a timeout — not a clean EOF): the follower must open a
/// *second* connection, which only happens if it timed out rather than hanging.
#[test]
fn silent_leader_read_timeout_triggers_reconnect() {
    use armdb::replication::protocol::{ShardInfo, VAR_PROTOCOL_VERSION, read_frame, write_frame};
    use std::io::BufReader;
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let server_addr: SocketAddr = listener.local_addr().unwrap();
    listener.set_nonblocking(true).unwrap();

    let accepted = Arc::new(AtomicUsize::new(0));
    let accepted_mock = accepted.clone();
    let stop_mock = Arc::new(AtomicBool::new(false));
    let stop_mock2 = stop_mock.clone();

    let mock = std::thread::spawn(move || {
        // Hold every accepted connection open + silent so the follower sees a
        // read timeout rather than an EOF.
        let mut held = Vec::new();
        while !stop_mock2.load(Ordering::Relaxed) {
            match listener.accept() {
                Ok((stream, _)) => {
                    stream.set_nonblocking(false).ok();
                    let _ = stream.set_nodelay(true);
                    let mut reader = BufReader::new(stream.try_clone().unwrap());
                    let mut writer = stream.try_clone().unwrap();
                    // Complete the handshake: read SyncRequest, reply ShardInfo,
                    // then say nothing further.
                    if read_frame(&mut reader).is_ok() {
                        let info = ShardInfo {
                            protocol_version: VAR_PROTOCOL_VERSION,
                            shard_count: 1,
                            max_file_size: 256 * 1024 * 1024,
                        };
                        let _ = write_frame(&mut writer, &info.encode());
                        accepted_mock.fetch_add(1, Ordering::Relaxed);
                    }
                    held.push(stream);
                }
                Err(_) => std::thread::sleep(Duration::from_millis(20)),
            }
        }
    });

    let follower_dir = tempfile::tempdir().unwrap();
    let follower = Arc::new(
        VarTree::<[u8; 8]>::open(
            follower_dir.path(),
            Config::balanced().shard_count(1).hints(true).build(),
        )
        .unwrap(),
    );
    let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));
    let signal = ShutdownSignal::new();
    // heartbeat_interval_secs = 1 → 2s read timeout; fast reconnect backoff.
    let _client = follower
        .start_replication_client_with_options(
            server_addr,
            registry,
            signal.clone(),
            ReplicationClientOptions {
                reconnect_base_ms: 200,
                reconnect_max_ms: 200,
                heartbeat_interval_secs: 1,
            },
        )
        .unwrap();

    // A second accepted handshake proves the follower timed out on the silent
    // leader and reconnected instead of blocking forever.
    assert!(
        wait_until(Duration::from_secs(15), || {
            accepted.load(Ordering::Relaxed) >= 2
        }),
        "follower did not reconnect after read timeout (accepted={})",
        accepted.load(Ordering::Relaxed)
    );

    signal.shutdown();
    stop_mock.store(true, Ordering::Relaxed);
    let _ = mock.join();
}

// ---------------------------------------------------------------------------
// Test: dead_peer_releases_streaming_consumer  (F5)
// ---------------------------------------------------------------------------
/// A follower that dies silently (socket held open, no heartbeat ACKs)
/// used to keep the leader's streaming loop alive — it only exited on server
/// shutdown or a write error, which for a half-dead socket takes minutes of TCP
/// retransmission. Meanwhile the SPSC consumer stayed checked out, so a *real*
/// follower reconnecting was bounced with "shard already streaming". The
/// liveness worker sends a heartbeat after one idle interval and flags the peer
/// dead if the next interval elapses without an ACK; streaming then releases
/// the consumer.
///
/// We connect a mock follower that handshakes, drains catch-up, then goes
/// silent while holding the socket open. After the dead-peer window a second
/// mock follower reconnects and must be *accepted* (catch-up frames) rather than
/// rejected with an `Error`.
#[test]
fn dead_peer_releases_streaming_consumer() {
    use armdb::replication::protocol::{
        AckMessage, MessageType, SyncRequest, VAR_PROTOCOL_VERSION, read_frame, write_frame,
    };
    use std::io::BufReader;
    use std::net::TcpStream;
    use std::sync::atomic::{AtomicBool, Ordering};

    let leader_dir = tempfile::tempdir().unwrap();
    let leader = VarTree::<[u8; 8]>::open(
        leader_dir.path(),
        Config::balanced().shard_count(1).hints(true).build(),
    )
    .unwrap();
    // A few entries so a `from_gsn = 1` follower runs catch-up and receives a
    // CaughtUp frame (a positive "consumer acquired" signal).
    for i in 0u64..5 {
        leader.put(&i.to_be_bytes(), b"x").unwrap();
    }

    let addr = next_bind_addr();
    let leader_signal = ShutdownSignal::new();
    // heartbeat 1s → one probe + one unanswered interval → dead within ~2s.
    let _server = leader
        .start_replication_server_with_options(
            addr,
            leader_signal.clone(),
            ReplicationServerOptions {
                heartbeat_interval_secs: 1,
                ..Default::default()
            },
        )
        .unwrap();

    // Mock follower #1: handshake, drain to CaughtUp, then go silent (holding
    // the socket open) to simulate a peer that died without a TCP RST.
    let stop_dead = Arc::new(AtomicBool::new(false));
    let stop_dead2 = stop_dead.clone();
    let dead = std::thread::spawn(move || {
        let stream = TcpStream::connect(addr).unwrap();
        stream.set_nodelay(true).ok();
        let mut reader = BufReader::new(stream.try_clone().unwrap());
        let mut writer = stream.try_clone().unwrap();
        let req = SyncRequest {
            protocol_version: VAR_PROTOCOL_VERSION,
            shard_id: 0,
            from_gsn: 1,
            key_len: 8,
        };
        write_frame(&mut writer, &req.encode()).unwrap();
        // ShardInfo, then drain catch-up until CaughtUp.
        assert_eq!(
            read_frame(&mut reader).unwrap().msg_type,
            MessageType::ShardInfo
        );
        loop {
            if read_frame(&mut reader).unwrap().msg_type == MessageType::CaughtUp {
                break;
            }
        }
        // Now silent: no acks, no reads. Hold the socket open.
        while !stop_dead2.load(Ordering::Relaxed) {
            std::thread::sleep(Duration::from_millis(50));
        }
        drop(stream);
    });

    // Give the leader time to (a) reach streaming with the dead peer, then
    // (b) detect the unanswered heartbeat and release the consumer.
    std::thread::sleep(Duration::from_secs(5));

    // Mock follower #2: a genuine reconnect. It must be accepted — the second
    // frame after ShardInfo must NOT be an Error ("shard already streaming").
    let stream = TcpStream::connect(addr).unwrap();
    stream.set_nodelay(true).ok();
    stream
        .set_read_timeout(Some(Duration::from_secs(10)))
        .unwrap();
    let mut reader = BufReader::new(stream.try_clone().unwrap());
    let mut writer = stream.try_clone().unwrap();
    let req = SyncRequest {
        protocol_version: VAR_PROTOCOL_VERSION,
        shard_id: 0,
        from_gsn: 1,
        key_len: 8,
    };
    write_frame(&mut writer, &req.encode()).unwrap();
    assert_eq!(
        read_frame(&mut reader).unwrap().msg_type,
        MessageType::ShardInfo
    );
    let mut frame = read_frame(&mut reader).unwrap();
    assert_ne!(
        frame.msg_type,
        MessageType::Error,
        "reconnect rejected — streaming consumer was not released after the peer died (F5)"
    );
    loop {
        if frame.msg_type == MessageType::Heartbeat {
            write_frame(
                &mut writer,
                &AckMessage {
                    shard_id: 0,
                    last_gsn: 5,
                }
                .encode(),
            )
            .unwrap();
            break;
        }
        frame = read_frame(&mut reader).unwrap();
        assert_ne!(frame.msg_type, MessageType::Error);
    }

    stop_dead.store(true, Ordering::Relaxed);
    let _ = dead.join();
    leader_signal.shutdown();
}

// ---------------------------------------------------------------------------
// Test: replicates_var_value_larger_than_read_ahead_buffer  (#78)
// ---------------------------------------------------------------------------
/// A Var value larger than the log reader's 64 KiB read-ahead buffer
/// (`READ_AHEAD_SIZE`) and larger than the server's 64 KiB `BATCH_MAX_BYTES`
/// must replicate intact over BOTH the catch-up (log-reader) path and the live
/// streaming (SPSC ring) path.
///
/// Regression guard for #78: the pre–WP-K streaming reader read entries through
/// a fixed 64 KiB buffer whose `peek_bytes_from` returned `None` when an entry
/// did not fit, so it silently skipped (CRC-miss + step-forward) any entry over
/// 64 KiB — a large Var value never reached the follower. The WP-K k-way-merge
/// rewrite reads each entry in full via `read_exact_at(offset, entry_len)` and
/// the batch loop emits an oversized single-entry batch, so large values now
/// replicate on both paths. This test asserts that end to end, byte for byte.
#[test]
fn replicates_var_value_larger_than_read_ahead_buffer() {
    // Distinct, recognizable patterns so a zeroed/truncated tail is caught, not
    // just a missing key. Both are > 64 KiB (READ_AHEAD_SIZE == BATCH_MAX_BYTES).
    let catchup_value: Vec<u8> = (0..200 * 1024).map(|j| (j % 251) as u8).collect();
    let stream_value: Vec<u8> = (0..130 * 1024).map(|j| ((j * 7 + 3) % 251) as u8).collect();
    let catchup_key = 1u64.to_be_bytes();
    let stream_key = 2u64.to_be_bytes();

    let leader_dir = tempfile::tempdir().unwrap();
    let follower_dir = tempfile::tempdir().unwrap();
    let addr = next_bind_addr();

    let leader = VarTree::<[u8; 8]>::open(leader_dir.path(), test_cfg()).unwrap();
    // Written before the follower connects → exercised by the catch-up path.
    leader.put(&catchup_key, &catchup_value).unwrap();

    let leader_signal = ShutdownSignal::new();
    let _server = leader
        .start_replication_server(addr, leader_signal.clone())
        .unwrap();

    let follower = Arc::new(VarTree::<[u8; 8]>::open(follower_dir.path(), test_cfg()).unwrap());
    let registry = Arc::new(ReplicationRegistry::new(follower.as_replication_target()));
    let follower_signal = ShutdownSignal::new();
    let _client = follower
        .start_replication_client(addr, registry, follower_signal.clone())
        .unwrap();

    // Catch-up path: the > 64 KiB value must arrive and match byte for byte.
    assert!(
        wait_until(Duration::from_secs(15), || {
            follower.contains(&catchup_key)
        }),
        "follower did not catch up the large value within 15s"
    );
    let got = follower.get(&catchup_key).expect("catch-up value present");
    assert_eq!(
        got.as_ref(),
        catchup_value.as_slice(),
        "catch-up large value corrupted: got {} bytes, expected {}",
        got.len(),
        catchup_value.len()
    );

    // Live streaming path: write a second > 64 KiB value after catch-up.
    leader.put(&stream_key, &stream_value).unwrap();
    assert!(
        wait_until(Duration::from_secs(10), || {
            follower.contains(&stream_key)
        }),
        "follower did not stream the large value within 10s"
    );
    let got = follower.get(&stream_key).expect("streamed value present");
    assert_eq!(
        got.as_ref(),
        stream_value.as_slice(),
        "streamed large value corrupted: got {} bytes, expected {}",
        got.len(),
        stream_value.len()
    );

    follower_signal.shutdown();
    leader_signal.shutdown();
}