walrus-rust 0.2.0

A high-performance Write-Ahead Log (WAL) implementation in Rust
Documentation
use std::cell::RefCell;
use std::fs;
use std::path::PathBuf;
use std::sync::OnceLock;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};


#[macro_export]
macro_rules! test_println {
    ($($arg:tt)*) => {
        if std::env::var("WALRUS_QUIET").is_err() {
            println!($($arg)*);
        }
    };
}

#[macro_export]
macro_rules! test_eprintln {
    ($($arg:tt)*) => {
        if std::env::var("WALRUS_QUIET").is_err() {
            eprintln!($($arg)*);
        }
    };
}

static BASE_DIR: OnceLock<PathBuf> = OnceLock::new();
static TEST_COUNTER: AtomicU64 = AtomicU64::new(0);

#[derive(Default)]
struct ThreadKeyState {
    active: Option<String>,
    last: Option<String>,
}

thread_local! {
    static THREAD_KEYS: RefCell<ThreadKeyState> = RefCell::new(ThreadKeyState::default());
}

fn ensure_base_dir() -> PathBuf {
    BASE_DIR
        .get_or_init(|| {
            let unique = format!(
                "walrus-test-run-{}-{}",
                std::process::id(),
                SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_nanos()
            );
            let dir = std::env::temp_dir().join(unique);
            let _ = fs::remove_dir_all(&dir);
            fs::create_dir_all(&dir).expect("failed to create walrus test root");
            unsafe {
                std::env::set_var("WALRUS_QUIET", "1");
                std::env::set_var("WALRUS_DATA_DIR", &dir);
            }
            dir
        })
        .clone()
}

fn next_namespace_key(counter: u64) -> String {
    format!(
        "test-key-{:x}-{:x}-{:x}",
        std::process::id(),
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos(),
        counter
    )
}

#[allow(dead_code)]
pub fn sanitize_key(key: &str) -> String {
    let mut sanitized: String = key
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || matches!(c, '-' | '_' | '.') {
                c
            } else {
                '_'
            }
        })
        .collect();

    if sanitized.trim_matches('_').is_empty() {
        sanitized = format!("ns_{:x}", checksum64(key.as_bytes()));
    }
    sanitized
}

#[allow(dead_code)]
fn checksum64(data: &[u8]) -> u64 {
    const FNV_OFFSET: u64 = 0xcbf29ce484222325;
    const FNV_PRIME: u64 = 0x00000100000001B3;
    let mut hash = FNV_OFFSET;
    for &b in data {
        hash ^= b as u64;
        hash = hash.wrapping_mul(FNV_PRIME);
    }
    hash
}

pub struct TestEnv {
    key: String,
    dir: PathBuf,
}

impl TestEnv {
    pub fn new() -> Self {
        let counter = TEST_COUNTER.fetch_add(1, Ordering::Relaxed);
        let base = ensure_base_dir();
        let key = next_namespace_key(counter);
        let dir = base.join(sanitize_key(&key));

        let _ = fs::remove_dir_all(&dir);
        fs::create_dir_all(&dir).expect("failed to create per-test walrus dir");

        THREAD_KEYS.with(|state| {
            let mut st = state.borrow_mut();
            st.active = Some(key.clone());
            st.last = Some(key.clone());
        });
        walrus_rust::wal::__set_thread_namespace_for_tests(&key);

        Self { key, dir }
    }

    pub fn namespace_key(&self) -> &str {
        &self.key
    }

    pub fn unique_key(&self, base: &str) -> String {
        format!("{}-{}", base, self.key)
    }
}

impl Drop for TestEnv {
    fn drop(&mut self) {
        let _ = fs::remove_dir_all(&self.dir);
        THREAD_KEYS.with(|state| {
            let mut st = state.borrow_mut();
            if st.active.as_deref() == Some(self.key.as_str()) {
                st.active = None;
            }
            st.last = Some(self.key.clone());
        });
        walrus_rust::wal::__clear_thread_namespace_for_tests();
    }
}

pub fn current_wal_dir() -> PathBuf {
    let mut base = ensure_base_dir();
    let key = THREAD_KEYS.with(|state| {
        let st = state.borrow();
        st.active
            .as_ref()
            .or(st.last.as_ref())
            .cloned()
            .unwrap_or_else(|| "default".to_string())
    });
    base.push(sanitize_key(&key));
    base
}

#[allow(dead_code)]
pub fn wal_root_dir() -> PathBuf {
    ensure_base_dir()
}