keyvaluedb-sqlite 0.2.4

//! WAL sizing, checkpointing, and open-under-contention behavior.
//!
//! Reproduces the field failure where WAL files sat pinned at the
//! wal_autocheckpoint high-water mark (~4MB at the sqlite default of 1000
//! pages) forever: passive checkpoints completed and restarted the WAL, but
//! nothing ever truncated the file, and a journal_size_limit above the
//! watermark never fired. The fixes hold the watermark at 256 pages, truncate
//! to 1MB on restart, and surface the checkpoint verdict instead of
//! discarding it.

use keyvaluedb::KeyValueDB;
use keyvaluedb_sqlite::{Database, DatabaseConfig, RepairCheck};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;

fn wal_path(path: &Path) -> PathBuf {
    let mut os = path.as_os_str().to_owned();
    os.push("-wal");
    PathBuf::from(os)
}

fn wal_len(path: &Path) -> u64 {
    std::fs::metadata(wal_path(path))
        .map(|m| m.len())
        .unwrap_or(0)
}

async fn commit_rows(db: &Database, start: u32, n: u32) {
    let value = vec![0xa5u8; 400];
    for i in start..start + n {
        let mut tx = db.transaction();
        tx.put(0, format!("key_{i:06}").as_bytes(), &value);
        db.write(tx).await.unwrap();
    }
}

/// The autocheckpoint keeps the WAL near its 256-page watermark under a
/// single-row-commit write load, and an explicit truncate checkpoint zeroes
/// the file.
#[tokio::test]
async fn wal_stays_bounded_and_truncates() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("db");
    let db = Database::open(&path, DatabaseConfig::new()).unwrap();

    // ~3 pages per commit; thousands of pages through a 256-page watermark
    commit_rows(&db, 0, 2000).await;

    // Watermark 256 pages = ~1.05MB; allow slack for headers and overshoot.
    // The pre-fix settings would let this reach the sqlite default watermark
    // of ~4.1MB and stay there.
    let during = wal_len(&path);
    assert!(
        during < 2 * 1024 * 1024,
        "WAL grew past the autocheckpoint watermark: {} bytes",
        during
    );

    let result = db.checkpoint().await.unwrap();
    assert!(
        result.complete(),
        "checkpoint did not complete: {:?}",
        result
    );
    assert_eq!(wal_len(&path), 0, "truncate checkpoint left WAL bytes");
}

/// The config knobs reach the connections: the pre-fix settings (sqlite
/// default 1000-page watermark, 6MB truncate limit) let the same write load
/// grow the WAL past where the defaults hold it
#[tokio::test]
async fn wal_config_overrides_apply() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("db");
    let db = Database::open(
        &path,
        DatabaseConfig::new()
            .with_wal_autocheckpoint_pages(1000)
            .with_wal_journal_size_limit(6 * 1024 * 1024),
    )
    .unwrap();

    commit_rows(&db, 0, 2000).await;

    let during = wal_len(&path);
    assert!(
        during > 2 * 1024 * 1024,
        "WAL stayed under the default watermark despite a raised override: {} bytes",
        during
    );
}

/// A read transaction on another connection pins the WAL: it grows past the
/// watermark and the checkpoint verdict says so instead of failing silently.
/// Releasing the reader lets the next checkpoint truncate.
#[tokio::test]
async fn pinned_reader_grows_wal_and_checkpoint_reports_it() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("db");
    let db = Database::open(&path, DatabaseConfig::new()).unwrap();
    commit_rows(&db, 0, 10).await;

    // Second connection with an open read snapshot, as another process or a
    // leaked reader would hold
    let raw = rusqlite::Connection::open(&path).unwrap();
    raw.execute_batch("BEGIN").unwrap();
    let _: i64 = raw
        .query_row("SELECT count(*) FROM column_0", [], |r| r.get(0))
        .unwrap();

    commit_rows(&db, 10, 1000).await;

    let pinned = wal_len(&path);
    assert!(
        pinned > 2 * 1024 * 1024,
        "expected the pinned WAL to outgrow the watermark, got {} bytes",
        pinned
    );

    let result = db.checkpoint().await.unwrap();
    assert!(
        !result.complete(),
        "checkpoint claimed completion under a pinned reader: {:?}",
        result
    );

    raw.execute_batch("COMMIT").unwrap();
    drop(raw);

    let result = db.checkpoint().await.unwrap();
    assert!(
        result.complete(),
        "checkpoint did not complete: {:?}",
        result
    );
    assert_eq!(wal_len(&path), 0);
}

/// Opening a database that carries a fat WAL from an unclean stop replays it
/// and truncates the file.
#[tokio::test]
async fn open_replays_and_truncates_stale_wal() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("db");
    let db = Database::open(&path, DatabaseConfig::new()).unwrap();

    // Pin the WAL so the frames stay in it, then snapshot the file trio as a
    // kill -9 would leave them
    let raw = rusqlite::Connection::open(&path).unwrap();
    raw.execute_batch("BEGIN").unwrap();
    let _: i64 = raw
        .query_row("SELECT count(*) FROM column_0", [], |r| r.get(0))
        .unwrap();
    commit_rows(&db, 0, 500).await;

    let copy = dir.path().join("copy");
    std::fs::copy(&path, &copy).unwrap();
    std::fs::copy(wal_path(&path), wal_path(&copy)).unwrap();
    assert!(wal_len(&copy) > 1024 * 1024);

    let db2 = Database::open(&copy, DatabaseConfig::new()).unwrap();
    assert_eq!(wal_len(&copy), 0, "open did not truncate the stale WAL");
    let val = db2.get(0, b"key_000499").await.unwrap();
    assert_eq!(val.unwrap().len(), 400, "replayed row missing after open");
}

/// A full open, integrity check included, completes in bounded time while
/// another handle to the same file commits continuously, as a second core in
/// the same process (hot restart) or another process would.
#[tokio::test(flavor = "multi_thread")]
async fn open_completes_under_concurrent_writer() {
    let dir = tempfile::tempdir().unwrap();
    let path = dir.path().join("db");
    let db = Database::open(&path, DatabaseConfig::new()).unwrap();
    commit_rows(&db, 0, 100).await;

    let stop = Arc::new(AtomicBool::new(false));
    let writer = {
        let db = db.clone();
        let stop = stop.clone();
        tokio::spawn(async move {
            let mut i = 100u32;
            while !stop.load(Ordering::Relaxed) {
                commit_rows(&db, i, 1).await;
                i += 1;
                tokio::time::sleep(Duration::from_millis(1)).await;
            }
        })
    };

    let cfg = DatabaseConfig::new()
        .with_repair_on_corrupt(true)
        .with_repair_check(RepairCheck::Full);
    let path2 = path.clone();
    let opened = tokio::time::timeout(
        Duration::from_secs(30),
        tokio::task::spawn_blocking(move || Database::open(&path2, cfg)),
    )
    .await
    .expect("open wedged under a concurrent writer")
    .unwrap()
    .unwrap();

    let val = opened.get(0, b"key_000050").await.unwrap();
    assert!(val.is_some());

    stop.store(true, Ordering::Relaxed);
    writer.await.unwrap();
}