liminal-rs 0.16.2

use std::error::Error;
use std::fs;
use std::path::Path;
use std::sync::Arc;
use std::time::{Duration, Instant};

use haematite::{Database, DatabaseConfig, EventStore, ReclaimStart, RuntimeOptions, SyncCounts};
use tempfile::TempDir;

use super::bridge::block_on;
use super::{
    CountingStore, DurableStore, HaematiteStore, RoutedBatchItem, advance_head,
    head_checkpoint_key, recover_partition_head,
};

const STREAM: &str = "real-store-proof";
const TAIL_ROWS: u64 = 5;
const LIVE_ROWS: u64 = 512;
type TestResult<T = ()> = Result<T, Box<dyn Error>>;

struct Fixture {
    store: HaematiteStore,
    events: Arc<EventStore>,
    directory: TempDir,
}

fn runtime(reclaim_start: ReclaimStart) -> RuntimeOptions {
    RuntimeOptions {
        long_wait_interval: Duration::from_secs(5),
        reclaim_start,
    }
}

impl Fixture {
    fn new() -> TestResult<Self> {
        Self::with_reclaim_start(ReclaimStart::AtOpen)
    }

    fn with_reclaim_start(reclaim_start: ReclaimStart) -> TestResult<Self> {
        let directory = tempfile::tempdir()?;
        let database = Database::create(
            DatabaseConfig {
                data_dir: directory.path().to_path_buf(),
                shard_count: 1,
                distributed: None,
                executor_threads: None,
                node_cache_budget: Some(haematite::NodeCacheBudget::Unlimited),
                retention: haematite::Retention::reclaim(0),
            },
            runtime(reclaim_start),
        )?;
        let events = Arc::new(EventStore::new(database));
        Ok(Self {
            store: HaematiteStore::new(Arc::clone(&events)),
            events,
            directory,
        })
    }

    fn reopen(self) -> TestResult<Self> {
        let Self {
            store,
            events,
            directory,
        } = self;
        drop(store);
        drop(events);
        // Hold reclamation so the node census covers only the recovery call.
        let database = Database::open(directory.path(), runtime(ReclaimStart::Held))?;
        let events = Arc::new(EventStore::new(database));
        Ok(Self {
            store: HaematiteStore::new(Arc::clone(&events)),
            events,
            directory,
        })
    }

    fn close(self) -> TestResult {
        let Self {
            store,
            events,
            directory,
        } = self;
        drop(store);
        drop(events);
        directory.close()?;
        Ok(())
    }
}

fn payload(sequence: u64) -> Vec<u8> {
    sequence.to_be_bytes().to_vec()
}

fn seed(store: &dyn DurableStore, rows: u64) -> TestResult {
    let written = block_on(store.append_all(STREAM, (0..rows).map(payload).collect(), 0))??;
    assert_eq!(written, rows, "seed rows={rows} written={written}");
    block_on(store.flush())??;
    Ok(())
}

#[derive(Debug)]
struct Restart {
    rows: u64,
    recovered: u64,
    recorded: Option<u64>,
    node_gets: u64,
    milliseconds: f64,
}

fn restart(rows: u64) -> TestResult<Restart> {
    let fixture = Fixture::new()?;
    seed(&fixture.store, rows)?;
    block_on(advance_head(&fixture.store, STREAM, rows - TAIL_ROWS))??;
    block_on(fixture.store.flush())??;
    let fixture = fixture.reopen()?;
    let before = fixture.events.database().store_counts().node_gets;
    let start = Instant::now();
    let recovered = block_on(recover_partition_head(&fixture.store, STREAM))??;
    let milliseconds = start.elapsed().as_secs_f64() * 1_000.0;
    let node_gets = fixture.events.database().store_counts().node_gets - before;
    let recorded = block_on(fixture.store.read_value(&head_checkpoint_key(STREAM)))??;
    fixture.close()?;
    Ok(Restart {
        rows,
        recovered,
        recorded,
        node_gets,
        milliseconds,
    })
}

#[test]
fn a_restart_over_the_real_store_reads_the_same_at_ten_thousand_and_a_hundred_thousand_rows()
-> TestResult {
    let small = restart(10_000)?;
    let large = restart(100_000)?;
    let report = format!(
        "NET-015 restart: tail_rows={TAIL_ROWS} small={small:?} ({:.1} ms) large={large:?} ({:.1} ms)",
        small.milliseconds, large.milliseconds
    );
    assert_eq!(small.recovered, small.rows, "{report}");
    assert_eq!(large.recovered, large.rows, "{report}");
    assert_eq!(small.recorded, Some(small.rows - TAIL_ROWS), "{report}");
    assert_eq!(large.recorded, Some(large.rows - TAIL_ROWS), "{report}");
    // The witness is the store's own count: a restart reads the same nodes at ten
    // times the rows. Its milliseconds are reported, never asserted.
    assert_eq!(small.node_gets, large.node_gets, "{report}");
    println!("{report}");
    Ok(())
}

fn files_and_bytes(root: &Path) -> TestResult<(u64, u64)> {
    let mut files = 0;
    let mut bytes = 0;
    let mut pending = vec![root.to_path_buf()];
    while let Some(directory) = pending.pop() {
        for entry in fs::read_dir(directory)? {
            let entry = entry?;
            let kind = entry.file_type()?;
            if kind.is_dir() {
                pending.push(entry.path());
            } else if kind.is_file() {
                files += 1;
                bytes += entry.metadata()?.len();
            }
        }
    }
    Ok((files, bytes))
}

#[test]
fn a_prune_on_the_real_store_shrinks_the_bytes_on_disk_and_keeps_the_live_rows() -> TestResult {
    prune(ReclaimStart::AtOpen, "prune")
}

#[test]
fn a_held_store_released_the_way_the_server_releases_it_reclaims_a_prune() -> TestResult {
    prune(ReclaimStart::Held, "held-prune")
}

fn prune(reclaim_start: ReclaimStart, label: &str) -> TestResult {
    const ROWS: u64 = 50_000;
    let fixture = Fixture::with_reclaim_start(reclaim_start)?;
    seed(&fixture.store, ROWS)?;
    let before = files_and_bytes(fixture.directory.path())?;
    if reclaim_start == ReclaimStart::Held {
        fixture.events.database().release_reclaim()?;
    }
    let cut = ROWS - LIVE_ROWS;
    let deleted = block_on(fixture.store.truncate_before(STREAM, cut))??;
    block_on(fixture.store.flush())??;
    fixture
        .events
        .database()
        .quiesce_reclaim()
        .map_err(|error| {
            format!(
                "NET-015 {label}: rows={ROWS} cut={cut} deleted_rows={deleted} \
             live_rows={LIVE_ROWS} before_files={} before_bytes={} quiesce_error={error}",
                before.0, before.1
            )
        })?;
    let after = files_and_bytes(fixture.directory.path())?;
    let retained = block_on(fixture.store.read_from(
        STREAM,
        cut,
        usize::try_from(LIVE_ROWS)? + 1,
    ))??;
    let report = format!(
        "NET-015 {label}: rows={ROWS} cut={cut} deleted_rows={deleted} live_rows={LIVE_ROWS} \
         before_files={} before_bytes={} after_files={} after_bytes={} retained_rows={}",
        before.0,
        before.1,
        after.0,
        after.1,
        retained.len()
    );
    assert_eq!(deleted, cut, "{report}");
    assert_eq!(u64::try_from(retained.len())?, LIVE_ROWS, "{report}");
    for (entry, sequence) in retained.iter().zip(cut..ROWS) {
        assert_eq!(
            entry.sequence, sequence,
            "{report} expected_sequence={sequence}"
        );
        assert_eq!(
            entry.payload,
            payload(sequence),
            "{report} expected_sequence={sequence}"
        );
    }
    assert!(after.1 < before.1, "{report}");
    fixture.close()?;
    println!("{report}");
    Ok(())
}

#[derive(Debug)]
struct Batch {
    items: u64,
    syncs: SyncCounts,
    device_flushes: u64,
    full_syncs: u64,
    apply_routed_batch_calls: u64,
    apply_routed_batch_items: u64,
    append_calls: u64,
    append_all_calls: u64,
    flushes: u64,
}

fn batch(items: u64) -> TestResult<Batch> {
    let fixture = Fixture::new()?;
    let store = CountingStore::new(fixture.store.clone());
    let writes = (0..items)
        .map(|sequence| RoutedBatchItem {
            key: sequence.to_be_bytes().to_vec(),
            expected: None,
            value: payload(sequence),
            ttl: None,
        })
        .collect();
    let before = store.snapshot();
    let before_syncs = fixture.events.database().sync_counts();
    let before_device_flushes = fixture.events.database().device_flushes();
    let before_full_syncs = fixture.events.database().full_syncs();
    block_on(store.apply_routed_batch(STREAM, writes))??;
    let syncs = fixture.events.database().sync_counts().since(before_syncs);
    let counts = store.snapshot().since(&before);
    let result = Batch {
        items,
        syncs,
        device_flushes: fixture.events.database().device_flushes() - before_device_flushes,
        full_syncs: fixture.events.database().full_syncs() - before_full_syncs,
        apply_routed_batch_calls: counts.apply_routed_batch_calls,
        apply_routed_batch_items: counts.apply_routed_batch_items,
        append_calls: counts.append_calls,
        append_all_calls: counts.append_all_calls,
        flushes: counts.flushes,
    };
    for sequence in 0..items {
        let value = block_on(store.get_routed(STREAM, &sequence.to_be_bytes()))??;
        assert_eq!(value, Some(payload(sequence)), "{result:?} key={sequence}");
    }
    drop(store);
    fixture.close()?;
    Ok(result)
}

#[test]
fn a_routed_batch_is_one_commit_with_one_sync_census() -> TestResult {
    let small = batch(1)?;
    let large = batch(1_000)?;
    let report = format!("NET-015 batch: small={small:?} large={large:?}");
    for measured in [&small, &large] {
        assert_eq!(measured.syncs.barrier_device_flushes, 1, "{report}");
        assert_eq!(measured.apply_routed_batch_calls, 1, "{report}");
        assert_eq!(
            measured.apply_routed_batch_items, measured.items,
            "{report}"
        );
        assert_eq!(measured.append_calls, 0, "{report}");
        assert_eq!(measured.append_all_calls, 0, "{report}");
        assert_eq!(measured.flushes, 0, "{report}");
    }
    assert_eq!(small.syncs, large.syncs, "{report}");
    assert_eq!(small.device_flushes, large.device_flushes, "{report}");
    assert_eq!(small.full_syncs, large.full_syncs, "{report}");
    println!("{report}");
    Ok(())
}