o41-training-data 0.1.3

//! Opt-in numerical counters for a single-loader benchmark process.
//! Wall times overlap across workers and nested phases; do not add them.
use std::sync::{
    Mutex, OnceLock,
    atomic::{AtomicU64, Ordering::Relaxed},
};
use std::time::Instant;

const NAMES: [&str; 12] = [
    "s3",
    "decompress",
    "parse",
    "value",
    "arrow",
    "projection",
    "assembly",
    "shard_wait",
    "next_wait",
    "queue_wait",
    "export",
    "prefetch_wait",
];
static WALL: [AtomicU64; 12] = [const { AtomicU64::new(0) }; 12];
static CPU: [AtomicU64; 12] = [const { AtomicU64::new(0) }; 12];
static CALLS: [AtomicU64; 12] = [const { AtomicU64::new(0) }; 12];
static START: OnceLock<Instant> = OnceLock::new();
static ENABLED: OnceLock<bool> = OnceLock::new();
static REQUESTS: AtomicU64 = AtomicU64::new(0);
static BYTES: AtomicU64 = AtomicU64::new(0);
static TTFB: AtomicU64 = AtomicU64::new(0);
static ACTIVE: AtomicU64 = AtomicU64::new(0);
static PEAK: AtomicU64 = AtomicU64::new(0);
static BYTE_NS: AtomicU64 = AtomicU64::new(0);
#[derive(Default, serde::Serialize)]
struct ShardReads {
    fetches: u64,
    probes: u64,
    decoded_bytes: u64,
    projected_bytes: u64,
}
static SHARDS: OnceLock<Mutex<std::collections::BTreeMap<String, ShardReads>>> = OnceLock::new();
pub fn shard(uri: &str, probe: bool, decoded: u64, projected: u64) {
    if enabled() {
        let mut shards = SHARDS.get_or_init(Default::default).lock().unwrap();
        let entry = shards.entry(uri.to_owned()).or_default();
        entry.fetches += u64::from(!probe);
        entry.probes += u64::from(probe);
        entry.decoded_bytes += decoded;
        entry.projected_bytes += projected;
    }
}
fn enabled() -> bool {
    *ENABLED.get_or_init(|| std::env::var_os("O41_TRAINING_DATA_PROFILE").is_some())
}
fn cpu() -> u64 {
    let mut t = libc::timespec {
        tv_sec: 0,
        tv_nsec: 0,
    };
    unsafe {
        libc::clock_gettime(libc::CLOCK_THREAD_CPUTIME_ID, &mut t);
    }
    t.tv_sec as u64 * 1_000_000_000 + t.tv_nsec as u64
}
pub struct Timer(usize, Option<(Instant, u64)>);
impl Timer {
    pub fn new(phase: usize) -> Self {
        Self(
            phase,
            enabled().then(|| {
                START.get_or_init(Instant::now);
                (Instant::now(), cpu())
            }),
        )
    }
}
impl Drop for Timer {
    fn drop(&mut self) {
        if let Some((wall, start_cpu)) = self.1 {
            WALL[self.0].fetch_add(wall.elapsed().as_nanos() as u64, Relaxed);
            CPU[self.0].fetch_add(cpu() - start_cpu, Relaxed);
            CALLS[self.0].fetch_add(1, Relaxed);
        }
    }
}
pub struct Request(Option<(Instant, u64)>);
impl Request {
    pub fn new(bytes: u64) -> Self {
        Self(enabled().then(|| {
            REQUESTS.fetch_add(1, Relaxed);
            let active = ACTIVE.fetch_add(bytes, Relaxed) + bytes;
            PEAK.fetch_max(active, Relaxed);
            (Instant::now(), bytes)
        }))
    }
    pub fn headers(&self) {
        if let Some((start, _)) = self.0 {
            TTFB.fetch_add(start.elapsed().as_nanos() as u64, Relaxed);
        }
    }
    pub fn bytes(&self, n: usize) {
        if self.0.is_some() {
            BYTES.fetch_add(n as u64, Relaxed);
        }
    }
}
impl Drop for Request {
    fn drop(&mut self) {
        if let Some((start, bytes)) = self.0 {
            ACTIVE.fetch_sub(bytes, Relaxed);
            BYTE_NS.fetch_add(
                (start.elapsed().as_nanos() * bytes as u128).min(u64::MAX as u128) as u64,
                Relaxed,
            );
        }
    }
}
pub fn report() {
    if !enabled() {
        return;
    }
    let elapsed = START.get().map_or(0., |t| t.elapsed().as_secs_f64());
    let shards = SHARDS.get_or_init(Default::default).lock().unwrap();
    let unique = shards.values().filter(|s| s.fetches > 0).count();
    let repeated: u64 = shards.values().map(|s| s.fetches.saturating_sub(1)).sum();
    eprintln!(
        "{}",
        serde_json::json!({"shards": *shards,
        "unique_shard_fetches": unique, "repeated_shard_fetches": repeated})
    );
    let phases: serde_json::Map<String, serde_json::Value> = NAMES.iter().enumerate().map(|(i, name)| (
        name.to_string(), serde_json::json!({"wall_s": WALL[i].load(Relaxed) as f64 / 1e9, "cpu_s": CPU[i].load(Relaxed) as f64 / 1e9, "calls": CALLS[i].load(Relaxed)})
    )).collect();
    eprintln!(
        "{}",
        serde_json::json!({"profile": phases, "elapsed_s": elapsed, "requests": REQUESTS.load(Relaxed), "bytes": BYTES.load(Relaxed), "headers_s": TTFB.load(Relaxed) as f64 / 1e9, "peak_outstanding_bytes": PEAK.load(Relaxed), "mean_outstanding_bytes": BYTE_NS.load(Relaxed) as f64 / 1e9 / elapsed})
    );
}