keyvaluedb-sqlite 0.2.3

//! Open-time corruption repair.

use keyvaluedb::KeyValueDB;
use keyvaluedb_sqlite::{Database, DatabaseConfig};

fn cfg(repair: bool) -> DatabaseConfig {
    DatabaseConfig::new()
        .with_columns(2)
        .with_repair_on_corrupt(repair)
}

async fn fill(db: &Database, col: u32, n: u32) {
    let mut tx = db.transaction();
    for i in 0..n {
        tx.put(
            col,
            format!("key_{i:04}").as_bytes(),
            format!("value_{i}").as_bytes(),
        );
    }
    db.write(tx).await.unwrap();
}

#[tokio::test]
async fn clean_open_has_no_repair_report() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("db");
    let db = Database::open(&path, cfg(false)).unwrap();
    fill(&db, 0, 10).await;
    drop(db);

    let db = Database::open(&path, cfg(true)).unwrap();
    assert!(db.repair_report().is_none());
    assert_eq!(
        db.get(0, b"key_0003").await.unwrap().unwrap(),
        b"value_3".to_vec()
    );
}

#[tokio::test]
async fn repair_salvages_readable_rows() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("db");
    let db = Database::open(&path, cfg(false)).unwrap();
    fill(&db, 0, 500).await;
    fill(&db, 1, 100).await;
    drop(db);

    // Trash one b-tree page in the middle of the file
    let mut bytes = std::fs::read(&path).unwrap();
    let page_size = u16::from_be_bytes([bytes[16], bytes[17]]) as usize;
    let target = (bytes.len() / page_size) / 2;
    let start = target * page_size;
    bytes[start..start + page_size].fill(0xFF);
    std::fs::write(&path, bytes).unwrap();

    let db = Database::open(&path, cfg(true)).unwrap();
    let report = db.repair_report().expect("repair should have run").clone();
    assert!(report.rows_recovered > 0, "report: {report:?}");
    assert!(report.corrupt_path.exists());

    // Most rows survive; the trashed page's rows are the acceptable loss
    let mut survivors = 0;
    for i in 0..500u32 {
        if db
            .get(0, format!("key_{i:04}").as_bytes())
            .await
            .unwrap()
            .is_some()
        {
            survivors += 1;
        }
    }
    assert!(survivors > 300, "survivors: {survivors}");

    // The repaired database accepts writes and reopens clean
    fill(&db, 0, 10).await;
    drop(db);
    let db = Database::open(&path, cfg(true)).unwrap();
    assert!(db.repair_report().is_none());
}

/// Craft the corruption class seen on every real corpse: table rows missing
/// from the UNIQUE autoindex, with stale index entries keeping the entry count
/// equal. quick_check passes on such a file; only integrity_check reports it.
fn detach_index_and_swap_rows(path: &std::path::Path, keys: &[&str]) {
    let conn = rusqlite::Connection::open(path).unwrap();
    let (index_rootpage, table_sql): (i64, String) = conn
        .query_row(
            "SELECT i.rootpage, t.sql FROM sqlite_master i, sqlite_master t \
             WHERE i.name='sqlite_autoindex_column_0_1' AND t.name='column_0'",
            [],
            |row| Ok((row.get(0)?, row.get(1)?)),
        )
        .unwrap();
    let no_unique_sql = table_sql.replace(" UNIQUE", "");
    assert_ne!(no_unique_sql, table_sql);
    let schema_version: i64 = conn
        .pragma_query_value(None, "schema_version", |r| r.get(0))
        .unwrap();

    conn.pragma_update(None, "writable_schema", "ON").unwrap();
    conn.execute(
        "UPDATE sqlite_master SET sql=?1 WHERE name='column_0'",
        [&no_unique_sql],
    )
    .unwrap();
    conn.execute(
        "DELETE FROM sqlite_master WHERE name='sqlite_autoindex_column_0_1'",
        [],
    )
    .unwrap();
    conn.pragma_update(None, "schema_version", schema_version + 1)
        .unwrap();
    conn.pragma_update(None, "writable_schema", "OFF").unwrap();
    drop(conn);

    // Through the UNIQUE-less schema the index is untouched: new rows never
    // reach it, and deleting one old row per new one leaves a stale entry
    // behind so the entry count still matches (what quick_check compares)
    let conn = rusqlite::Connection::open(path).unwrap();
    for key in keys {
        conn.execute(
            "INSERT INTO column_0 ([key], value) VALUES (?1, x'00')",
            [key],
        )
        .unwrap();
        conn.execute(
            "DELETE FROM column_0 WHERE id = (SELECT MIN(id) FROM column_0)",
            [],
        )
        .unwrap();
    }

    // Restore the original schema; the index is now missing the new rows
    conn.pragma_update(None, "writable_schema", "ON").unwrap();
    conn.execute(
        "UPDATE sqlite_master SET sql=?1 WHERE name='column_0'",
        [&table_sql],
    )
    .unwrap();
    conn.execute(
        "INSERT INTO sqlite_master (type, name, tbl_name, rootpage, sql) \
         VALUES ('index', 'sqlite_autoindex_column_0_1', 'column_0', ?1, NULL)",
        [index_rootpage],
    )
    .unwrap();
    conn.pragma_update(None, "schema_version", schema_version + 2)
        .unwrap();
    conn.pragma_update(None, "writable_schema", "OFF").unwrap();
    conn.pragma_update(None, "wal_checkpoint", "TRUNCATE")
        .unwrap();
    drop(conn);

    // Prove the crafted file matches the real corpses' signature
    let conn = rusqlite::Connection::open(path).unwrap();
    let quick: String = conn
        .query_row("PRAGMA quick_check(8)", [], |r| r.get(0))
        .unwrap();
    assert_eq!(quick, "ok", "crafted file must pass quick_check");
    let integrity: String = conn
        .query_row("PRAGMA integrity_check(8)", [], |r| r.get(0))
        .unwrap();
    assert!(
        integrity.contains("missing from index"),
        "crafted file must fail integrity_check: {integrity}"
    );
}

#[tokio::test]
async fn repair_index_missing_rows() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("db");
    let db = Database::open(&path, cfg(false)).unwrap();
    fill(&db, 0, 50).await;
    drop(db);

    detach_index_and_swap_rows(&path, &["stray_a", "stray_b"]);

    let db = Database::open(&path, cfg(true)).unwrap();
    let report = db
        .repair_report()
        .expect("index mismatch should trigger repair");
    assert!(report.rows_recovered >= 48, "report: {report:?}");
    assert!(report.corrupt_path.exists());
    drop(db);

    // Repaired database is integrity-clean
    let conn = rusqlite::Connection::open(&path).unwrap();
    let integrity: String = conn
        .query_row("PRAGMA integrity_check(8)", [], |r| r.get(0))
        .unwrap();
    assert_eq!(integrity, "ok");
}

/// Databases over the byte limit skip the open-time check (it reads every
/// page and would hold up open), and runtime repair still covers them.
#[tokio::test]
async fn byte_limit_skips_open_check_runtime_repair_covers() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("db");
    let db = Database::open(&path, cfg(false)).unwrap();
    fill(&db, 0, 50).await;
    drop(db);

    detach_index_and_swap_rows(&path, &["stray_a"]);

    // Over the limit: opens without checking, no repair
    let config = cfg(true).with_repair_check_byte_limit(1);
    let db = Database::open(&path, config.clone()).unwrap();
    assert!(
        db.repair_report().is_none(),
        "open must have skipped the check"
    );
    drop(db);

    // Under the limit: the same file repairs at open
    let config = cfg(true).with_repair_check_byte_limit(1 << 30);
    let db = Database::open(&path, config).unwrap();
    assert!(
        db.repair_report().is_some(),
        "open must have checked and repaired"
    );
}

/// Corruption hit by an operation mid-run repairs the database and the
/// operation's retry succeeds, without the caller seeing an error.
#[tokio::test]
async fn runtime_corruption_repairs_and_retries() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("db");
    let db = Database::open(&path, cfg(false)).unwrap();
    fill(&db, 0, 200).await;
    drop(db);

    // Trash the UNIQUE autoindex root page so index-driven reads hit
    // SQLITE_CORRUPT while the table rows stay readable
    {
        let conn = rusqlite::Connection::open(&path).unwrap();
        let rootpage: i64 = conn
            .query_row(
                "SELECT rootpage FROM sqlite_master WHERE name='sqlite_autoindex_column_0_1'",
                [],
                |r| r.get(0),
            )
            .unwrap();
        drop(conn);
        let mut bytes = std::fs::read(&path).unwrap();
        let page_size = u16::from_be_bytes([bytes[16], bytes[17]]) as usize;
        let start = (rootpage as usize - 1) * page_size;
        bytes[start..start + page_size].fill(0xFF);
        std::fs::write(&path, bytes).unwrap();
    }

    // RepairCheck::None lets the damage past the open gate, standing in for
    // corruption that appears while the database is live
    let config = cfg(true).with_repair_check(keyvaluedb_sqlite::RepairCheck::None);
    let db = Database::open(&path, config).unwrap();
    assert!(db.repair_report().is_none(), "open must not have repaired");

    // The read hits the corrupt index, repairs, retries, and succeeds
    let val = db.get(0, b"key_0100").await.unwrap();
    assert_eq!(val.unwrap(), b"value_100".to_vec());
    let report = db.repair_report().expect("runtime repair should have run");
    assert!(report.rows_recovered > 0, "report: {report:?}");
    assert!(report.corrupt_path.exists());

    // The repaired database takes writes and is integrity-clean
    fill(&db, 0, 10).await;
    drop(db);
    let conn = rusqlite::Connection::open(&path).unwrap();
    let integrity: String = conn
        .query_row("PRAGMA integrity_check(8)", [], |r| r.get(0))
        .unwrap();
    assert_eq!(integrity, "ok");
}

/// Racing operations during a runtime repair either repair-and-retry
/// themselves or wait out the repair and retry; none of them surface an error.
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn runtime_corruption_concurrent_ops() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("db");
    let db = Database::open(&path, cfg(false)).unwrap();
    fill(&db, 0, 200).await;
    drop(db);

    {
        let conn = rusqlite::Connection::open(&path).unwrap();
        let rootpage: i64 = conn
            .query_row(
                "SELECT rootpage FROM sqlite_master WHERE name='sqlite_autoindex_column_0_1'",
                [],
                |r| r.get(0),
            )
            .unwrap();
        drop(conn);
        let mut bytes = std::fs::read(&path).unwrap();
        let page_size = u16::from_be_bytes([bytes[16], bytes[17]]) as usize;
        let start = (rootpage as usize - 1) * page_size;
        bytes[start..start + page_size].fill(0xFF);
        std::fs::write(&path, bytes).unwrap();
    }

    let config = cfg(true)
        .with_repair_check(keyvaluedb_sqlite::RepairCheck::None)
        .with_num_conns(4);
    let db = Database::open(&path, config).unwrap();

    let mut tasks = Vec::new();
    for i in 0..20u32 {
        let db = db.clone();
        tasks.push(tokio::spawn(async move {
            let key = format!("key_{:04}", i * 7);
            db.get(0, key.as_bytes()).await.unwrap().unwrap()
        }));
    }
    for t in tasks {
        t.await.unwrap();
    }
    assert!(db.repair_report().is_some());
}

/// Run by hand against a copy of a real damaged store:
/// `KVDB_REPAIR_CORPSE=/path/to/copy cargo test -p keyvaluedb-sqlite --test repair -- --ignored`
#[tokio::test]
#[ignore = "needs KVDB_REPAIR_CORPSE pointing at a damaged database copy"]
async fn repair_real_corpse() {
    let src = std::env::var("KVDB_REPAIR_CORPSE").expect("set KVDB_REPAIR_CORPSE");
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("db");
    std::fs::copy(&src, &path).unwrap();

    let db = Database::open(
        &path,
        DatabaseConfig::new()
            .with_columns(1)
            .with_repair_on_corrupt(true),
    )
    .unwrap();
    let report = db.repair_report().expect("corpse should trigger repair");
    eprintln!("repair report: {report:?}");
    assert!(report.rows_recovered > 0);
}