liteboxfs 0.2.0

mod common;

use std::{
    fs,
    io::{self, Write},
    time::{Duration, Instant},
};

use criterion::{
    BenchmarkGroup, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main,
    measurement::WallTime,
};
use liteboxfs::{Connection, CreateOptions, FileKind, Owner};
use rand::{RngCore, SeedableRng, rngs::SmallRng};

use common::{KIB, TEST_INPUTS};

const TEST_SAMPLE_SIZE: usize = 20;
const TEST_CHUNK_SIZE: usize = 32 * KIB;

// To prevent unbounded memory usage, we need each test iteration to execute in a separate database
// transaction that is rolled back, both for the LiteboxFS tests and the baseline raw SQLite tests.
// To accomplish this, we need to tell Criterion to let us do the timing ourselves.
fn write_sqlite(group: &mut BenchmarkGroup<'_, WallTime>, mut conn: rusqlite::Connection) {
    let mut rng = SmallRng::from_os_rng();

    conn.execute_batch(
        r#"
        CREATE TABLE blocks (
            id INTEGER PRIMARY KEY,
            hash BLOB NOT NULL,
            data BLOB NOT NULL
        );
        "#,
    )
    .unwrap();

    for input in TEST_INPUTS {
        group.throughput(Throughput::Bytes(input.len as u64));

        // The purpose of this test is to measure the impact of overhead in the LiteboxFS
        // implementation over the theoretical maximum throughput one could expect based on the
        // fundamental design of the system (writing BLAKE3-hashed chunks to a SQLite
        // database).
        group.bench_with_input(
            BenchmarkId::new("theoretical maximum write throughput", input.to_string()),
            &input.len,
            |b, len| {
                b.iter_custom(|iters| {
                    let mut total = Duration::from_secs(0);

                    for _ in 0..iters {
                        let tx = conn.transaction().unwrap();

                        let mut stmt = tx
                            .prepare("INSERT INTO blocks (hash, data) VALUES (?, ?);")
                            .unwrap();

                        let mut source = vec![0u8; *len];
                        rng.fill_bytes(&mut source);

                        let start = Instant::now();

                        // Writing the entire input buffer to a single SQLite BLOB would actually be
                        // *slower* than LiteboxFS.
                        for chunk in source.chunks(TEST_CHUNK_SIZE) {
                            let hash = blake3::hash(chunk);
                            stmt.execute(rusqlite::params![hash.as_bytes(), chunk])
                                .unwrap();
                        }

                        total += start.elapsed();

                        drop(stmt);
                        tx.rollback().unwrap();
                    }

                    total
                })
            },
        );
    }
}

fn write_litebox(group: &mut BenchmarkGroup<'_, WallTime>, mut conn: Connection) {
    let mut rng = SmallRng::from_os_rng();

    for input in TEST_INPUTS {
        group.throughput(Throughput::Bytes(input.len as u64));

        group.bench_with_input(
            BenchmarkId::new("actual write throughput", input.to_string()),
            &input.len,
            |b, len| {
                b.iter_custom(|iters| {
                    let mut total = Duration::from_secs(0);

                    for _ in 0..iters {
                        let mut tx = conn.tx().unwrap();
                        let mut fs = tx.fs().unwrap();
                        let mut file = fs
                            .create("test.txt", FileKind::Regular, Owner::ROOT)
                            .unwrap();

                        let mut source = vec![0u8; *len];
                        rng.fill_bytes(&mut source);

                        let start = Instant::now();

                        // We need to flush here explicitly because we stop timing before the file's
                        // drop handler runs.
                        io::copy(&mut source.as_slice(), &mut file).unwrap();
                        file.flush().unwrap();

                        total += start.elapsed();

                        drop(file);
                        drop(fs);
                        tx.rollback().unwrap();
                    }

                    total
                });
            },
        );
    }
}

fn write_in_memory(c: &mut Criterion) {
    let mut group = c.benchmark_group("write performance (in-memory)");
    group.sample_size(TEST_SAMPLE_SIZE);

    let conn = rusqlite::Connection::open_in_memory().unwrap();
    write_sqlite(&mut group, conn);

    let conn = Connection::open_in_memory(&CreateOptions::default()).unwrap();
    write_litebox(&mut group, conn);

    group.finish();
}

fn write_in_memory_chunking(c: &mut Criterion) {
    let mut group = c.benchmark_group("write performance (in-memory, chunking)");
    group.sample_size(TEST_SAMPLE_SIZE);

    let opts = CreateOptions::default().chunking(true);
    let conn = Connection::open_in_memory(&opts).unwrap();
    write_litebox(&mut group, conn);

    group.finish();
}

fn write_on_disk(c: &mut Criterion) {
    let mut group = c.benchmark_group("write performance (on-disk)");
    group.sample_size(TEST_SAMPLE_SIZE);

    let conn = rusqlite::Connection::open("bench.sqlite").unwrap();
    write_sqlite(&mut group, conn);

    let conn = Connection::create_new("bench.litebox", &CreateOptions::default()).unwrap();
    write_litebox(&mut group, conn);

    group.finish();

    fs::remove_file("bench.sqlite").ok();
    fs::remove_file("bench.litebox").ok();
}

fn write_on_disk_chunking(c: &mut Criterion) {
    let mut group = c.benchmark_group("write performance (on-disk, chunking)");
    group.sample_size(TEST_SAMPLE_SIZE);

    let opts = CreateOptions::default().chunking(true);
    let conn = Connection::create_new("bench.litebox", &opts).unwrap();
    write_litebox(&mut group, conn);

    group.finish();

    fs::remove_file("bench.litebox").ok();
}

criterion_group!(
    write,
    write_in_memory,
    write_on_disk,
    write_in_memory_chunking,
    write_on_disk_chunking
);

criterion_main!(write);