kithara-test-utils 0.0.1-alpha4

kithara test runtime: probe/hang/mock helpers (no-op in release).
use std::collections::HashMap;

use kithara_platform::time::Instant;
use tracing::{Subscriber, field::Visit};
use tracing_subscriber::{Layer, layer::Context, registry::LookupSpan};

use super::{event::ProbeEvent, recorder::SharedLog};

/// Build a `tracing_subscriber` `Layer` that captures probe events
/// into the shared log. Composed into the global subscriber by
/// [`crate::test::init_tracing`].
#[must_use]
pub fn probe_layer<S>() -> impl Layer<S>
where
    S: Subscriber + for<'a> LookupSpan<'a>,
{
    ProbeLayer {
        log: super::recorder::shared_log(),
    }
}

struct ProbeLayer {
    log: SharedLog,
}

impl<S: Subscriber + for<'a> LookupSpan<'a>> Layer<S> for ProbeLayer {
    fn on_event(&self, event: &tracing::Event<'_>, _ctx: Context<'_, S>) {
        let target = event.metadata().target();
        if !target.ends_with("_probe") {
            return;
        }
        let mut visitor = ProbeVisitor::default();
        event.record(&mut visitor);
        let evt = ProbeEvent {
            target: target.to_string(),
            at: Instant::now(),
            fields: visitor.numeric,
            string_fields: visitor.strings,
        };
        if let Ok(mut log) = self.log.lock() {
            log.push(evt);
        }
    }
}

#[derive(Default)]
struct ProbeVisitor {
    numeric: HashMap<String, u64>,
    strings: HashMap<String, String>,
}

impl Visit for ProbeVisitor {
    fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
        self.numeric
            .insert(field.name().to_string(), u64::from(value));
    }
    fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) {
        let formatted = format!("{value:?}");
        if let Ok(parsed) = formatted.parse::<u64>() {
            self.numeric.insert(field.name().to_string(), parsed);
        } else if let Ok(parsed) = formatted.parse::<i64>() {
            if parsed >= 0 {
                self.numeric
                    .insert(field.name().to_string(), u64::try_from(parsed).unwrap_or(0));
            } else {
                self.strings.insert(field.name().to_string(), formatted);
            }
        } else {
            self.strings.insert(field.name().to_string(), formatted);
        }
    }
    fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
        if value >= 0 {
            self.numeric
                .insert(field.name().to_string(), u64::try_from(value).unwrap_or(0));
        }
    }
    fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
        self.strings
            .insert(field.name().to_string(), value.to_string());
    }
    fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
        self.numeric.insert(field.name().to_string(), value);
    }
}