use std::sync::Once;
use tracing::Span;
use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
pub static INIT: Once = Once::new();
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct CorrelationId(pub uuid::Uuid);
impl CorrelationId {
pub fn new() -> Self {
Self(uuid::Uuid::new_v4())
}
pub fn from_uuid(uuid: uuid::Uuid) -> Self {
Self(uuid)
}
#[must_use]
pub fn into_inner(self) -> uuid::Uuid {
self.0
}
#[must_use]
pub fn as_str(&self) -> String {
self.0.to_string()
}
}
impl Default for CorrelationId {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Display for CorrelationId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
impl From<uuid::Uuid> for CorrelationId {
fn from(uuid: uuid::Uuid) -> Self {
Self(uuid)
}
}
pub fn add_correlation_id(id: CorrelationId) -> impl Fn(&Span) + Clone + Send + Sync {
let id_str = id.0.to_string();
move |span: &Span| {
span.record("correlation_id", id_str.as_str());
}
}
pub fn init_tracing(filter: Option<&str>) {
INIT.call_once(|| {
let filter = filter
.and_then(|f| EnvFilter::try_from(f).ok())
.unwrap_or_else(|| {
EnvFilter::try_from("info").unwrap_or_else(|_| EnvFilter::new("info"))
});
tracing_subscriber::registry()
.with(filter)
.with(fmt::layer().with_target(true).with_thread_ids(true))
.init();
});
}
pub fn init_tracing_json(filter: Option<&str>) {
INIT.call_once(|| {
let filter = filter
.and_then(|f| EnvFilter::try_from(f).ok())
.unwrap_or_else(|| {
EnvFilter::try_from("info").unwrap_or_else(|_| EnvFilter::new("info"))
});
tracing_subscriber::registry()
.with(filter)
.with(
fmt::layer()
.json()
.with_target(true)
.with_thread_ids(true)
.with_file(true)
.with_line_number(true),
)
.init();
});
}
pub fn init_tracing_pretty() {
INIT.call_once(|| {
let filter = EnvFilter::try_from("debug").unwrap_or_else(|_| EnvFilter::new("debug"));
tracing_subscriber::registry()
.with(filter)
.with(fmt::layer().pretty().with_target(true))
.init();
});
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_correlation_id_new() {
let id = CorrelationId::new();
assert_ne!(id.0, uuid::Uuid::nil());
}
#[test]
fn test_correlation_id_default() {
let id1 = CorrelationId::default();
let id2 = CorrelationId::default();
assert_ne!(id1.0, id2.0);
}
#[test]
fn test_correlation_id_display() {
let id = CorrelationId::new();
let display = format!("{}", id);
assert_eq!(display, id.0.to_string());
}
#[test]
fn test_correlation_id_as_str() {
let id = CorrelationId::new();
let s = id.as_str();
assert_eq!(s.len(), 36); }
#[test]
fn test_correlation_id_from_uuid() {
let uuid = uuid::Uuid::new_v4();
let id = CorrelationId::from_uuid(uuid);
assert_eq!(id.0, uuid);
}
#[test]
fn test_correlation_id_into_inner() {
let uuid = uuid::Uuid::new_v4();
let id = CorrelationId::from_uuid(uuid);
assert_eq!(id.into_inner(), uuid);
}
#[test]
fn test_add_correlation_id() {
let id = CorrelationId::new();
let hook = add_correlation_id(id);
let span = tracing::info_span!("test_span");
hook(&span);
assert!(!span.is_none());
}
}