use std::cell::Cell;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering::Relaxed};
use std::time::Instant;
#[derive(Copy, Clone)]
#[repr(usize)]
pub(crate) enum P {
Preprocess = 0,
Phase1Triggers,
BackendDispatch,
Hot,
Confirmed,
Phase2Prefilter,
Phase2KeywordAc,
Phase2SharedAc,
Phase2AnchoredVerify,
Phase2WholeChunk,
Generic,
Entropy,
Ml,
Decode,
}
const N: usize = 14;
const NAMES: [&str; N] = [
"preprocess",
"phase1",
"backend-dispatch",
"hot",
"confirmed",
"phase2:prefilter",
"phase2:keyword-ac",
"phase2:shared-ac",
"phase2:verify",
"phase2:whole-chunk",
"generic",
"entropy",
"ml",
"decode",
];
const ZEROS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N];
static NS: [AtomicU64; N] = ZEROS;
static CALLS: [AtomicU64; N] = ZEROS;
static NS_DECODE: [AtomicU64; N] = ZEROS;
static ROOT_BYTES: AtomicU64 = AtomicU64::new(0);
static ROOT_FILES: AtomicU64 = AtomicU64::new(0);
static PROFILE_ENABLED: AtomicBool = AtomicBool::new(false);
static PERF_TRACE_ENABLED: AtomicBool = AtomicBool::new(false);
pub fn set_profile_enabled(enabled: bool) {
PROFILE_ENABLED.store(enabled, Relaxed);
}
pub fn set_perf_trace_enabled(enabled: bool) {
PERF_TRACE_ENABLED.store(enabled, Relaxed);
}
pub(crate) fn enabled() -> bool {
PROFILE_ENABLED.load(Relaxed)
}
#[cfg(any(feature = "simd", feature = "gpu"))]
pub(crate) fn perf_trace_enabled() -> bool {
PERF_TRACE_ENABLED.load(Relaxed)
}
thread_local! {
static IN_DECODE: Cell<bool> = const { Cell::new(false) };
}
#[cfg(feature = "decode")]
pub(crate) fn set_in_decode(on: bool) -> bool {
IN_DECODE.with(|c| c.replace(on))
}
#[inline]
pub(crate) fn in_decode() -> bool {
IN_DECODE.with(Cell::get)
}
pub(crate) struct Guard {
p: usize,
start: Option<Instant>,
}
#[inline]
#[must_use]
pub(crate) fn span(p: P) -> Guard {
Guard {
p: p as usize,
start: if enabled() {
Some(Instant::now())
} else {
None
},
}
}
impl Drop for Guard {
#[inline]
fn drop(&mut self) {
if let Some(start) = self.start {
let ns = start.elapsed().as_nanos() as u64;
NS[self.p].fetch_add(ns, Relaxed);
CALLS[self.p].fetch_add(1, Relaxed);
if IN_DECODE.with(Cell::get) {
NS_DECODE[self.p].fetch_add(ns, Relaxed);
}
}
}
}
pub(crate) fn add_bytes(bytes: u64) {
if enabled() {
ROOT_BYTES.fetch_add(bytes, Relaxed);
}
}
pub(crate) fn add_files(files: u64) {
if enabled() {
ROOT_FILES.fetch_add(files, Relaxed);
}
}
fn read_reset() -> ([u64; N], [u64; N], [u64; N], u64, u64) {
let ns = std::array::from_fn(|i| NS[i].swap(0, Relaxed));
let calls = std::array::from_fn(|i| CALLS[i].swap(0, Relaxed));
let ns_decode = std::array::from_fn(|i| NS_DECODE[i].swap(0, Relaxed));
let bytes = ROOT_BYTES.swap(0, Relaxed);
let files = ROOT_FILES.swap(0, Relaxed);
(ns, calls, ns_decode, bytes, files)
}
pub fn reset() {
let _ = read_reset(); crate::engine::scan_inner_profile::scan_inner_profile_reset();
crate::engine::scan_postprocess::decode_profile_reset();
crate::decode::extract_profile_reset();
crate::decode::decoder_profile_reset();
crate::engine::phase2_generic::generic_profile_reset();
crate::engine::phase2::phase2_mark_stats_reset();
crate::engine::phase2::hs_mark_timing_reset();
crate::engine::scan_postprocess::ml_batch_profile_reset();
crate::gpu::ml_split_profile_reset();
}
const PHASE2_CAPTURE_LEAVES: [usize; 5] = [
P::Phase2Prefilter as usize,
P::Phase2KeywordAc as usize,
P::Phase2SharedAc as usize,
P::Phase2AnchoredVerify as usize,
P::Phase2WholeChunk as usize,
];
const PHASE2_LEAVES: [usize; 9] = [
P::Hot as usize,
P::Confirmed as usize,
P::Phase2Prefilter as usize,
P::Phase2KeywordAc as usize,
P::Phase2SharedAc as usize,
P::Phase2AnchoredVerify as usize,
P::Phase2WholeChunk as usize,
P::Generic as usize,
P::Entropy as usize,
];
pub fn dump(label: &str) {
if !enabled() {
eprintln!("[profile {label}] scanner profile switch is off; no data");
return;
}
let (ns, calls, ns_decode, bytes, files) = read_reset();
let ms = |i: usize| ns[i] as f64 / 1e6;
let sum = |ids: &[usize]| ids.iter().map(|&i| ns[i]).sum::<u64>();
let phase2_ns = sum(&PHASE2_LEAVES) + ns[P::Ml as usize];
let capture_ns = sum(&PHASE2_CAPTURE_LEAVES);
let scan_ns = ns[P::Preprocess as usize]
+ ns[P::Phase1Triggers as usize]
+ ns[P::BackendDispatch as usize]
+ phase2_ns
+ ns[P::Decode as usize];
let scan_ms = scan_ns as f64 / 1e6;
let pct = |part: u64, whole: u64| {
if whole > 0 {
100.0 * part as f64 / whole as f64
} else {
0.0
}
};
eprintln!("=== keyhog profile [{label}] ===");
let thru = if scan_ms > 0.0 {
(bytes as f64 / 1e6) / (scan_ms / 1000.0)
} else {
0.0
};
eprintln!(
"SCAN {scan_ms:>9.1} ms summed across workers · {} files · {:.2} MiB · {:.1} MB/s (pass-time sum)",
files,
bytes as f64 / (1024.0 * 1024.0),
thru
);
let leaf = |i: usize, parent_ns: u64, indent: &str| {
let c = calls[i];
let dec = ns_decode[i];
eprintln!(
"{indent}{:<16} {:>9.1} ms {:>5.1}% parent {:>6.1}% scan calls={:<8} {:>6.0} ns/call decode={:>4.1}%",
NAMES[i],
ms(i),
pct(ns[i], parent_ns),
pct(ns[i], scan_ns),
c,
if c > 0 { ns[i] as f64 / c as f64 } else { 0.0 },
pct(dec, ns[i].max(1)),
);
};
let parent = |name: &str, total: u64, indent: &str| {
eprintln!(
"{indent}{:<16} {:>9.1} ms {:>5.1}% scan",
name,
total as f64 / 1e6,
pct(total, scan_ns),
);
};
let mark: crate::engine::phase2::MarkSnapshot = crate::engine::phase2::phase2_mark_stats();
let hs_split: crate::engine::phase2::HsMarkSplit =
crate::engine::phase2::hs_mark_timing_snapshot();
leaf(P::Preprocess as usize, scan_ns, " ");
leaf(P::Phase1Triggers as usize, scan_ns, " ");
leaf(P::BackendDispatch as usize, scan_ns, " ");
parent("phase2", phase2_ns, " ");
leaf(P::Hot as usize, phase2_ns, " ");
leaf(P::Confirmed as usize, phase2_ns, " ");
parent("phase2-capture", capture_ns, " ");
for &i in &PHASE2_CAPTURE_LEAVES {
leaf(i, capture_ns, " ");
if i == P::Phase2Prefilter as usize && mark.calls > 0 {
let line = crate::engine::phase2::format_mark_decomposition(&mark);
if mark.is_consistent() {
eprintln!(" ↳ {line}");
} else {
eprintln!(
" ↳ {line} ⚠ INCONSISTENT: gate-skip + hs + regexset ({}) != calls ({}), prefilter call accounting bug",
mark.gate_skips + mark.served_total(),
mark.calls
);
}
if hs_split.any_recorded() {
eprintln!(
" ↳ {}",
crate::engine::phase2::format_hs_mark_split(&hs_split)
);
}
}
}
leaf(P::Generic as usize, phase2_ns, " ");
leaf(P::Entropy as usize, phase2_ns, " ");
leaf(P::Ml as usize, phase2_ns, " ");
leaf(P::Decode as usize, scan_ns, " ");
let decode_total: u64 = (0..N).map(|i| ns_decode[i]).sum();
eprintln!(
" (of all leaf time, {:.1}% was recorded inside decode sub-chunk rescans)",
pct(decode_total, scan_ns),
);
crate::engine::scan_inner_profile::scan_inner_profile_dump();
crate::engine::scan_postprocess::decode_profile_dump();
crate::decode::extract_profile_dump();
crate::decode::decoder_profile_dump();
crate::engine::phase2_generic::generic_profile_dump();
crate::engine::scan_postprocess::ml_batch_profile_dump();
crate::gpu::ml_split_profile_dump();
crate::engine::phase2::phase2_mark_stats_reset();
crate::engine::phase2::hs_mark_timing_reset();
}