use std::sync::atomic::{AtomicBool, Ordering};
#[cfg(feature = "metrics")]
use std::sync::Arc;
use tracing::Subscriber;
use tracing_subscriber::layer::{Layer, SubscriberExt};
use tracing_subscriber::{fmt, EnvFilter};
use crate::config::{LogFormat, TelemetryConfig};
use crate::error::TelemetryError;
#[cfg(feature = "otlp")]
type OtelProvider = opentelemetry_sdk::trace::SdkTracerProvider;
#[cfg(not(feature = "otlp"))]
type OtelProvider = ();
pub struct Telemetry {
#[cfg(feature = "metrics")]
metrics: Arc<metrics_kit::Registry>,
#[cfg(feature = "otlp")]
tracer_provider: Option<OtelProvider>,
shutdown: AtomicBool,
}
impl std::fmt::Debug for Telemetry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Telemetry").finish_non_exhaustive()
}
}
impl Telemetry {
#[allow(clippy::needless_pass_by_value)]
pub fn init(config: TelemetryConfig) -> Result<Self, TelemetryError> {
#[cfg(feature = "metrics")]
let metrics = Arc::new(metrics_kit::Registry::with_max_series(
config.metrics_budget,
));
#[cfg(feature = "otlp")]
let (subscriber, provider) = build_pipeline(&config)?;
#[cfg(not(feature = "otlp"))]
let (subscriber, _) = build_pipeline(&config)?;
if tracing::subscriber::set_global_default(subscriber).is_err() {
#[cfg(feature = "otlp")]
if let Some(provider) = provider {
let _ = provider.shutdown();
}
return Err(TelemetryError::AlreadyInitialized);
}
Ok(Self {
#[cfg(feature = "metrics")]
metrics,
#[cfg(feature = "otlp")]
tracer_provider: provider,
shutdown: AtomicBool::new(false),
})
}
#[cfg(feature = "metrics")]
pub fn metrics(&self) -> Arc<metrics_kit::Registry> {
Arc::clone(&self.metrics)
}
pub fn shutdown(&self) -> Result<(), TelemetryError> {
if self.shutdown.swap(true, Ordering::AcqRel) {
return Ok(());
}
#[cfg(feature = "otlp")]
if let Some(provider) = &self.tracer_provider {
provider
.shutdown()
.map_err(|e| TelemetryError::InitFailed(e.to_string()))?;
}
Ok(())
}
}
impl Drop for Telemetry {
fn drop(&mut self) {
if self.shutdown.swap(true, Ordering::AcqRel) {
return; }
#[cfg(feature = "otlp")]
if let Some(provider) = self.tracer_provider.take() {
let _ = provider.force_flush();
}
}
}
pub fn build_subscriber(
config: &TelemetryConfig,
) -> Result<Box<dyn Subscriber + Send + Sync>, TelemetryError> {
build_pipeline(config).map(|(subscriber, _)| subscriber)
}
pub(crate) fn resolve_filter(config: &TelemetryConfig) -> Result<EnvFilter, TelemetryError> {
match std::env::var("RUST_LOG") {
Ok(directive) if !directive.trim().is_empty() => build_filter(&directive),
_ => build_filter(&config.log_level),
}
}
pub(crate) fn build_filter(directive: &str) -> Result<EnvFilter, TelemetryError> {
EnvFilter::try_new(directive).map_err(|e| {
TelemetryError::InitFailed(format!("invalid env-filter directive {directive:?}: {e}"))
})
}
fn build_pipeline(
config: &TelemetryConfig,
) -> Result<(Box<dyn Subscriber + Send + Sync>, Option<OtelProvider>), TelemetryError> {
let filter = resolve_filter(config)?;
#[cfg(not(feature = "otlp"))]
{
let subscriber = tracing_subscriber::registry()
.with(filter)
.with(build_fmt_layer(config.log_format)?);
return Ok((Box::new(subscriber), None));
}
#[cfg(feature = "otlp")]
{
let provider = build_tracer_provider(config)?;
let otel_layer = provider.as_ref().map(|p| {
use opentelemetry::trace::TracerProvider as _;
let tracer = p.tracer(config.service_name.clone());
tracing_opentelemetry::layer().with_tracer(tracer)
});
let subscriber = tracing_subscriber::registry()
.with(filter)
.with(build_fmt_layer(config.log_format)?)
.with(otel_layer);
Ok((Box::new(subscriber), provider))
}
}
#[allow(clippy::unnecessary_wraps)]
fn build_fmt_layer<S>(format: LogFormat) -> Result<Box<dyn Layer<S> + Send + Sync>, TelemetryError>
where
S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
{
match format {
LogFormat::Json => {
#[cfg(feature = "json")]
{
Ok(fmt::layer().json().with_ansi(false).boxed())
}
#[cfg(not(feature = "json"))]
{
Err(TelemetryError::InitFailed(
"LogFormat::Json requested but the `json` feature is not compiled in; \
enable the default `json` feature or select LogFormat::Pretty"
.to_owned(),
))
}
}
LogFormat::Pretty => Ok(fmt::layer().boxed()),
}
}
#[cfg(feature = "otlp")]
fn build_tracer_provider(config: &TelemetryConfig) -> Result<Option<OtelProvider>, TelemetryError> {
use opentelemetry_otlp::WithExportConfig;
let Some(endpoint) = config.otlp_endpoint.as_deref() else {
return Ok(None);
};
let resource = opentelemetry_sdk::Resource::builder()
.with_attributes([
opentelemetry::KeyValue::new("service.name", config.service_name.clone()),
opentelemetry::KeyValue::new("service.version", config.service_version.clone()),
])
.build();
let exporter = opentelemetry_otlp::SpanExporter::builder()
.with_http()
.with_endpoint(endpoint)
.build()
.map_err(|e| TelemetryError::InitFailed(format!("OTLP exporter build failed: {e}")))?;
let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
.with_resource(resource)
.with_sampler(opentelemetry_sdk::trace::Sampler::ParentBased(Box::new(
opentelemetry_sdk::trace::Sampler::TraceIdRatioBased(f64::from(config.sample_rate)),
)))
.with_batch_exporter(exporter)
.build();
Ok(Some(provider))
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used, clippy::expect_used)]
use super::*;
use tracing::level_filters::LevelFilter;
static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
#[cfg(feature = "otlp")]
fn raw_telemetry(provider: Option<OtelProvider>) -> Telemetry {
Telemetry {
#[cfg(feature = "metrics")]
metrics: Arc::new(metrics_kit::Registry::new()),
tracer_provider: provider,
shutdown: AtomicBool::new(false),
}
}
#[cfg(not(feature = "otlp"))]
fn raw_telemetry() -> Telemetry {
Telemetry {
shutdown: AtomicBool::new(false),
}
}
fn with_rust_log<T>(value: Option<&str>, f: impl FnOnce() -> T) -> T {
let guard = ENV_LOCK.lock().unwrap();
let saved = std::env::var("RUST_LOG").ok();
match value {
Some(v) => std::env::set_var("RUST_LOG", v),
None => std::env::remove_var("RUST_LOG"),
}
let out = f();
match saved {
Some(v) => std::env::set_var("RUST_LOG", v),
None => std::env::remove_var("RUST_LOG"),
}
drop(guard);
out
}
#[test]
fn build_filter_parses_directives_and_levels() {
assert_eq!(
build_filter("info").unwrap().max_level_hint(),
Some(LevelFilter::INFO)
);
assert_eq!(
build_filter("warn,hyper=debug").unwrap().max_level_hint(),
Some(LevelFilter::DEBUG)
);
assert!(matches!(
build_filter("hyper=notalevel"),
Err(TelemetryError::InitFailed(_))
));
assert!(matches!(
build_filter("a=b=c"),
Err(TelemetryError::InitFailed(_))
));
}
#[test]
fn resolve_filter_uses_config_when_rust_log_unset() {
with_rust_log(None, || {
let cfg = TelemetryConfig::new("svc").log_level("warn");
assert_eq!(
resolve_filter(&cfg).unwrap().max_level_hint(),
Some(LevelFilter::WARN)
);
});
}
#[test]
fn resolve_filter_rust_log_wins_over_config() {
with_rust_log(Some("debug"), || {
let cfg = TelemetryConfig::new("svc").log_level("info");
assert_eq!(
resolve_filter(&cfg).unwrap().max_level_hint(),
Some(LevelFilter::DEBUG)
);
});
}
#[test]
fn resolve_filter_treats_empty_rust_log_as_unset() {
with_rust_log(Some(" "), || {
let cfg = TelemetryConfig::new("svc").log_level("error");
assert_eq!(
resolve_filter(&cfg).unwrap().max_level_hint(),
Some(LevelFilter::ERROR)
);
});
}
#[test]
fn resolve_filter_propagates_invalid_rust_log() {
with_rust_log(Some("a=b=c"), || {
let cfg = TelemetryConfig::new("svc").log_level("info");
assert!(matches!(
resolve_filter(&cfg),
Err(TelemetryError::InitFailed(_))
));
});
}
#[test]
fn build_subscriber_smoke_pretty() {
let cfg = TelemetryConfig::new("svc").log_format(LogFormat::Pretty);
let subscriber = build_subscriber(&cfg).unwrap();
assert_eq!(subscriber.max_level_hint(), Some(LevelFilter::INFO));
}
#[cfg(feature = "json")]
#[test]
fn build_subscriber_smoke_json() {
let cfg = TelemetryConfig::new("svc");
let subscriber = build_subscriber(&cfg).unwrap();
assert_eq!(subscriber.max_level_hint(), Some(LevelFilter::INFO));
}
#[test]
fn init_reports_already_initialized_when_global_is_taken() {
let _ = tracing::subscriber::set_global_default(tracing::subscriber::NoSubscriber::new());
let err =
Telemetry::init(TelemetryConfig::new("svc").log_format(LogFormat::Pretty)).unwrap_err();
assert!(matches!(err, TelemetryError::AlreadyInitialized));
}
#[cfg(feature = "otlp")]
#[test]
fn init_with_otlp_reports_already_initialized_without_leaking_provider() {
let _ = tracing::subscriber::set_global_default(tracing::subscriber::NoSubscriber::new());
let cfg = TelemetryConfig::new("svc").otlp_endpoint("http://127.0.0.1:4317");
let err = Telemetry::init(cfg).unwrap_err();
assert!(matches!(err, TelemetryError::AlreadyInitialized));
}
#[cfg(feature = "otlp")]
#[test]
fn shutdown_is_idempotent() {
let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder().build();
let t = raw_telemetry(Some(provider));
t.shutdown().unwrap();
t.shutdown().unwrap();
drop(t);
}
#[cfg(feature = "otlp")]
#[test]
fn shutdown_maps_provider_error_to_init_failed_then_stays_ok() {
let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder().build();
provider.shutdown().unwrap(); let t = raw_telemetry(Some(provider));
assert!(matches!(
t.shutdown().unwrap_err(),
TelemetryError::InitFailed(_)
));
t.shutdown().unwrap();
}
#[cfg(feature = "otlp")]
#[test]
fn drop_is_a_best_effort_flush_short_circuit() {
let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder().build();
let t = raw_telemetry(Some(provider));
t.shutdown().unwrap(); drop(t);
drop(raw_telemetry(Some(
opentelemetry_sdk::trace::SdkTracerProvider::builder().build(),
)));
}
#[cfg(not(feature = "otlp"))]
#[test]
fn shutdown_without_otlp_is_trivially_idempotent() {
let t = raw_telemetry();
t.shutdown().unwrap();
t.shutdown().unwrap();
drop(t);
}
}