use rustlavel_core::events::Event;
use rustlavel_core::Json;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
tokio::task_local! {
static CURRENT: Arc<Mutex<Collected>>;
}
#[derive(Debug, Clone)]
pub struct Timing {
pub kind: String,
pub label: String,
pub duration: Option<Duration>,
pub fields: Vec<(String, String)>,
}
#[derive(Debug, Default, Clone)]
pub struct Collected {
pub timings: Vec<Timing>,
pub started: Option<Instant>,
}
impl Collected {
pub fn of(&self, kind: &str) -> Vec<&Timing> {
self.timings.iter().filter(|timing| timing.kind == kind).collect()
}
pub fn total(&self, kind: &str) -> Duration {
self.of(kind).iter().filter_map(|timing| timing.duration).sum()
}
pub fn repeated_queries(&self) -> Vec<(String, usize)> {
let mut counts: Vec<(String, usize)> = Vec::new();
for query in self.of("db.query") {
match counts.iter_mut().find(|(sql, _)| *sql == query.label) {
Some((_, count)) => *count += 1,
None => counts.push((query.label.clone(), 1)),
}
}
counts.retain(|(_, count)| *count > 1);
counts.sort_by_key(|(_, count)| std::cmp::Reverse(*count));
counts
}
}
pub struct Collector;
impl Collector {
pub fn install() {
static ONCE: std::sync::Once = std::sync::Once::new();
ONCE.call_once(|| {
rustlavel_core::events::subscribe(|event: &Event| {
let _ = CURRENT.try_with(|collected| {
if let Ok(mut collected) = collected.lock() {
collected.timings.push(timing_from(event));
}
});
});
});
}
pub async fn collect<F, T>(work: F) -> (T, Collected)
where
F: std::future::Future<Output = T>,
{
let collected = Arc::new(Mutex::new(Collected {
timings: Vec::new(),
started: Some(Instant::now()),
}));
let handle = Arc::clone(&collected);
let result = CURRENT.scope(collected, work).await;
let gathered = handle
.lock()
.map(|inner| inner.clone())
.unwrap_or_else(|poisoned| poisoned.into_inner().clone());
(result, gathered)
}
}
fn label_for(event: &Event) -> String {
let field = |name: &str| event.field(name).and_then(Json::as_str).unwrap_or("").to_string();
match event.kind {
"db.query" => field("sql"),
"cache.hit" | "cache.miss" => field("key"),
"http.client" => format!("{} {}", field("method"), field("url")),
"mail.sent" => field("to"),
"ai.call" => format!("{} {}", field("provider"), field("model")),
"mcp.call" => field("tool"),
"queue.pushed" | "queue.processed" | "queue.failed" => field("job"),
_ => event
.fields
.iter()
.next()
.map(|(key, value)| format!("{key}={}", brief(value)))
.unwrap_or_default(),
}
}
fn timing_from(event: &Event) -> Timing {
Timing {
kind: event.kind.to_string(),
label: label_for(event),
duration: event.duration,
fields: event
.fields
.iter()
.filter(|(key, _)| !is_sensitive(key))
.map(|(key, value)| (key.clone(), brief(value)))
.collect(),
}
}
fn is_sensitive(key: &str) -> bool {
const MARKERS: &[&str] = &[
"password", "passwd", "secret", "token", "authorization", "api_key", "apikey",
"api-key", "credential", "private_key", "session", "cookie",
];
let lowered = key.to_ascii_lowercase();
MARKERS.iter().any(|marker| lowered.contains(marker))
}
fn brief(value: &Json) -> String {
let text = match value {
Json::String(text) => text.clone(),
other => other.to_string(),
};
if text.chars().count() > 200 {
format!("{}…", text.chars().take(200).collect::<String>())
} else {
text
}
}
#[cfg(test)]
mod tests {
use super::*;
fn query(sql: &str, ms: u64) -> Timing {
Timing {
kind: "db.query".into(),
label: sql.into(),
duration: Some(Duration::from_millis(ms)),
fields: Vec::new(),
}
}
#[test]
fn totals_only_the_kind_asked_for() {
let collected = Collected {
timings: vec![
query("select 1", 3),
query("select 2", 4),
Timing {
kind: "cache.hit".into(),
label: "users:1".into(),
duration: Some(Duration::from_millis(50)),
fields: Vec::new(),
},
],
started: None,
};
assert_eq!(collected.of("db.query").len(), 2);
assert_eq!(collected.total("db.query"), Duration::from_millis(7));
assert_eq!(collected.total("cache.hit"), Duration::from_millis(50));
}
#[test]
fn finds_the_query_that_repeats() {
let mut timings = vec![query("select * from posts", 2)];
for _ in 0..30 {
timings.push(query("select * from users where id = ?", 1));
}
let collected = Collected { timings, started: None };
let repeated = collected.repeated_queries();
assert_eq!(repeated.len(), 1, "only the repeated one is reported");
assert_eq!(repeated[0], ("select * from users where id = ?".to_string(), 30));
}
#[test]
fn a_query_run_once_is_not_reported_as_repeated() {
let collected =
Collected { timings: vec![query("select 1", 1), query("select 2", 1)], started: None };
assert!(collected.repeated_queries().is_empty());
}
#[test]
fn sensitive_field_names_never_reach_the_page() {
for key in [
"password", "PASSWORD", "api_key", "authorization", "x-api-key",
"session_id", "cookie", "db_credential", "private_key",
] {
assert!(is_sensitive(key), "{key} should be withheld");
}
for key in ["sql", "rows", "duration_ms", "route", "status", "key"] {
assert!(!is_sensitive(key), "{key} is not sensitive and is worth showing");
}
}
#[test]
fn a_long_value_is_cut_rather_than_filling_the_page() {
let long = Json::String("x".repeat(5_000));
let shown = brief(&long);
assert!(shown.chars().count() <= 201, "got {} characters", shown.chars().count());
assert!(shown.ends_with('…'));
}
#[tokio::test]
async fn collects_what_happened_inside_the_scope_and_nothing_outside_it() {
Collector::install();
Event::new("db.query").with("sql", "select 'before'").dispatch();
let ((), collected) = Collector::collect(async {
Event::new("db.query").with("sql", "select 'inside'").dispatch();
Event::new("cache.hit").with("key", "users:1").dispatch();
})
.await;
Event::new("db.query").with("sql", "select 'after'").dispatch();
let queries = collected.of("db.query");
assert_eq!(queries.len(), 1, "only the query inside the scope");
assert_eq!(queries[0].label, "select 'inside'");
assert_eq!(collected.of("cache.hit").len(), 1);
}
#[tokio::test]
async fn two_requests_do_not_see_each_other() {
Collector::install();
let first = Collector::collect(async {
Event::new("db.query").with("sql", "first").dispatch();
tokio::task::yield_now().await;
Event::new("db.query").with("sql", "first again").dispatch();
});
let second = Collector::collect(async {
Event::new("db.query").with("sql", "second").dispatch();
tokio::task::yield_now().await;
});
let (((), one), ((), two)) = tokio::join!(first, second);
assert_eq!(one.of("db.query").len(), 2);
assert_eq!(two.of("db.query").len(), 1);
assert_eq!(two.of("db.query")[0].label, "second");
}
}