use crate::{ObservabilityError, Result};
use opentelemetry::trace::TracerProvider;
use opentelemetry::KeyValue;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::{trace as sdktrace, Resource};
use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
pub struct TracingConfig {
pub service_name: String,
pub otlp_endpoint: String,
pub environment: Option<String>,
pub version: Option<String>,
}
impl TracingConfig {
pub fn new(service_name: impl Into<String>, otlp_endpoint: impl Into<String>) -> Self {
Self {
service_name: service_name.into(),
otlp_endpoint: otlp_endpoint.into(),
environment: None,
version: None,
}
}
pub fn with_environment(mut self, env: impl Into<String>) -> Self {
self.environment = Some(env.into());
self
}
pub fn with_version(mut self, version: impl Into<String>) -> Self {
self.version = Some(version.into());
self
}
}
pub fn init_distributed_tracing(service_name: &str, endpoint: &str) -> Result<()> {
let config = TracingConfig::new(service_name, endpoint);
init_distributed_tracing_with_config(config)
}
pub fn init_distributed_tracing_with_config(config: TracingConfig) -> Result<()> {
let mut resource_attrs = vec![
KeyValue::new("service.name", config.service_name.clone()),
];
if let Some(env) = &config.environment {
resource_attrs.push(KeyValue::new("deployment.environment", env.clone()));
}
if let Some(version) = &config.version {
resource_attrs.push(KeyValue::new("service.version", version.clone()));
}
let resource = Resource::builder()
.with_attributes(resource_attrs)
.build();
let exporter = opentelemetry_otlp::SpanExporter::builder()
.with_http()
.with_endpoint(&config.otlp_endpoint)
.build()
.map_err(|e| ObservabilityError::TracingInit(format!("Failed to create OTLP exporter: {}", e)))?;
let provider = sdktrace::SdkTracerProvider::builder()
.with_batch_exporter(exporter)
.with_resource(resource)
.build();
let tracer = provider.tracer(config.service_name.clone());
opentelemetry::global::set_tracer_provider(provider);
let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer);
let filter = EnvFilter::try_from_default_env()
.unwrap_or_else(|_| EnvFilter::new("info"));
tracing_subscriber::registry()
.with(filter)
.with(
tracing_subscriber::fmt::layer()
.with_target(false)
.with_thread_ids(true)
.with_level(true)
.json(),
)
.with(otel_layer)
.init();
tracing::info!(
service = config.service_name,
endpoint = config.otlp_endpoint,
"Distributed tracing initialized"
);
Ok(())
}
pub fn shutdown_tracing() {
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tracing_config_builder() {
let config = TracingConfig::new("test-service", "http://localhost:4318")
.with_environment("test")
.with_version("abc123");
assert_eq!(config.service_name, "test-service");
assert_eq!(config.otlp_endpoint, "http://localhost:4318");
assert_eq!(config.environment, Some("test".to_string()));
assert_eq!(config.version, Some("abc123".to_string()));
}
}