1use serde::{Deserialize, Serialize};
9use std::cell::Cell;
10use std::fmt;
11use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
12use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
13use sysinfo::{ProcessesToUpdate, System};
14
15pub const PROFILE_SCHEMA: &str = "keyhog-profile-v1";
17
18#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
20#[serde(rename_all = "kebab-case")]
21#[repr(usize)]
22pub enum Stage {
23 SourceAcquire = 0,
24 SourceWalk,
25 SourceRead,
26 Preprocess,
27 Phase1Triggers,
28 BackendDispatch,
29 HotPatterns,
30 ConfirmedPatterns,
31 Phase2Prefilter,
32 Phase2KeywordAc,
33 Phase2SharedAc,
34 Phase2AnchoredVerify,
35 Phase2WholeChunk,
36 GenericDetection,
37 Entropy,
38 MachineLearning,
39 Decode,
40 Suppression,
41 LiveVerification,
42 Reporting,
43}
44
45impl Stage {
46 pub const ALL: [Self; 20] = [
48 Self::SourceAcquire,
49 Self::SourceWalk,
50 Self::SourceRead,
51 Self::Preprocess,
52 Self::Phase1Triggers,
53 Self::BackendDispatch,
54 Self::HotPatterns,
55 Self::ConfirmedPatterns,
56 Self::Phase2Prefilter,
57 Self::Phase2KeywordAc,
58 Self::Phase2SharedAc,
59 Self::Phase2AnchoredVerify,
60 Self::Phase2WholeChunk,
61 Self::GenericDetection,
62 Self::Entropy,
63 Self::MachineLearning,
64 Self::Decode,
65 Self::Suppression,
66 Self::LiveVerification,
67 Self::Reporting,
68 ];
69
70 #[inline]
71 const fn index(self) -> usize {
72 self as usize
73 }
74
75 pub const fn as_str(self) -> &'static str {
77 match self {
78 Self::SourceAcquire => "source-acquire",
79 Self::SourceWalk => "source-walk",
80 Self::SourceRead => "source-read",
81 Self::Preprocess => "preprocess",
82 Self::Phase1Triggers => "phase1-triggers",
83 Self::BackendDispatch => "backend-dispatch",
84 Self::HotPatterns => "hot-patterns",
85 Self::ConfirmedPatterns => "confirmed-patterns",
86 Self::Phase2Prefilter => "phase2-prefilter",
87 Self::Phase2KeywordAc => "phase2-keyword-ac",
88 Self::Phase2SharedAc => "phase2-shared-ac",
89 Self::Phase2AnchoredVerify => "phase2-anchored-verify",
90 Self::Phase2WholeChunk => "phase2-whole-chunk",
91 Self::GenericDetection => "generic-detection",
92 Self::Entropy => "entropy",
93 Self::MachineLearning => "machine-learning",
94 Self::Decode => "decode",
95 Self::Suppression => "suppression",
96 Self::LiveVerification => "live-verification",
97 Self::Reporting => "reporting",
98 }
99 }
100}
101
102const STAGE_COUNT: usize = Stage::ALL.len();
103const ZERO_COUNTERS: [AtomicU64; STAGE_COUNT] = [const { AtomicU64::new(0) }; STAGE_COUNT];
104static ENABLED: AtomicBool = AtomicBool::new(false);
105static SESSION_ACTIVE: AtomicBool = AtomicBool::new(false);
106static ELAPSED_NS: [AtomicU64; STAGE_COUNT] = ZERO_COUNTERS;
107static CALLS: [AtomicU64; STAGE_COUNT] = ZERO_COUNTERS;
108static ATTRIBUTED_NS: [AtomicU64; STAGE_COUNT] = ZERO_COUNTERS;
109static SESSION_ELAPSED_NS: [AtomicU64; STAGE_COUNT] = ZERO_COUNTERS;
110static SESSION_CALLS: [AtomicU64; STAGE_COUNT] = ZERO_COUNTERS;
111static SESSION_ATTRIBUTED_NS: [AtomicU64; STAGE_COUNT] = ZERO_COUNTERS;
112static INPUT_BYTES: AtomicU64 = AtomicU64::new(0);
113static INPUT_UNITS: AtomicU64 = AtomicU64::new(0);
114static SESSION_INPUT_BYTES: AtomicU64 = AtomicU64::new(0);
115static SESSION_INPUT_UNITS: AtomicU64 = AtomicU64::new(0);
116static RUN_SEQUENCE: AtomicU64 = AtomicU64::new(0);
117
118#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
120#[repr(u8)]
121pub enum Attribution {
122 #[default]
123 Root = 0,
124 Decoded = 1,
125}
126
127thread_local! {
128 static ATTRIBUTION: Cell<Attribution> = const { Cell::new(Attribution::Root) };
129}
130
131pub fn set_attribution(attribution: Attribution) -> Attribution {
133 ATTRIBUTION.with(|slot| slot.replace(attribution))
134}
135
136#[inline]
138pub fn enabled() -> bool {
139 ENABLED.load(Ordering::Relaxed)
140}
141
142pub fn set_enabled(enabled: bool) {
148 ENABLED.store(enabled, Ordering::Relaxed);
149}
150
151#[must_use]
153pub struct Span {
154 stage: Stage,
155 started: Option<Instant>,
156 session_recording: bool,
157}
158
159impl Span {
160 pub fn is_recording(&self) -> bool {
162 self.started.is_some()
163 }
164}
165
166#[inline]
168pub fn span(stage: Stage) -> Span {
169 let recording = enabled();
170 Span {
171 stage,
172 started: recording.then(Instant::now),
173 session_recording: recording && SESSION_ACTIVE.load(Ordering::Relaxed),
174 }
175}
176
177impl Drop for Span {
178 #[inline]
179 fn drop(&mut self) {
180 let Some(started) = self.started else {
181 return;
182 };
183 let elapsed = u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX);
184 let index = self.stage.index();
185 ELAPSED_NS[index].fetch_add(elapsed, Ordering::Relaxed);
186 CALLS[index].fetch_add(1, Ordering::Relaxed);
187 if self.session_recording {
188 SESSION_ELAPSED_NS[index].fetch_add(elapsed, Ordering::Relaxed);
189 SESSION_CALLS[index].fetch_add(1, Ordering::Relaxed);
190 }
191 if ATTRIBUTION.with(|slot| slot.get()) == Attribution::Decoded {
192 ATTRIBUTED_NS[index].fetch_add(elapsed, Ordering::Relaxed);
193 if self.session_recording {
194 SESSION_ATTRIBUTED_NS[index].fetch_add(elapsed, Ordering::Relaxed);
195 }
196 }
197 }
198}
199
200#[inline]
202pub fn add_input_bytes(bytes: u64) {
203 if enabled() {
204 INPUT_BYTES.fetch_add(bytes, Ordering::Relaxed);
205 if SESSION_ACTIVE.load(Ordering::Relaxed) {
206 SESSION_INPUT_BYTES.fetch_add(bytes, Ordering::Relaxed);
207 }
208 }
209}
210
211#[inline]
213pub fn add_input_units(units: u64) {
214 if enabled() {
215 INPUT_UNITS.fetch_add(units, Ordering::Relaxed);
216 if SESSION_ACTIVE.load(Ordering::Relaxed) {
217 SESSION_INPUT_UNITS.fetch_add(units, Ordering::Relaxed);
218 }
219 }
220}
221
222#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
224pub struct StageMeasurement {
225 pub stage: Stage,
226 pub elapsed_ns: u64,
227 pub calls: u64,
228 pub attributed_ns: u64,
229}
230
231pub fn take_stage_measurements() -> Vec<StageMeasurement> {
233 take_measurements_from(&ELAPSED_NS, &CALLS, &ATTRIBUTED_NS)
234}
235
236fn take_measurements_from(
237 elapsed: &[AtomicU64; STAGE_COUNT],
238 calls: &[AtomicU64; STAGE_COUNT],
239 attributed: &[AtomicU64; STAGE_COUNT],
240) -> Vec<StageMeasurement> {
241 Stage::ALL
242 .into_iter()
243 .filter_map(|stage| {
244 let index = stage.index();
245 let elapsed_ns = elapsed[index].swap(0, Ordering::Relaxed);
246 let calls = calls[index].swap(0, Ordering::Relaxed);
247 let attributed_ns = attributed[index].swap(0, Ordering::Relaxed);
248 (elapsed_ns != 0 || calls != 0 || attributed_ns != 0).then_some(StageMeasurement {
249 stage,
250 elapsed_ns,
251 calls,
252 attributed_ns,
253 })
254 })
255 .collect()
256}
257
258fn take_session_stage_measurements() -> Vec<StageMeasurement> {
259 take_measurements_from(&SESSION_ELAPSED_NS, &SESSION_CALLS, &SESSION_ATTRIBUTED_NS)
260}
261
262pub fn take_input_totals() -> (u64, u64) {
264 (
265 INPUT_BYTES.swap(0, Ordering::Relaxed),
266 INPUT_UNITS.swap(0, Ordering::Relaxed),
267 )
268}
269
270fn take_session_input_totals() -> (u64, u64) {
271 (
272 SESSION_INPUT_BYTES.swap(0, Ordering::Relaxed),
273 SESSION_INPUT_UNITS.swap(0, Ordering::Relaxed),
274 )
275}
276
277pub fn reset() {
279 let _ = take_stage_measurements();
280 INPUT_BYTES.store(0, Ordering::Relaxed);
281 INPUT_UNITS.store(0, Ordering::Relaxed);
282}
283
284fn reset_session() {
285 let _ = take_session_stage_measurements();
286 SESSION_INPUT_BYTES.store(0, Ordering::Relaxed);
287 SESSION_INPUT_UNITS.store(0, Ordering::Relaxed);
288}
289
290#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
292#[serde(rename_all = "kebab-case")]
293pub enum CacheState {
294 #[default]
295 Unknown,
296 Disabled,
297 Cold,
298 Warm,
299}
300
301#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
303#[serde(rename_all = "kebab-case")]
304pub enum DaemonState {
305 #[default]
306 Off,
307 Client,
308 Worker,
309 Mass,
310}
311
312#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
314#[serde(rename_all = "kebab-case")]
315pub enum RunState {
316 Created,
317 Acquiring,
318 Scanning,
319 Verifying,
320 Reporting,
321 Completed,
322 Failed,
323}
324
325#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
327pub struct RunIdentity {
328 pub run_id: String,
329 pub binary_version: String,
330 pub detector_digest: String,
331 pub config_digest: String,
332 pub source_kind: String,
333 pub workload_class: String,
334 pub backend_requested: String,
335 pub backend_selected: Option<String>,
336 pub cache_state: CacheState,
337 pub daemon_state: DaemonState,
338 pub scanner_threads: usize,
339 pub reader_threads: Option<usize>,
340 pub logical_cpus: usize,
341}
342
343impl RunIdentity {
344 pub fn new(
346 binary_version: impl Into<String>,
347 detector_digest: impl Into<String>,
348 config_digest: impl Into<String>,
349 source_kind: impl Into<String>,
350 workload_class: impl Into<String>,
351 backend_requested: impl Into<String>,
352 ) -> Self {
353 let sequence = RUN_SEQUENCE.fetch_add(1, Ordering::Relaxed);
354 let unix_ns = SystemTime::now()
355 .duration_since(UNIX_EPOCH)
356 .unwrap_or_default()
357 .as_nanos();
358 Self {
359 run_id: format!("{}-{unix_ns}-{sequence}", std::process::id()),
360 binary_version: binary_version.into(),
361 detector_digest: detector_digest.into(),
362 config_digest: config_digest.into(),
363 source_kind: source_kind.into(),
364 workload_class: workload_class.into(),
365 backend_requested: backend_requested.into(),
366 backend_selected: None,
367 cache_state: CacheState::Unknown,
368 daemon_state: DaemonState::Off,
369 scanner_threads: 0,
370 reader_threads: None,
371 logical_cpus: std::thread::available_parallelism().map_or(1, usize::from),
372 }
373 }
374}
375
376#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
378pub struct StateTransition {
379 pub state: RunState,
380 pub elapsed_ns: u64,
381}
382
383#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
385pub struct ResourceSample {
386 pub state: RunState,
387 pub elapsed_ns: u64,
388 pub snapshot: ResourceSnapshot,
389}
390
391#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)]
393pub struct ResourceSnapshot {
394 pub cpu_time_ms: Option<u64>,
395 pub resident_bytes: Option<u64>,
396 pub virtual_bytes: Option<u64>,
397 pub thread_count: Option<u64>,
398}
399
400fn process_resources() -> ResourceSnapshot {
401 let Ok(pid) = sysinfo::get_current_pid() else {
402 return ResourceSnapshot::default();
403 };
404 let mut system = System::new();
405 let pids = [pid];
406 system.refresh_processes(ProcessesToUpdate::Some(&pids), true);
407 let Some(process) = system.process(pid) else {
408 return ResourceSnapshot::default();
409 };
410 ResourceSnapshot {
411 cpu_time_ms: Some(process.accumulated_cpu_time()),
412 resident_bytes: Some(process.memory()),
413 virtual_bytes: Some(process.virtual_memory()),
414 thread_count: process.tasks().map(|tasks| tasks.len() as u64),
415 }
416}
417
418#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
420pub struct ResourceUsage {
421 pub start: ResourceSnapshot,
422 pub finish: ResourceSnapshot,
423 pub max_observed_resident_bytes: Option<u64>,
424 pub max_observed_threads: Option<u64>,
425 pub aggregate_cpu_percent: Option<f64>,
426}
427
428fn max_option(left: Option<u64>, right: Option<u64>) -> Option<u64> {
429 match (left, right) {
430 (Some(left), Some(right)) => Some(left.max(right)),
431 (left, right) => left.or(right),
432 }
433}
434
435fn resource_usage(
436 start: ResourceSnapshot,
437 finish: ResourceSnapshot,
438 wall: Duration,
439 samples: &[ResourceSample],
440) -> ResourceUsage {
441 let aggregate_cpu_percent = start
442 .cpu_time_ms
443 .zip(finish.cpu_time_ms)
444 .filter(|(start, finish)| finish >= start)
445 .and_then(|(start, finish)| {
446 let wall_ms = wall.as_secs_f64() * 1_000.0;
447 (wall_ms > 0.0).then_some((finish - start) as f64 * 100.0 / wall_ms)
448 });
449 let max_observed_resident_bytes = samples
450 .iter()
451 .filter_map(|sample| sample.snapshot.resident_bytes)
452 .fold(
453 max_option(start.resident_bytes, finish.resident_bytes),
454 |maximum, value| max_option(maximum, Some(value)),
455 );
456 let max_observed_threads = samples
457 .iter()
458 .filter_map(|sample| sample.snapshot.thread_count)
459 .fold(
460 max_option(start.thread_count, finish.thread_count),
461 |maximum, value| max_option(maximum, Some(value)),
462 );
463 ResourceUsage {
464 max_observed_resident_bytes,
465 max_observed_threads,
466 start,
467 finish,
468 aggregate_cpu_percent,
469 }
470}
471
472#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
474pub struct RunProfile {
475 pub schema: String,
476 pub identity: RunIdentity,
477 pub status: RunState,
478 pub wall_time_ns: u64,
479 pub input_bytes: u64,
480 pub input_units: u64,
481 pub stages: Vec<StageMeasurement>,
482 pub transitions: Vec<StateTransition>,
483 pub resource_samples: Vec<ResourceSample>,
484 pub resources: ResourceUsage,
485}
486
487impl RunProfile {
488 pub fn to_json_pretty(&self) -> serde_json::Result<String> {
490 serde_json::to_string_pretty(self)
491 }
492
493 pub fn render_text(&self) -> String {
495 let reader_threads = self
496 .identity
497 .reader_threads
498 .map_or_else(|| "auto".to_owned(), |threads| threads.to_string());
499 let mut output = format!(
500 "KeyHog profile {}\n\
501 state={} source={} workload={} backend_requested={} backend_selected={} cache={} daemon={} wall_ms={:.3}\n\
502 version={} detector_digest={} config_digest={}\n\
503 input_bytes={} input_units={} scanner_threads={} reader_threads={} logical_cpus={}\n",
504 self.identity.run_id,
505 state_name(self.status),
506 self.identity.source_kind,
507 self.identity.workload_class,
508 self.identity.backend_requested,
509 self.identity
510 .backend_selected
511 .as_deref()
512 .unwrap_or("unselected"),
513 cache_name(self.identity.cache_state),
514 daemon_name(self.identity.daemon_state),
515 self.wall_time_ns as f64 / 1_000_000.0,
516 self.identity.binary_version,
517 self.identity.detector_digest,
518 self.identity.config_digest,
519 self.input_bytes,
520 self.input_units,
521 self.identity.scanner_threads,
522 reader_threads,
523 self.identity.logical_cpus,
524 );
525 for stage in &self.stages {
526 output.push_str(&format!(
527 " {:<24} {:>10.3} ms calls={} attributed_ms={:.3}\n",
528 stage.stage.as_str(),
529 stage.elapsed_ns as f64 / 1_000_000.0,
530 stage.calls,
531 stage.attributed_ns as f64 / 1_000_000.0,
532 ));
533 }
534 if let Some(cpu) = self.resources.aggregate_cpu_percent {
535 output.push_str(&format!("resources aggregate_cpu={cpu:.1}%"));
536 } else {
537 output.push_str("resources aggregate_cpu=unavailable");
538 }
539 if let Some(rss) = self.resources.max_observed_resident_bytes {
540 output.push_str(&format!(" max_observed_rss_bytes={rss}"));
541 }
542 if let Some(threads) = self.resources.max_observed_threads {
543 output.push_str(&format!(" max_observed_threads={threads}"));
544 }
545 output.push('\n');
546 output
547 }
548}
549
550fn state_name(state: RunState) -> &'static str {
551 match state {
552 RunState::Created => "created",
553 RunState::Acquiring => "acquiring",
554 RunState::Scanning => "scanning",
555 RunState::Verifying => "verifying",
556 RunState::Reporting => "reporting",
557 RunState::Completed => "completed",
558 RunState::Failed => "failed",
559 }
560}
561
562fn cache_name(state: CacheState) -> &'static str {
563 match state {
564 CacheState::Unknown => "unknown",
565 CacheState::Disabled => "disabled",
566 CacheState::Cold => "cold",
567 CacheState::Warm => "warm",
568 }
569}
570
571fn daemon_name(state: DaemonState) -> &'static str {
572 match state {
573 DaemonState::Off => "off",
574 DaemonState::Client => "client",
575 DaemonState::Worker => "worker",
576 DaemonState::Mass => "mass",
577 }
578}
579
580#[derive(Clone, Copy, Debug, Eq, PartialEq)]
582pub struct SessionActive;
583
584impl fmt::Display for SessionActive {
585 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
586 formatter.write_str("a KeyHog profile session is already active in this process")
587 }
588}
589
590impl std::error::Error for SessionActive {}
591
592pub struct Session {
598 identity: Option<RunIdentity>,
599 started: Instant,
600 resources_at_start: ResourceSnapshot,
601 transitions: Vec<StateTransition>,
602 resource_samples: Vec<ResourceSample>,
603 finished: bool,
604}
605
606impl Session {
607 pub fn start(identity: RunIdentity) -> Result<Self, SessionActive> {
609 SESSION_ACTIVE
610 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
611 .map_err(|_| SessionActive)?;
612 reset();
613 reset_session();
614 set_enabled(true);
615 let started = Instant::now();
616 let resources_at_start = process_resources();
617 Ok(Self {
618 identity: Some(identity),
619 started,
620 resources_at_start,
621 transitions: vec![StateTransition {
622 state: RunState::Created,
623 elapsed_ns: 0,
624 }],
625 resource_samples: vec![ResourceSample {
626 state: RunState::Created,
627 elapsed_ns: 0,
628 snapshot: resources_at_start,
629 }],
630 finished: false,
631 })
632 }
633
634 pub fn identity_mut(&mut self) -> &mut RunIdentity {
636 self.identity
637 .as_mut()
638 .expect("unfinished profile owns identity")
639 }
640
641 pub fn transition(&mut self, state: RunState) {
643 let elapsed_ns = u64::try_from(self.started.elapsed().as_nanos()).unwrap_or(u64::MAX);
644 self.transitions.push(StateTransition { state, elapsed_ns });
645 self.resource_samples.push(ResourceSample {
646 state,
647 elapsed_ns,
648 snapshot: process_resources(),
649 });
650 }
651
652 pub fn finish(mut self, status: RunState) -> RunProfile {
654 self.transition(status);
655 set_enabled(false);
656 let wall = self.started.elapsed();
657 let finish_resources = self
658 .resource_samples
659 .last()
660 .map_or(self.resources_at_start, |sample| sample.snapshot);
661 let (input_bytes, input_units) = take_session_input_totals();
662 let stages = take_session_stage_measurements();
663 reset();
664 let resource_samples = std::mem::take(&mut self.resource_samples);
665 let resources = resource_usage(
666 self.resources_at_start,
667 finish_resources,
668 wall,
669 &resource_samples,
670 );
671 let profile = RunProfile {
672 schema: PROFILE_SCHEMA.to_string(),
673 identity: self
674 .identity
675 .take()
676 .expect("unfinished profile owns identity"),
677 status,
678 wall_time_ns: u64::try_from(wall.as_nanos()).unwrap_or(u64::MAX),
679 input_bytes,
680 input_units,
681 stages,
682 transitions: std::mem::take(&mut self.transitions),
683 resource_samples,
684 resources,
685 };
686 self.finished = true;
687 SESSION_ACTIVE.store(false, Ordering::Release);
688 profile
689 }
690}
691
692impl Drop for Session {
693 fn drop(&mut self) {
694 if !self.finished {
695 set_enabled(false);
696 reset();
697 reset_session();
698 SESSION_ACTIVE.store(false, Ordering::Release);
699 }
700 }
701}