use std::sync::{Mutex, OnceLock};
use crate::error::KumoError;
static TRACER_PROVIDER: OnceLock<Mutex<Option<opentelemetry_sdk::trace::SdkTracerProvider>>> =
OnceLock::new();
pub async fn init(
service_name: impl Into<String>,
otlp_endpoint: impl Into<String>,
) -> Result<(), KumoError> {
use opentelemetry::KeyValue;
use opentelemetry::trace::TracerProvider as _;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::{Resource, trace::SdkTracerProvider};
use tracing_subscriber::prelude::*;
let service_name = service_name.into();
let endpoint = otlp_endpoint.into();
let exporter = opentelemetry_otlp::SpanExporter::builder()
.with_tonic()
.with_endpoint(&endpoint)
.build()
.map_err(|e| KumoError::store_msg(format!("otel exporter: {e}")))?;
let provider = SdkTracerProvider::builder()
.with_batch_exporter(exporter)
.with_resource(
Resource::builder()
.with_attribute(KeyValue::new("service.name", service_name.clone()))
.build(),
)
.build();
opentelemetry::global::set_tracer_provider(provider.clone());
let provider_slot = TRACER_PROVIDER.get_or_init(|| Mutex::new(None));
if let Ok(mut current) = provider_slot.lock() {
*current = Some(provider.clone());
}
let otel_layer = tracing_opentelemetry::layer().with_tracer(provider.tracer("kumo"));
tracing_subscriber::registry()
.with(tracing_subscriber::EnvFilter::from_default_env())
.with(tracing_subscriber::fmt::layer())
.with(otel_layer)
.try_init()
.map_err(|e| KumoError::store_msg(format!("tracing subscriber: {e}")))?;
tracing::info!(
service = %service_name,
endpoint = %endpoint,
"otel initialized"
);
Ok(())
}
pub fn shutdown() {
if let Some(provider_slot) = TRACER_PROVIDER.get()
&& let Ok(mut provider) = provider_slot.lock()
&& let Some(provider) = provider.take()
{
let _ = provider.shutdown();
}
}