armdb 0.9.0

sharded bitcask key-value storage optimized for NVMe
Documentation
use std::collections::{HashMap, HashSet};

use crate::armour::db::Db;
use crate::error::{DbError, DbResult};

use super::registry::{BuiltRegistry, scan_target_actual};
use super::validate::hex;

/// Options for [`repair`].
#[derive(Debug, Clone, Default)]
pub struct RepairOptions {
    /// Delete orphan index rows on targets marked [`super::SchemaRegistry::complete_index`].
    pub delete_orphans: bool,
}

/// Per-index repair counters.
#[derive(Debug, Clone, Default)]
pub struct IndexRepairStats {
    pub name: &'static str,
    pub puts: u64,
    pub rewrites: u64,
    pub deletes: u64,
    pub declined_orphans: u64,
}

/// A source collection whose scan skipped rows it could not read or decode.
#[derive(Debug, Clone)]
pub struct UndecodableSource {
    pub name: &'static str,
    pub rows: u64,
    /// Targets whose orphan deletion was suppressed because of it.
    pub protected_targets: Vec<&'static str>,
}

/// Aggregate result of [`repair`].
#[derive(Debug, Clone, Default)]
pub struct RepairReport {
    pub per_index: Vec<IndexRepairStats>,
    /// Sources that could not be scanned in full. Their index rows are missing
    /// from the expected set, so orphan deletion is held back for every target
    /// they feed — see [`repair`].
    pub undecodable_sources: Vec<UndecodableSource>,
}

/// Apply the same scan as `validate` and fix index drift: put missing rows,
/// rewrite mismatches, optionally delete orphans on complete targets.
///
/// Aborts with [`DbError::IndexKeyCollision`] if two source owners project the
/// same `(target, key)` — collisions are unrepairable and must not leave an
/// arbitrary scan-order winner on disk.
///
/// # Undecodable source rows
///
/// A source scan silently skips rows whose value cannot be read or decoded (Var
/// and LSM families; see `Finding::UndecodableRows`). Those rows contribute
/// nothing to the expected set, so their live index entries would look like
/// orphans. Deleting them would destroy the only remaining record of a row whose
/// value is already unreadable, so **orphan deletion is suppressed** for every
/// target fed by such a source, regardless of
/// [`RepairOptions::delete_orphans`] and `complete_index`. Missing rows are
/// still put and mismatches still rewritten — only the destructive half is held
/// back. The affected sources are listed in
/// [`RepairReport::undecodable_sources`].
pub fn repair(db: &Db, reg: &BuiltRegistry, opts: RepairOptions) -> DbResult<RepairReport> {
    let _ = db;
    use super::expected::{self, BuildIssue};

    // Pass 1: the expected set — shared with `validate`, so the two cannot
    // disagree about what an index should hold. Any issue aborts here, before
    // anything below has written a byte.
    let (expected, skipped) = expected::build(reg, &mut |issue| match issue {
        BuildIssue::Collision { index, key, owners } => Err(DbError::IndexKeyCollision {
            index,
            key: hex(&key),
            owners: owners.to_vec(),
        }),
        BuildIssue::Projection { error, .. } => Err(error),
    })?;

    // Targets whose expected set is known to be short, so an "orphan" there may
    // be a live row of a source record that can no longer be read.
    let protected: HashSet<&'static str> = skipped
        .iter()
        .flat_map(|s| s.targets.iter().copied())
        .collect();
    let undecodable_sources: Vec<UndecodableSource> = skipped
        .into_iter()
        .map(|s| UndecodableSource {
            name: s.name,
            rows: s.rows,
            protected_targets: s.targets,
        })
        .collect();

    let mut actual_by_target: HashMap<&'static str, HashMap<Vec<u8>, Vec<u8>>> = HashMap::new();
    let mut stats_by_name: HashMap<&'static str, IndexRepairStats> = HashMap::new();
    let all_targets = expected::declared_targets(reg);

    // Pass 2: put what is missing, rewrite what differs. Keyed by target, not by
    // declaration — the expected set already unions every source, and `actual`
    // is a pre-mutation snapshot, so walking a shared target once per
    // declaration would re-put each row and double-count the report.
    for target_name in &all_targets {
        let Some(rows) = expected.rows(target_name) else {
            continue;
        };
        let Some(t) = expected::target_decl(reg, target_name) else {
            continue;
        };
        let actual = scan_target_actual(reg, &mut actual_by_target, t.name, &t.value_to_bytes);
        for (key, (value, _owner)) in rows {
            let action = match actual.get(key) {
                None => Some(true),
                Some(actual_bytes) => match (t.value_eq)(value, actual_bytes) {
                    Ok(false) => Some(false),
                    Ok(true) => None,
                    Err(e) => return Err(e),
                },
            };
            if let Some(missing) = action {
                t.index.put_bytes(key, value)?;
                let stats = stats_by_name
                    .entry(t.name)
                    .or_insert_with(|| IndexRepairStats {
                        name: t.name,
                        ..Default::default()
                    });
                if missing {
                    stats.puts += 1;
                } else {
                    stats.rewrites += 1;
                }
            }
        }
    }

    // Pass 3: orphans — index rows no source claims.
    for target_name in all_targets {
        let Some(tdecl) = expected::target_decl(reg, target_name) else {
            continue;
        };
        let actual = scan_target_actual(
            reg,
            &mut actual_by_target,
            target_name,
            &tdecl.value_to_bytes,
        );
        let complete =
            reg.complete_indices.contains(target_name) && !protected.contains(target_name);
        let orphans: Vec<Vec<u8>> = actual
            .keys()
            .filter(|kb| !expected.contains(target_name, kb))
            .cloned()
            .collect();
        for kb in orphans {
            let stats = stats_by_name
                .entry(target_name)
                .or_insert_with(|| IndexRepairStats {
                    name: target_name,
                    ..Default::default()
                });
            if complete && opts.delete_orphans {
                tdecl.index.delete_bytes(&kb)?;
                stats.deletes += 1;
            } else {
                stats.declined_orphans += 1;
            }
        }
    }

    let mut per_index: Vec<_> = stats_by_name.into_values().collect();
    per_index.sort_by_key(|s| s.name);
    Ok(RepairReport {
        per_index,
        undecodable_sources,
    })
}