diff --git a/crates/metrics/src/lib.rs b/crates/metrics/src/lib.rs index e28ee58..154fdc6 100644 --- a/crates/metrics/src/lib.rs +++ b/crates/metrics/src/lib.rs @@ -1,7 +1,16 @@ -pub use {future::*, once_cell::sync::Lazy, opentelemetry as otel, task::*}; +pub use { + future::*, + once_cell::sync::Lazy, + opentelemetry as otel, + opentelemetry_sdk as otel_sdk, + task::*, +}; use { - opentelemetry_sdk::metrics::SdkMeterProvider, - otel::metrics::{Meter, MeterProvider}, + opentelemetry_sdk::metrics::{MeterProviderBuilder, SdkMeterProvider}, + otel::{ + global, + metrics::{Meter, MeterProvider}, + }, prometheus::{Error as PrometheusError, Registry, TextEncoder}, std::{ sync::{Arc, Mutex}, @@ -17,19 +26,27 @@ const DEFAULT_SERVICE_NAME: &str = "unknown_service"; static SERVICE_NAME: Mutex> = Mutex::new(None); +pub type MeterProviderBuilderFn = fn(MeterProviderBuilder) -> MeterProviderBuilder; +static METER_PROVIDER_BUILDER_FN: Mutex> = Mutex::new(None); + static METRICS_CORE: Lazy> = Lazy::new(|| { let service_name = SERVICE_NAME.lock().unwrap().unwrap_or(DEFAULT_SERVICE_NAME); + let meter_provider_builder = *METER_PROVIDER_BUILDER_FN.lock().unwrap(); let registry = Registry::new(); let prometheus_exporter = opentelemetry_prometheus::exporter() .with_registry(registry.clone()) .build() .unwrap(); - let provider = SdkMeterProvider::builder() - .with_reader(prometheus_exporter) - .build(); + let mut builder = SdkMeterProvider::builder().with_reader(prometheus_exporter); + if let Some(buidler_fn) = meter_provider_builder { + builder = buidler_fn(builder); + }; + let provider = builder.build(); let meter = provider.meter(service_name); + global::set_meter_provider(provider); + Arc::new(ServiceMetrics { registry, meter }) }); @@ -64,6 +81,22 @@ impl ServiceMetrics { Lazy::force(&METRICS_CORE); } + /// Initializes service metrics with the specified name and meter provider + /// builder. + /// + /// # Panics + /// + /// Panics if either prometheus exporter or opentelemetry meter fails to + /// initialize. + pub fn init_with_meter_builder( + name: &'static str, + meter_provider_builder: MeterProviderBuilderFn, + ) { + *SERVICE_NAME.lock().unwrap() = Some(name); + *METER_PROVIDER_BUILDER_FN.lock().unwrap() = Some(meter_provider_builder); + Lazy::force(&METRICS_CORE); + } + /// Generates export data in Prometheus format, serialized into string. pub fn export() -> Result { let data = Self::get().registry.gather();