datum-mq 0.10.4

use std::{
    sync::atomic::{AtomicBool, AtomicU64, Ordering},
    time::Instant,
};

use rustix::time::{ClockId, clock_gettime};

#[derive(Debug, Clone, Copy)]
pub enum ProfileBucket {
    PollFetch,
    RecordConversion,
    ChannelHandoff,
    SourceEmit,
    DatumConsume,
    CommitBookkeeping,
}

#[derive(Debug, Clone, Copy, Default)]
pub struct ProfileBucketSnapshot {
    pub calls: u64,
    pub wall_ns: u64,
    pub cpu_ns: u64,
}

#[derive(Debug, Clone, Copy, Default)]
pub struct ProfileSnapshot {
    pub poll_fetch: ProfileBucketSnapshot,
    pub record_conversion: ProfileBucketSnapshot,
    pub channel_handoff: ProfileBucketSnapshot,
    pub source_emit: ProfileBucketSnapshot,
    pub datum_consume: ProfileBucketSnapshot,
    pub commit_bookkeeping: ProfileBucketSnapshot,
}

#[derive(Debug, Default)]
struct BucketCounters {
    calls: AtomicU64,
    wall_ns: AtomicU64,
    cpu_ns: AtomicU64,
}

#[derive(Debug, Default)]
struct ProfileCounters {
    poll_fetch: BucketCounters,
    record_conversion: BucketCounters,
    channel_handoff: BucketCounters,
    source_emit: BucketCounters,
    datum_consume: BucketCounters,
    commit_bookkeeping: BucketCounters,
}

static ENABLED: AtomicBool = AtomicBool::new(false);
static COUNTERS: ProfileCounters = ProfileCounters {
    poll_fetch: BucketCounters {
        calls: AtomicU64::new(0),
        wall_ns: AtomicU64::new(0),
        cpu_ns: AtomicU64::new(0),
    },
    record_conversion: BucketCounters {
        calls: AtomicU64::new(0),
        wall_ns: AtomicU64::new(0),
        cpu_ns: AtomicU64::new(0),
    },
    channel_handoff: BucketCounters {
        calls: AtomicU64::new(0),
        wall_ns: AtomicU64::new(0),
        cpu_ns: AtomicU64::new(0),
    },
    source_emit: BucketCounters {
        calls: AtomicU64::new(0),
        wall_ns: AtomicU64::new(0),
        cpu_ns: AtomicU64::new(0),
    },
    datum_consume: BucketCounters {
        calls: AtomicU64::new(0),
        wall_ns: AtomicU64::new(0),
        cpu_ns: AtomicU64::new(0),
    },
    commit_bookkeeping: BucketCounters {
        calls: AtomicU64::new(0),
        wall_ns: AtomicU64::new(0),
        cpu_ns: AtomicU64::new(0),
    },
};

pub fn set_enabled(enabled: bool) {
    ENABLED.store(enabled, Ordering::Relaxed);
}

pub fn set_native_enabled(enabled: bool) {
    crate::native::profile::set_enabled(enabled);
}

pub fn reset() {
    COUNTERS.reset();
}

pub fn reset_native() {
    crate::native::profile::reset();
}

#[must_use]
pub fn enabled() -> bool {
    ENABLED.load(Ordering::Relaxed)
}

#[must_use]
pub fn native_enabled() -> bool {
    crate::native::profile::enabled()
}

pub fn measure<T, F>(bucket: ProfileBucket, f: F) -> T
where
    F: FnOnce() -> T,
{
    if !enabled() {
        return f();
    }
    let wall_start = Instant::now();
    let cpu_start = thread_cpu_ns();
    let result = f();
    record(
        bucket,
        wall_start.elapsed().as_nanos() as u64,
        thread_cpu_ns().saturating_sub(cpu_start),
    );
    result
}

pub fn measure_native_datum_consume<T, F>(f: F) -> T
where
    F: FnOnce() -> T,
{
    crate::native::profile::measure(crate::native::profile::ProfileBucket::DatumConsume, f)
}

pub fn snapshot() -> ProfileSnapshot {
    ProfileSnapshot {
        poll_fetch: snapshot_bucket(&COUNTERS.poll_fetch),
        record_conversion: snapshot_bucket(&COUNTERS.record_conversion),
        channel_handoff: snapshot_bucket(&COUNTERS.channel_handoff),
        source_emit: snapshot_bucket(&COUNTERS.source_emit),
        datum_consume: snapshot_bucket(&COUNTERS.datum_consume),
        commit_bookkeeping: snapshot_bucket(&COUNTERS.commit_bookkeeping),
    }
}

pub fn native_profile_rows() -> [(&'static str, ProfileBucketSnapshot, &'static str); 9] {
    let snapshot = crate::native::profile::snapshot();
    [
        (
            "socket fetch/request",
            native_bucket(snapshot.socket_fetch),
            "Fetch request write plus response read on broker task",
        ),
        (
            "Fetch response parse",
            native_bucket(snapshot.response_parse),
            "FetchResponse frame decode and compact-field skips",
        ),
        (
            "record-batch decode",
            native_bucket(snapshot.record_decode),
            "CRC-32C plus record varints/key/header skips",
        ),
        (
            "payload copy",
            native_bucket(snapshot.payload_copy),
            "copy payload slices into one owned output buffer",
        ),
        (
            "offset bookkeeping",
            native_bucket(snapshot.offset_bookkeeping),
            "emit/commit watermark tracking",
        ),
        (
            "worker -> Source channel handoff",
            native_bucket(snapshot.channel_handoff),
            "std sync_channel send/try_send",
        ),
        (
            "Datum Source emit/envelope",
            native_bucket(snapshot.source_emit),
            "resource read plus adapter boundary",
        ),
        (
            "Datum sink consumption",
            native_bucket(snapshot.datum_consume),
            "benchmark counter over payload batch",
        ),
        (
            "commit bookkeeping",
            native_bucket(snapshot.commit_bookkeeping),
            "native broker commit request path",
        ),
    ]
}

fn record(bucket: ProfileBucket, wall_ns: u64, cpu_ns: u64) {
    let counters = match bucket {
        ProfileBucket::PollFetch => &COUNTERS.poll_fetch,
        ProfileBucket::RecordConversion => &COUNTERS.record_conversion,
        ProfileBucket::ChannelHandoff => &COUNTERS.channel_handoff,
        ProfileBucket::SourceEmit => &COUNTERS.source_emit,
        ProfileBucket::DatumConsume => &COUNTERS.datum_consume,
        ProfileBucket::CommitBookkeeping => &COUNTERS.commit_bookkeeping,
    };
    counters.calls.fetch_add(1, Ordering::Relaxed);
    counters.wall_ns.fetch_add(wall_ns, Ordering::Relaxed);
    counters.cpu_ns.fetch_add(cpu_ns, Ordering::Relaxed);
}

fn native_bucket(bucket: crate::native::profile::ProfileBucketSnapshot) -> ProfileBucketSnapshot {
    ProfileBucketSnapshot {
        calls: bucket.calls,
        wall_ns: bucket.wall_ns,
        cpu_ns: bucket.cpu_ns,
    }
}

fn snapshot_bucket(counters: &BucketCounters) -> ProfileBucketSnapshot {
    ProfileBucketSnapshot {
        calls: counters.calls.load(Ordering::Relaxed),
        wall_ns: counters.wall_ns.load(Ordering::Relaxed),
        cpu_ns: counters.cpu_ns.load(Ordering::Relaxed),
    }
}

impl ProfileCounters {
    fn reset(&self) {
        self.poll_fetch.reset();
        self.record_conversion.reset();
        self.channel_handoff.reset();
        self.source_emit.reset();
        self.datum_consume.reset();
        self.commit_bookkeeping.reset();
    }
}

impl BucketCounters {
    fn reset(&self) {
        self.calls.store(0, Ordering::Relaxed);
        self.wall_ns.store(0, Ordering::Relaxed);
        self.cpu_ns.store(0, Ordering::Relaxed);
    }
}

fn thread_cpu_ns() -> u64 {
    let ts = clock_gettime(ClockId::ThreadCPUTime);
    u64::try_from(ts.tv_sec)
        .unwrap_or(0)
        .saturating_mul(1_000_000_000)
        .saturating_add(u64::try_from(ts.tv_nsec).unwrap_or(0))
}