nxd_core/

profiling.rs

//! Opt-in, process-local phase timing for operator performance diagnosis.

use serde::Serialize;
use std::borrow::Cow;
use std::cell::RefCell;
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Mutex, OnceLock};
use std::time::{Duration, Instant};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Format {
	Human,
	Structured,
}

#[derive(Clone)]
struct ActiveSpan {
	id: u64,
	worker: Option<String>,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct Event {
	id: u64,
	phase: Cow<'static, str>,
	start_offset_ms: u128,
	duration_ms: u128,
	parent: Option<u64>,
	worker: Option<String>,
}

struct State {
	format: Format,
	origin: Instant,
	events: Vec<Event>,
}

static STATE: OnceLock<Mutex<Option<State>>> = OnceLock::new();
static NEXT_ID: AtomicU64 = AtomicU64::new(1);

thread_local! {
	static ACTIVE: RefCell<Vec<ActiveSpan>> = const { RefCell::new(Vec::new()) };
}

pub fn set_format(format: Option<Format>, origin: Instant) {
	*STATE.get_or_init(|| Mutex::new(None)).lock().expect("profile state lock poisoned") =
		format.map(|format| State { format, origin, events: Vec::new() });
	ACTIVE.with(|active| active.borrow_mut().clear());
}

pub struct Span {
	id: u64,
	name: Cow<'static, str>,
	started: Option<Instant>,
	start_offset: Duration,
	parent: Option<u64>,
	worker: Option<String>,
	format: Option<Format>,
}

impl Span {
	pub fn new(name: &'static str) -> Self {
		Self::start(name.into(), None)
	}

	pub fn named(name: impl Into<String>) -> Self {
		Self::start(name.into().into(), None)
	}

	pub fn named_worker(name: impl Into<String>, worker: impl Into<String>) -> Self {
		Self::start(name.into().into(), Some(worker.into()))
	}

	fn start(name: Cow<'static, str>, worker_name: Option<String>) -> Self {
		let id = NEXT_ID.fetch_add(1, Ordering::Relaxed);
		let mut format = None;
		let mut start_offset = Duration::ZERO;
		let mut worker = None;
		if let Some(state) = STATE.get() {
			let mut state = state.lock().expect("profile state lock poisoned");
			if let Some(state) = state.as_mut() {
				format = Some(state.format);
				start_offset = state.origin.elapsed();
				worker = worker_name;
			}
		}
		let (parent, inherited_worker) = ACTIVE.with(|active| {
			active.borrow().last().map_or((None, None), |span| (Some(span.id), span.worker.clone()))
		});
		worker = worker.or(inherited_worker);
		let started = format.map(|_| Instant::now());
		if started.is_some() {
			ACTIVE.with(|active| active.borrow_mut().push(ActiveSpan { id, worker: worker.clone() }));
		}
		Self { id, name, started, start_offset, parent, worker, format }
	}
}

impl Drop for Span {
	fn drop(&mut self) {
		let Some(started) = self.started else { return };
		let duration = started.elapsed();
		ACTIVE.with(|active| {
			let mut active = active.borrow_mut();
			if let Some(index) = active.iter().rposition(|span| span.id == self.id) {
				active.remove(index);
			}
		});
		match self.format {
			Some(Format::Human) => {
				eprintln!("PROFILE phase={} duration_ms={}", self.name, duration.as_millis());
			}
			Some(Format::Structured) => record(Event {
				id: self.id,
				phase: self.name.clone(),
				start_offset_ms: self.start_offset.as_millis(),
				duration_ms: duration.as_millis(),
				parent: self.parent,
				worker: self.worker.clone(),
			}),
			None => {}
		}
	}
}

pub fn record_duration(name: &'static str, start_offset: Duration, duration: Duration) {
	let Some(state) = STATE.get() else { return };
	let format =
		state.lock().expect("profile state lock poisoned").as_ref().map(|state| state.format);
	match format {
		Some(Format::Human) => {
			eprintln!("PROFILE phase={name} duration_ms={}", duration.as_millis());
		}
		Some(Format::Structured) => record(Event {
			id: NEXT_ID.fetch_add(1, Ordering::Relaxed),
			phase: name.into(),
			start_offset_ms: start_offset.as_millis(),
			duration_ms: duration.as_millis(),
			parent: None,
			worker: None,
		}),
		None => {}
	}
}

fn record(event: Event) {
	if let Some(state) = STATE.get()
		&& let Some(state) = state.lock().expect("profile state lock poisoned").as_mut()
	{
		state.events.push(event);
	}
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct Summary<'a> {
	phase: &'a str,
	count: usize,
	total_duration_ms: u128,
}

#[derive(Serialize)]
#[serde(rename_all = "camelCase")]
struct Envelope<'a> {
	version: u8,
	stability: &'static str,
	events: &'a [Event],
	summary: Vec<Summary<'a>>,
}

fn summarize(events: &[Event]) -> Vec<Summary<'_>> {
	let mut totals = BTreeMap::<&str, (usize, u128)>::new();
	for event in events {
		let total = totals.entry(&event.phase).or_default();
		total.0 += 1;
		total.1 += event.duration_ms;
	}
	let mut summary = totals
		.into_iter()
		.map(|(phase, (count, total_duration_ms))| Summary { phase, count, total_duration_ms })
		.collect::<Vec<_>>();
	summary.sort_by(|left, right| {
		right.total_duration_ms.cmp(&left.total_duration_ms).then(left.phase.cmp(right.phase))
	});
	summary
}

fn structured_value(state: &mut State) -> serde_json::Value {
	state.events.sort_by_key(|event| (event.start_offset_ms, event.id));
	let summary = summarize(&state.events);
	serde_json::to_value(Envelope {
		version: 1,
		stability: "unstable",
		events: &state.events,
		summary,
	})
	.expect("profile envelope contains serializable values")
}

/// Return the current structured profile without consuming terminal emission.
pub fn structured_snapshot() -> Option<serde_json::Value> {
	let mut state = STATE.get()?.lock().expect("profile state lock poisoned");
	let state = state.as_mut()?;
	(state.format == Format::Structured).then(|| structured_value(state))
}

pub fn finish() {
	let Some(state) = STATE.get() else { return };
	let mut state = state.lock().expect("profile state lock poisoned");
	let Some(mut state) = state.take() else { return };
	if state.format != Format::Structured {
		return;
	}
	eprintln!("{}", structured_value(&mut state));
}

#[cfg(test)]
mod tests {
	use super::*;

	#[test]
	fn structured_envelope_marks_instability_and_aggregates_slowest_first() {
		set_format(Some(Format::Structured), Instant::now());
		let parent = Span::named_worker("repeat", "host-a");
		let parent_id = parent.id;
		{
			let _child = Span::new("repeat");
		}
		drop(parent);
		let value = structured_snapshot().unwrap();
		assert_eq!(structured_snapshot(), Some(value.clone()));
		assert_eq!(value["stability"], "unstable");
		assert_eq!(value["version"], 1);
		let child = value["events"]
			.as_array()
			.unwrap()
			.iter()
			.find(|event| event["parent"] == parent_id)
			.unwrap();
		assert_eq!(child["worker"], "host-a");
		assert_eq!(value["summary"][0]["count"], 2);
	}
}