rushdb-storage 3.0.7

//! Opt-in read accounting for profiling: counts every backend's (redb and
//! in-memory) point reads, range opens, rows yielded by ranges and seeks per
//! key family (`keys::family_label`). Off by default; when off, a read costs
//! one relaxed atomic load. [`take`] returns every thread's reads (profile
//! one workload at a time); [`take_local`]
//! returns only the calling thread's, which stays exact while other work
//! runs concurrently (tests).

use std::cell::RefCell;
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Mutex;

static ENABLED: AtomicBool = AtomicBool::new(false);
static COUNTS: Mutex<BTreeMap<(&'static str, String), u64>> = Mutex::new(BTreeMap::new());

thread_local! {
    static LOCAL: RefCell<BTreeMap<(&'static str, String), u64>> =
        const { RefCell::new(BTreeMap::new()) };
}

/// Turn counting on or off.
pub fn enable(on: bool) {
    ENABLED.store(on, Ordering::Relaxed);
}

/// Count one backend read of kind `op` (`get`, `scan`, `row`, `seek`) at
/// `key`.
pub(crate) fn record(op: &'static str, key: &[u8]) {
    if !ENABLED.load(Ordering::Relaxed) {
        return;
    }
    let family = rushdb_types::keys::family_label(key);
    LOCAL.with(|local| *local.borrow_mut().entry((op, family.clone())).or_default() += 1);
    let mut counts = COUNTS
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    *counts.entry((op, family)).or_default() += 1;
}

/// Add `counts` to the calling thread's tally: reads a helper thread did on
/// this thread's behalf ([`crate::parallel::run_parts`]).
pub(crate) fn add_local(counts: BTreeMap<(&'static str, String), u64>) {
    if counts.is_empty() {
        return;
    }
    LOCAL.with(|local| {
        let mut local = local.borrow_mut();
        for (key, count) in counts {
            *local.entry(key).or_default() += count;
        }
    });
}

/// Take and reset the calling thread's counts.
#[must_use]
pub fn take_local() -> BTreeMap<(&'static str, String), u64> {
    LOCAL.with(|local| std::mem::take(&mut *local.borrow_mut()))
}

/// Take and reset the counts: `(op, family) -> reads`.
#[must_use]
pub fn take() -> BTreeMap<(&'static str, String), u64> {
    std::mem::take(
        &mut *COUNTS
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner),
    )
}