armdb 0.11.2

#![cfg(all(
    feature = "armour",
    feature = "replication",
    feature = "typed-tree",
    feature = "rapira-codec"
))]

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};

use armdb::armour::Db;
use armdb::fixed_replication::{FixedReplicationClient, FixedReplicationTarget};
use armdb::replication::ReplicationRegistry;
use armdb::{
    CollectionMeta, Config, DbError, FixedConfig, FixedTree, NoHook, RapiraCodec, ShutdownSignal,
    TtlPolicy, VarTree, WriteHook,
};
use armour_core::GetType;
use rapira::Rapira;
use tempfile::tempdir;

#[derive(Clone, Debug, PartialEq, Rapira, GetType)]
struct Row {
    expires_at_ms: u64,
}

impl CollectionMeta for Row {
    type SelfId = [u8; 8];
    const NAME: &'static str = "ttl_repl_rows";
    const TTL: Option<TtlPolicy<Self>> = Some(TtlPolicy {
        revision: 1,
        deadline: |row, _| (row.expires_at_ms > 0).then_some(row.expires_at_ms),
    });
}

struct VarExpiryHook;

impl WriteHook<[u8; 8]> for VarExpiryHook {
    const EXPIRES: bool = true;

    fn deadline(&self, _key: &[u8; 8], value: &[u8]) -> Option<u64> {
        let arr: [u8; 8] = value.get(..8)?.try_into().ok()?;
        let ms = u64::from_be_bytes(arr);
        (ms > 0).then_some(ms)
    }

    fn on_write(&self, _key: &[u8; 8], _old: Option<&[u8]>, _new: Option<&[u8]>) {}
}

struct FixedExpiryHook;

impl WriteHook<[u8; 8]> for FixedExpiryHook {
    const EXPIRES: bool = true;

    fn deadline(&self, _key: &[u8; 8], value: &[u8]) -> Option<u64> {
        let arr: [u8; 8] = value.get(..8)?.try_into().ok()?;
        let ms = u64::from_be_bytes(arr);
        (ms > 0).then_some(ms)
    }

    fn on_write(&self, _key: &[u8; 8], _old: Option<&[u8]>, _new: Option<&[u8]>) {}
}

fn sweep_config() -> Config {
    Config::balanced().shard_count(1).hints(true).build()
}

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

fn deadline_bytes(ms: u64) -> Vec<u8> {
    ms.to_be_bytes().to_vec()
}

#[test]
fn follower_has_sweep_disabled() {
    let dir = tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let tree = db
        .open_typed_tree::<Row, RapiraCodec, NoHook>(sweep_config(), NoHook, &[])
        .unwrap();
    tree.put(&key(1), Row { expires_at_ms: 1 }).unwrap();

    tree.mark_replication_follower();
    let out = tree.sweep_expired(10, 4096).unwrap();
    assert_eq!(out.deleted, 0);
    assert!(tree.get(&key(1)).is_some());
}

#[test]
fn registry_registration_disables_sweep() {
    let dir = tempdir().unwrap();
    let follower = Arc::new(
        VarTree::<[u8; 8], VarExpiryHook>::open_hooked(dir.path(), sweep_config(), VarExpiryHook)
            .unwrap(),
    );
    follower.put(&key(1), &deadline_bytes(1)).unwrap();
    follower.build_ttl_index().unwrap();

    let _registry = ReplicationRegistry::new(follower.as_replication_target());

    let out = follower.sweep_expired(10, 4096).unwrap();
    assert_eq!(out.deleted, 0);
    assert!(follower.get(&key(1)).is_some());
}

#[test]
fn fixed_client_start_disables_sweep() {
    let dir = tempdir().unwrap();
    let follower = Arc::new(
        FixedTree::<[u8; 8], 8, FixedExpiryHook>::open_with_hook(
            dir.path(),
            FixedConfig::test(),
            FixedExpiryHook,
        )
        .unwrap(),
    );
    follower.put(&key(1), &key(1)).unwrap();
    follower.build_ttl_index().unwrap();

    let target: Arc<dyn FixedReplicationTarget> = follower.clone();
    let stop = ShutdownSignal::new();
    let _client =
        FixedReplicationClient::start("127.0.0.1:1".parse().unwrap(), target, stop.clone())
            .unwrap();

    let out = follower.sweep_expired(10, 4096).unwrap();
    assert_eq!(out.deleted, 0);
    assert!(follower.get(&key(1)).is_some());

    stop.shutdown();
}

#[test]
fn set_ttl_sweep_false_waits_for_in_flight_pass() {
    let dir = tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let tree = db
        .open_typed_tree::<Row, RapiraCodec, NoHook>(sweep_config(), NoHook, &[])
        .unwrap();
    tree.put(&key(1), Row { expires_at_ms: 1 }).unwrap();

    let started = Arc::new(AtomicBool::new(false));
    let started2 = Arc::clone(&started);

    std::thread::scope(|s| {
        s.spawn(|| {
            let _pass = tree.test_hold_sweep_pass().expect("sweep enabled");
            started2.store(true, Ordering::Release);
            std::thread::sleep(Duration::from_millis(200));
        });

        while !started.load(Ordering::Acquire) {
            std::thread::yield_now();
        }

        let t0 = Instant::now();
        tree.set_ttl_sweep(false).unwrap();
        assert!(
            t0.elapsed() >= Duration::from_millis(150),
            "set_ttl_sweep(false) returned before in-flight pass ended"
        );
    });
}

#[test]
fn enable_after_follower_without_rebuild_is_client_error() {
    let dir = tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let tree = db
        .open_typed_tree::<Row, RapiraCodec, NoHook>(sweep_config(), NoHook, &[])
        .unwrap();
    tree.mark_replication_follower();

    let err = db.set_ttl_sweep(Row::NAME, true).unwrap_err();
    assert!(matches!(err, DbError::Client(_)));
    let err = tree.set_ttl_sweep(true).unwrap_err();
    assert!(matches!(err, DbError::Client(_)));
}

#[test]
fn rebuild_then_enable_allows_sweep() {
    let dir = tempdir().unwrap();
    let db = Db::open_test(dir.path()).unwrap();
    let tree = db
        .open_typed_tree::<Row, RapiraCodec, NoHook>(sweep_config(), NoHook, &[])
        .unwrap();
    tree.put(&key(1), Row { expires_at_ms: 1 }).unwrap();
    tree.mark_replication_follower();

    db.rebuild_ttl_index(Row::NAME).unwrap();
    db.set_ttl_sweep(Row::NAME, true).unwrap();

    let out = tree.sweep_expired(10, 4096).unwrap();
    assert_eq!(out.deleted, 1);
    assert!(tree.get(&key(1)).is_none());
}