cberner / redb

An embedded key-value database in pure Rust
https://www.redb.org
Apache License 2.0
3.07k stars 137 forks source link

Possible test for crash consistency #786

Closed cherryman closed 2 months ago

cherryman commented 3 months ago

Wanted to stress-test crash consistency before using so decided to take the approach below. Sharing in case it's useful to anybody, feel free to close this issue 🙂. Repeatedly ran the program below using parallel.

$ seq 10000000 | parallel -j50 ./target/release/woohoo

Cargo.toml:

[package]
name = "woohoo"
version = "0.1.0"
edition = "2021"

[dependencies]
nix = { version = "0.28.0", features = ["mman", "fs", "process", "signal"] }
rand = "0.8.5"
rand_pcg = "0.3.1"
redb = "2.0.0"
tempfile = "3"

main.rs:

use nix::sys::mman::{mmap_anonymous, munmap, MapFlags, ProtFlags};
use nix::sys::signal::{kill, Signal};
use nix::unistd::{fork, ForkResult};
use rand::{Rng as _, SeedableRng as _};
use redb::{Durability, ReadableTable, TableDefinition};
use std::num::NonZeroUsize;
use std::sync::atomic::{AtomicU64, Ordering};

const TBL: TableDefinition<u64, &[u8]> = TableDefinition::new("default");

fn main() {
    let jnl_path = tempfile::tempdir().unwrap();
    let jnl_path = jnl_path.path().join("journal.db");
    let mem_size = std::mem::size_of::<AtomicU64>();
    let data_ptr = unsafe {
        // SAFETY: The resulting pointer is non-null and page-aligned.
        //
        // From `mmap(7)`:
        //
        // > If addr is NULL, then the kernel chooses the (page-aligned) address
        // > at which to create the mapping; this is the most portable method of
        // > creating a new  mapping
        mmap_anonymous(
            None,
            NonZeroUsize::new(mem_size).unwrap(),
            ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
            MapFlags::MAP_SHARED,
        )
        .unwrap()
    };
    let data = unsafe {
        let ptr = data_ptr.cast::<AtomicU64>().as_ptr();

        // SAFETY: `ptr` is valid and aligned; `.write()` has move semantics.
        ptr.write(AtomicU64::new(0));

        // SAFETY: `ptr` is valid; mutation uses `UnsafeCell`.
        ptr.as_ref().unwrap()
    };
    let mut rng = rand_pcg::Pcg64Mcg::from_entropy();
    let xs = (0..1000)
        .map(|_| {
            let len = rng.gen_range(0..2048);
            let mut x = vec![0; len];
            rng.fill(&mut x[..]);
            x
        })
        .collect::<Vec<Vec<u8>>>();

    match unsafe {
        // SAFETY: No locks prior to the fork are used.
        fork().expect("Failed to fork")
    } {
        ForkResult::Child => {
            let db = redb::Database::builder().create(jnl_path).unwrap();
            // Starting at 1 since `data` starts at 0.
            for k in 1u64.. {
                let mut tx = db.begin_write().unwrap();
                let mut t = tx.open_table(TBL).unwrap();
                t.insert(&k, &xs[k as usize % xs.len()][..]).unwrap();
                drop(t);
                tx.set_durability(Durability::Immediate);
                tx.commit().unwrap();
                data.store(k, Ordering::SeqCst);
            }
        }
        ForkResult::Parent { child } => {
            std::thread::yield_now();
            std::thread::sleep(std::time::Duration::from_millis(500));
            kill(child, Signal::SIGKILL).unwrap();
            let last_id = data.load(Ordering::SeqCst);
            std::thread::sleep(std::time::Duration::from_millis(50));
            let db = redb::Database::builder().create(jnl_path).unwrap();

            // To be correct, last must be either equal to or one ahead
            // of the value in `data`.
            let tx = db.begin_read().unwrap();
            let t = tx.open_table(TBL).unwrap();
            let last = t.last().unwrap().map(|(k, _)| k.value()).unwrap_or(0);
            assert!(last == last_id || last - 1 == last_id);

            for x in t.range(0u64..).unwrap() {
                let (k, v) = x.unwrap();
                let (k, v) = (k.value(), v.value());
                assert_eq!(v, xs[k as usize % xs.len()]);
            }
        }
    }

    unsafe {
        // SAFETY: Values passed here are valid based on safety of
        // operations above. `data` is also no longer aliased.
        munmap(data_ptr, mem_size).unwrap();
    }
}
cberner commented 3 months ago

Neat. Glad it passed :) The fuzzer in fuzz/ folder also does some crash testing