coordinode-lsm-tree 7.2.0

use super::*;

fn stats_with_avg(avg_entry_on_disk_bytes: u64) -> StorageStats {
    StorageStats {
        used_bytes: 0,
        capacity_bytes: None,
        available_bytes: None,
        compaction_possible: true,
        full_compaction_bytes: 0,
        tight_compaction_bytes: 0,
        item_count: 0,
        table_count: 0,
        avg_entry_on_disk_bytes,
        avg_key_bytes: None,
        avg_value_bytes: None,
        reclaimable_bytes_estimate: 0,
        status: StorageStatus::Healthy,
        blob_references: BlobReferenceStats::default(),
    }
}

/// `span_stats` over spans written as `(file, first, last)` byte strings.
fn spans_of(spans: &[(u64, &str, &str)]) -> BlobReferenceStats {
    let keys: Vec<(u64, crate::UserKey, crate::UserKey)> = spans
        .iter()
        .map(|&(id, first, last)| (id, first.into(), last.into()))
        .collect();
    span_stats(
        &mut keys
            .iter()
            .map(|(id, first, last)| (*id, first, last))
            .collect::<Vec<_>>(),
        crate::comparator::default_comparator().as_ref(),
    )
}

/// Spans are closed: two files whose spans meet at one key both cover it, so
/// a scan through that key reads both.
#[test]
fn spans_meeting_at_one_key_overlap() {
    assert_eq!(
        spans_of(&[(1, "a", "c"), (2, "c", "e")]),
        BlobReferenceStats { count: 2, depth: 2 }
    );
    assert_eq!(
        spans_of(&[(1, "a", "b"), (2, "c", "e")]),
        BlobReferenceStats { count: 2, depth: 1 }
    );
}

/// One blob file referenced by two tables is one file, however its spans
/// overlap: depth counts distinct files over a key, not references.
#[test]
fn one_file_referenced_twice_counts_once() {
    assert_eq!(
        spans_of(&[(7, "a", "c"), (7, "b", "d")]),
        BlobReferenceStats { count: 1, depth: 1 }
    );
    // Disjoint spans of one file, with another file between them.
    assert_eq!(
        spans_of(&[(7, "a", "b"), (7, "x", "y"), (8, "c", "d")]),
        BlobReferenceStats { count: 2, depth: 1 }
    );
}

/// Nested spans stack: a key inside all three covers all three files.
#[test]
fn nested_spans_stack() {
    assert_eq!(
        spans_of(&[(1, "a", "z"), (2, "m", "n"), (3, "m", "m")]),
        BlobReferenceStats { count: 3, depth: 3 }
    );
}

/// No spans: nothing referenced.
#[test]
fn no_spans_report_nothing() {
    assert_eq!(spans_of(&[]), BlobReferenceStats::default());
    assert_eq!(depths_of(&[]), Vec::new());
}

/// `span_depths` over spans written as `(file, first, last)` byte strings.
fn depths_of(spans: &[(u64, &str, &str)]) -> Vec<(u64, u64)> {
    let keys: Vec<(u64, crate::UserKey, crate::UserKey)> = spans
        .iter()
        .map(|&(id, first, last)| (id, first.into(), last.into()))
        .collect();
    span_depths(
        &mut keys
            .iter()
            .map(|(id, first, last)| (*id, first, last))
            .collect::<Vec<_>>(),
        crate::comparator::default_comparator().as_ref(),
    )
}

/// `pick_overlapping` over spans written as `(file, first, last)` byte
/// strings: the chosen files of `ranked`, ascending, and the cost spent.
fn picks_of(
    spans: &[(u64, &str, &str)],
    fixed: &[(u64, usize)],
    ranked: &[(u64, usize, u64)],
    budget: u64,
) -> (Vec<u64>, u64) {
    let keys: Vec<(u64, crate::UserKey, crate::UserKey)> = spans
        .iter()
        .map(|&(id, first, last)| (id, first.into(), last.into()))
        .collect();
    let (chosen, spent) = pick_overlapping(
        &mut keys
            .iter()
            .map(|(id, first, last)| (*id, first, last))
            .collect::<Vec<_>>(),
        fixed,
        ranked,
        budget,
        crate::comparator::default_comparator().as_ref(),
    );
    let mut ids: Vec<u64> = ranked
        .iter()
        .zip(chosen)
        .filter_map(|(&(id, _, _), chosen)| chosen.then_some(id))
        .collect();
    ids.sort_unstable();
    (ids, spent)
}

/// Every file of `spans` a candidate of one group costing one, in id order.
fn overlaps_of(spans: &[(u64, &str, &str)]) -> Vec<u64> {
    let mut ids: Vec<u64> = spans.iter().map(|&(id, _, _)| id).collect();
    ids.sort_unstable();
    ids.dedup();
    let ranked: Vec<(u64, usize, u64)> = ids.iter().map(|&id| (id, 0, 1)).collect();
    picks_of(spans, &[], &ranked, u64::MAX).0
}

/// Nested spans of any count give every file the full depth: the outermost
/// span's range covers more events than the largest power of two below the
/// event count, which the range-maximum table must still answer.
#[test]
fn nested_spans_give_every_file_the_full_depth() {
    for n in 1..=12u64 {
        let keys: Vec<(u64, String, String)> = (0..n)
            .map(|f| (f, format!("a{f:02}"), format!("z{:02}", n - f)))
            .collect();
        let spans: Vec<(u64, &str, &str)> = keys
            .iter()
            .map(|(id, first, last)| (*id, first.as_str(), last.as_str()))
            .collect();
        let want: Vec<(u64, u64)> = (0..n).map(|f| (f, n)).collect();
        assert_eq!(depths_of(&spans), want, "{n} nested spans");
    }
}

/// Files only merge within their group: overlapping files of different
/// groups are no partners, those of one group are.
#[test]
fn only_files_of_one_group_are_partners() {
    let spans = [(1, "a", "m"), (2, "c", "k"), (3, "d", "e")];
    // 1 and 3 share a group; 2 is alone in its own.
    let ranked = [(1, 0, 1), (2, 1, 1), (3, 0, 1)];
    assert_eq!(picks_of(&spans, &[], &ranked, u64::MAX), (vec![1, 3], 2));
}

/// A file relocated anyway is a partner at no cost: a candidate overlapping
/// it is chosen for its own cost alone.
#[test]
fn a_fixed_file_is_a_partner_at_no_cost() {
    let spans = [(5, "a", "z"), (1, "c", "d"), (2, "d", "f")];
    let ranked = [(1, 0, 10), (2, 0, 10)];
    assert_eq!(picks_of(&spans, &[(5, 0)], &ranked, 10), (vec![1], 10));
    // Of another group, it is no partner, and the two candidates only fit
    // together.
    assert_eq!(picks_of(&spans, &[(5, 1)], &ranked, 10), (vec![], 0));
    assert_eq!(picks_of(&spans, &[(5, 1)], &ranked, 20), (vec![1, 2], 20));
}

/// A candidate whose partners cannot come along leaves the budget to the
/// ones after it: a cheap file interleaved with two dear ones ranks first,
/// but a cheaper overlapping pair further down is what fits.
#[test]
fn a_partnerless_candidate_leaves_the_budget_to_a_pair() {
    let spans = [
        (1, "a", "m"),
        (2, "b", "n"),
        (3, "c", "o"),
        (4, "x", "y"),
        (5, "x", "z"),
    ];
    // Ranked 1, 4, 5, 2, 3; 1 overlaps only 2 and 3.
    let ranked = [(1, 0, 20), (4, 0, 20), (5, 0, 20), (2, 0, 40), (3, 0, 40)];
    assert_eq!(picks_of(&spans, &[], &ranked, 50), (vec![4, 5], 40));
}

/// Many disjoint pairs whose first members all rank ahead of their partners,
/// with a budget for one member of each: each pair is decided once, not the
/// whole plan once per rejected candidate.
#[test]
fn many_disjoint_pairs_are_decided_in_one_pass() {
    const PAIRS: u64 = 20_000;
    let keys: Vec<(u64, String, String)> = (0..PAIRS)
        .flat_map(|pair| {
            [
                (pair, format!("k{pair:06}a"), format!("k{pair:06}c")),
                (PAIRS + pair, format!("k{pair:06}b"), format!("k{pair:06}d")),
            ]
        })
        .collect();
    let spans: Vec<(u64, &str, &str)> = keys
        .iter()
        .map(|(id, first, last)| (*id, first.as_str(), last.as_str()))
        .collect();
    let ranked: Vec<(u64, usize, u64)> = (0..2 * PAIRS).map(|id| (id, 0, 1)).collect();
    // Room for every pair.
    let (chosen, spent) = picks_of(&spans, &[], &ranked, 2 * PAIRS);
    assert_eq!(chosen, (0..2 * PAIRS).collect::<Vec<_>>());
    assert_eq!(spent, 2 * PAIRS);
    // Room for half of them: the first pairs by rank, each whole.
    let (chosen, spent) = picks_of(&spans, &[], &ranked, PAIRS);
    let half: Vec<u64> = (0..PAIRS / 2).chain(PAIRS..PAIRS + PAIRS / 2).collect();
    assert_eq!(chosen, half);
    assert_eq!(spent, PAIRS);
}

/// A deep interleaving is chosen whole, and a file ending before the rest
/// open overlaps nothing and stays.
#[test]
fn a_deep_interleaving_is_chosen_whole() {
    let keys: Vec<(u64, String, String)> = (0..1_000u64)
        .map(|f| (f, format!("k{f:05}"), format!("z{f:05}")))
        .chain([(5_000, "a".to_owned(), "b".to_owned())])
        .collect();
    let spans: Vec<(u64, &str, &str)> = keys
        .iter()
        .map(|(id, first, last)| (*id, first.as_str(), last.as_str()))
        .collect();
    assert_eq!(overlaps_of(&spans), (0..1_000u64).collect::<Vec<_>>());
    assert_eq!(depths_of(&spans).first(), Some(&(0, 1_000)));
    assert_eq!(depths_of(&spans).last(), Some(&(5_000, 1)));
}

/// Only files whose spans meet another file's are chosen: a file off on its
/// own range is not, even beside overlapping ones.
#[test]
fn only_files_that_meet_another_are_chosen() {
    assert_eq!(
        overlaps_of(&[(1, "a", "m"), (2, "c", "k"), (9, "x", "z")]),
        vec![1, 2],
    );
    // Meeting at one key is overlapping; two spans of one file are not.
    assert_eq!(overlaps_of(&[(1, "a", "c"), (2, "c", "d")]), vec![1, 2]);
    assert_eq!(
        overlaps_of(&[(7, "a", "b"), (7, "c", "d")]),
        Vec::<u64>::new()
    );
    assert_eq!(
        overlaps_of(&[(1, "a", "b"), (9, "x", "z")]),
        Vec::<u64>::new()
    );
}

/// A span recorded with its bounds reversed covers the keys between them, not
/// nothing: it still overlaps the span beside it.
#[test]
fn a_reversed_span_is_taken_by_its_bounds() {
    assert_eq!(
        spans_of(&[(1, "m", "a"), (2, "c", "d")]),
        BlobReferenceStats { count: 2, depth: 2 }
    );
    assert_eq!(
        depths_of(&[(1, "m", "a"), (2, "c", "d")]),
        vec![(1, 2), (2, 2)]
    );
}

/// A span of a file inside another span of the same file adds nothing.
#[test]
fn a_span_nested_in_its_own_file_merges_away() {
    assert_eq!(
        spans_of(&[(7, "a", "z"), (7, "b", "c"), (8, "x", "y")]),
        BlobReferenceStats { count: 2, depth: 2 }
    );
}

/// A file's depth is the deepest point inside its own spans, not the tree's:
/// a file off to the side of a deep region keeps a depth of one, so a depth
/// limit never selects it.
#[test]
fn a_file_outside_the_deep_region_keeps_its_own_depth() {
    assert_eq!(
        depths_of(&[(1, "a", "m"), (2, "c", "k"), (3, "e", "g"), (4, "x", "z")]),
        vec![(1, 3), (2, 3), (3, 3), (4, 1)],
    );
}

/// A file whose spans sit in two regions takes the deeper one, and a span
/// that only touches a deep region at its edge key still shares its depth.
#[test]
fn a_file_takes_the_deepest_region_its_spans_reach() {
    assert_eq!(
        depths_of(&[(1, "a", "b"), (1, "p", "q"), (2, "p", "z"), (3, "q", "r")]),
        vec![(1, 3), (2, 3), (3, 3)],
    );
    // Only file 2 overlaps file 1's first span, and only at its last key.
    assert_eq!(
        depths_of(&[(1, "a", "c"), (2, "c", "d"), (3, "e", "f")]),
        vec![(1, 2), (2, 2), (3, 1)],
    );
}

#[test]
fn estimated_remaining_entries_divides_budget_by_average() {
    // budget / avg_entry_on_disk: 1000 bytes at 50 bytes/entry = 20 entries.
    let stats = stats_with_avg(50);
    assert_eq!(stats.estimated_remaining_entries(1000), 20);
    // Partial entries round down (integer division).
    assert_eq!(stats.estimated_remaining_entries(1049), 20);
    assert_eq!(stats.estimated_remaining_entries(0), 0);
}

#[test]
fn estimated_remaining_entries_is_zero_when_average_is_unknown() {
    // An empty tree has no average to extrapolate from, so any budget
    // yields 0 rather than dividing by zero.
    let stats = stats_with_avg(0);
    assert_eq!(stats.estimated_remaining_entries(1_000_000), 0);
}

#[test]
fn compute_on_empty_version_maps_compaction_flag_to_status() {
    use crate::TreeType;
    use crate::version::Version;

    // An empty version has no tables, so no file is stat-ed: the call is
    // pure and exercises only the status mapping and the zero-table path.
    let version = Version::new(0, TreeType::Standard);

    #[expect(
        clippy::unwrap_used,
        reason = "compute_storage_stats cannot fail on an empty in-memory version (no file to stat)"
    )]
    let busy = compute_storage_stats(&version, true, true).unwrap();
    assert_eq!(busy.status, StorageStatus::CompactionInProgress);
    assert_eq!(busy.used_bytes, 0);
    assert_eq!(busy.item_count, 0);
    assert_eq!(busy.table_count, 0);
    assert_eq!(busy.avg_key_bytes, None);
    assert_eq!(busy.estimated_remaining_entries(1_000_000), 0);

    #[expect(
        clippy::unwrap_used,
        reason = "compute_storage_stats cannot fail on an empty in-memory version (no file to stat)"
    )]
    let idle = compute_storage_stats(&version, false, true).unwrap();
    assert_eq!(idle.status, StorageStatus::Healthy);
}

#[test]
fn storage_statistics_is_object_safe_via_mock() -> crate::Result<()> {
    // A non-tree mock implements the trait, proving it is object-safe and usable
    // for planner / tiering tests without a real engine behind it.
    struct MockStats;
    impl StorageStatistics for MockStats {
        fn storage_stats(&self) -> crate::Result<StorageStats> {
            Ok(stats_with_avg(6))
        }
        fn level_segment_stats(&self) -> crate::Result<Vec<LevelStats>> {
            Ok(Vec::new())
        }
        fn compaction_debt(&self, _strategy: &dyn crate::compaction::CompactionStrategy) -> u64 {
            123
        }
        #[cfg(feature = "metrics")]
        fn cache_stats(&self) -> crate::CacheStats {
            crate::CacheStats {
                hits: 9,
                misses: 1,
                hit_rate: 0.9,
                size_bytes: 10,
                capacity_bytes: 100,
            }
        }
    }

    let mock = MockStats;
    let stats: &dyn StorageStatistics = &mock;
    assert_eq!(stats.storage_stats()?.avg_entry_on_disk_bytes, 6);
    assert!(stats.level_segment_stats()?.is_empty());
    let strategy = crate::compaction::leveled::Strategy::default();
    assert_eq!(stats.compaction_debt(&strategy), 123);
    #[cfg(feature = "metrics")]
    assert_eq!(stats.cache_stats().hits, 9);
    Ok(())
}

#[test]
fn storage_statistics_blanket_impl_over_real_tree() -> crate::Result<()> {
    // Drive the blanket `impl<T: AbstractTree> StorageStatistics for T` over a
    // real tree (the mock above bypasses it), and a non-leveled strategy through
    // the trait-default `pending_compaction_bytes` of 0.
    use crate::{AbstractTree, Config, SequenceNumberCounter, StorageStatistics};

    let dir = tempfile::tempdir()?;
    let tree = Config::new(
        dir.path(),
        SequenceNumberCounter::default(),
        SequenceNumberCounter::default(),
    )
    .open()?;
    for i in 0..50u32 {
        tree.insert(format!("k{i:04}"), "v", 0);
    }
    tree.flush_active_memtable(0)?;

    let s: &dyn StorageStatistics = &tree;
    assert_eq!(s.storage_stats()?.item_count, 50);
    let level_items: u64 = s.level_segment_stats()?.iter().map(|l| l.item_count).sum();
    assert_eq!(level_items, 50);
    // FIFO has no size-target debt notion, so the trait default 0 applies.
    let fifo = crate::compaction::fifo::Strategy::new(u64::MAX, None);
    assert_eq!(s.compaction_debt(&fifo), 0);
    #[cfg(feature = "metrics")]
    assert!(s.cache_stats().capacity_bytes > 0);
    Ok(())
}

/// The per-level accounting shares the tree-level basis: a restricted
/// table's live `.restrict-bound` sidecar counts in `LevelStats` /
/// `SegmentStats` too, or summing the levels stops reconciling with the
/// documented SST portion of `StorageStats::used_bytes` and tiering
/// consumers understate the footprint.
#[test]
#[expect(clippy::expect_used, reason = "test code")]
fn level_stats_count_a_restricted_tables_sidecar() -> crate::Result<()> {
    use crate::blob_tree::FragmentationMap;
    use crate::table::{Table, Writer};
    use crate::version::{BlobFileList, Level, Run, Version};
    use crate::{InternalValue, TreeType, ValueType};
    use std::sync::Arc;

    let dir = tempfile::tempdir()?;
    let fs: Arc<dyn crate::fs::Fs> = Arc::new(crate::fs::StdFs);
    let sst = dir.path().join("0");

    // A multi-block SST so a restriction bound lands on a real block boundary.
    let mut w = Writer::new(sst.clone(), 0, 0, Arc::clone(&fs))?.use_data_block_size(128);
    for i in 0..256u32 {
        w.write(InternalValue::from_components(
            format!("k{i:05}").into_bytes(),
            format!("v{i}").into_bytes(),
            u64::from(i) + 1,
            ValueType::Value,
        ))?;
    }
    let (_, checksum) = w.finish()?.expect("table written");

    let table = {
        let mut params = crate::table::RecoverParams::new(
            sst.clone(),
            checksum,
            0,
            Arc::clone(&fs),
            crate::comparator::default_comparator(),
            Arc::new(crate::Cache::with_capacity_bytes(1_000_000)),
        );
        params.descriptor_table = Some(Arc::new(crate::descriptor_table::DescriptorTable::new(10)));
        Table::recover(params)?
    };
    let restricted = table.reopen_restricted(b"k00100".to_vec().into())?;
    crate::restrict_bound::write(&*fs, &sst, None, 0, b"k00100", crate::fs::SyncMode::Normal)?;

    let run = Arc::new(Run::new(vec![restricted]).expect("non-empty run"));
    let version = Version::from_levels(
        1,
        TreeType::Standard,
        vec![Level::from_runs(vec![run])],
        BlobFileList::default(),
        FragmentationMap::default(),
    );

    let levels = compute_level_segment_stats(&version)?;
    let sum: u64 = levels.iter().map(|l| l.used_bytes).sum();
    let expected: u64 = version
        .iter_tables()
        .map(table_on_disk_bytes)
        .sum::<crate::Result<u64>>()?;
    let sst_alone = fs.metadata(&sst)?.len;
    assert!(
        expected > sst_alone,
        "the fixture's live sidecar contributes bytes ({expected} vs {sst_alone})",
    );
    assert_eq!(
        sum, expected,
        "per-level and tree-level accounting share one sidecar-aware basis",
    );
    Ok(())
}

/// A tight-space compaction reclaims the consumed prefix of its input IN PLACE:
/// the blocks are gone from the device while `len` still reports the original
/// size. Charging the logical length keeps those freed bytes on the quota
/// forever: under `storage_limit_bytes` the headroom never recovers and the
/// tree stays read-only even though the compaction succeeded.
#[test]
#[expect(clippy::expect_used, reason = "test code")]
fn a_punched_table_is_charged_its_allocated_bytes() -> crate::Result<()> {
    use crate::fs::{Fs, MemFs};
    use crate::table::{Table, Writer};
    use crate::{InternalValue, ValueType};
    use std::sync::Arc;

    let memfs = MemFs::new();
    let fs: Arc<dyn Fs> = Arc::new(memfs.clone());
    let dir = std::path::absolute("/punched_accounting")?;
    memfs.create_dir_all(&dir)?;
    let sst = dir.join("0");

    let mut w = Writer::new(sst.clone(), 0, 0, Arc::clone(&fs))?.use_data_block_size(128);
    for i in 0..256u32 {
        w.write(InternalValue::from_components(
            format!("k{i:05}").into_bytes(),
            format!("v{i}").into_bytes(),
            u64::from(i) + 1,
            ValueType::Value,
        ))?;
    }
    let (_, checksum) = w.finish()?.expect("table written");

    let table = Table::recover(crate::table::RecoverParams::new(
        sst.clone(),
        checksum,
        0,
        Arc::clone(&fs),
        crate::comparator::default_comparator(),
        Arc::new(crate::Cache::with_capacity_bytes(1_000_000)),
    ))?;
    let bound: crate::UserKey = b"k00100".to_vec().into();
    let punch_at = table.punch_offset_for(&bound)?;
    assert!(punch_at > 0, "the bound must fall past the first block");
    let logical = fs.metadata(&sst)?.len;

    let restricted = table.reopen_restricted(bound)?;
    // Reclaim the consumed prefix, exactly as the tight-space pass does.
    memfs.punch_hole(&sst, 0, punch_at)?;

    let charged = table_on_disk_bytes(&restricted)?;
    assert!(
        charged < logical,
        "the reclaimed prefix must leave the quota: charged {charged}, logical {logical}",
    );
    assert_eq!(
        charged,
        logical - punch_at,
        "what is charged is what is still allocated",
    );
    Ok(())
}

/// A restricted table's `metadata.item_count` still describes the WHOLE
/// original SST, prefix included. While a tight-space slice is in flight the
/// version holds both that restricted input and the output that now owns the
/// prefix, so summing the raw metadata counts those entries twice: the reported
/// item count inflates, the average entry size shrinks, and
/// `estimated_remaining_entries` overstates capacity.
#[test]
#[expect(clippy::expect_used, reason = "test code")]
fn a_restricted_table_contributes_only_its_live_entries() -> crate::Result<()> {
    use crate::blob_tree::FragmentationMap;
    use crate::table::{Table, Writer};
    use crate::version::{BlobFileList, Level, Run, Version};
    use crate::{InternalValue, TreeType, ValueType};
    use std::sync::Arc;

    let dir = tempfile::tempdir()?;
    let fs: Arc<dyn crate::fs::Fs> = Arc::new(crate::fs::StdFs);
    let sst = dir.path().join("0");

    let mut w = Writer::new(sst.clone(), 0, 0, Arc::clone(&fs))?.use_data_block_size(128);
    for i in 0..256u32 {
        w.write(InternalValue::from_components(
            format!("k{i:05}").into_bytes(),
            format!("v{i}").into_bytes(),
            u64::from(i) + 1,
            ValueType::Value,
        ))?;
    }
    let (_, checksum) = w.finish()?.expect("table written");

    let table = Table::recover(crate::table::RecoverParams::new(
        sst,
        checksum,
        0,
        Arc::clone(&fs),
        crate::comparator::default_comparator(),
        Arc::new(crate::Cache::with_capacity_bytes(1_000_000)),
    ))?;
    let whole = table.metadata.item_count;
    assert_eq!(whole, 256, "the fixture holds every entry");

    // Restrict away roughly the first half, as a completed slice would.
    let restricted = table.reopen_restricted(b"k00128".to_vec().into())?;
    let run = Arc::new(Run::new(vec![restricted]).expect("non-empty run"));
    let version = Version::from_levels(
        1,
        TreeType::Standard,
        vec![Level::from_runs(vec![run])],
        BlobFileList::default(),
        FragmentationMap::default(),
    );

    let stats = compute_storage_stats(&version, false, true)?;
    assert!(
        stats.item_count < whole,
        "the consumed prefix must not be counted: reported {} of {whole}",
        stats.item_count,
    );
    assert!(
        stats.item_count > whole / 4,
        "the live suffix must still be counted: reported {} of {whole}",
        stats.item_count,
    );
    Ok(())
}