use ahash::AHashMap;
use arc_swap::ArcSwap;
use opentelemetry::global::{BoxedTracer, ObjectSafeTracerProvider};
use opentelemetry::{InstrumentationScope, trace};
use std::sync::Arc;
use std::sync::LazyLock;
#[derive(Clone)]
pub struct InstanceTracerProvider {
provider: Arc<dyn ObjectSafeTracerProvider + Send + Sync>,
}
impl InstanceTracerProvider {
fn new<P, T, S>(provider: P) -> Self
where
S: trace::Span + Send + Sync + 'static,
T: trace::Tracer<Span = S> + Send + Sync + 'static,
P: trace::TracerProvider<Tracer = T> + Send + Sync + 'static,
{
InstanceTracerProvider {
provider: Arc::new(provider),
}
}
}
impl trace::TracerProvider for InstanceTracerProvider {
type Tracer = BoxedTracer;
fn tracer_with_scope(&self, scope: InstrumentationScope) -> Self::Tracer {
BoxedTracer::new(self.provider.boxed_tracer(scope))
}
}
type TracerProviders = AHashMap<String, InstanceTracerProvider>;
static TRACER_PROVIDER_MAP: LazyLock<ArcSwap<TracerProviders>> =
LazyLock::new(|| ArcSwap::from_pointee(AHashMap::new()));
pub fn add_provider(
name: &str,
provider: opentelemetry_sdk::trace::SdkTracerProvider,
) {
let mut m: TracerProviders = AHashMap::new();
for (name, provider) in TRACER_PROVIDER_MAP.load().iter() {
m.insert(name.to_string(), provider.clone());
}
m.insert(name.to_string(), InstanceTracerProvider::new(provider));
TRACER_PROVIDER_MAP.store(Arc::new(m));
}
#[inline]
pub fn get_provider(name: &str) -> Option<InstanceTracerProvider> {
TRACER_PROVIDER_MAP.load().get(name).cloned()
}