use opentelemetry::trace::TracerProvider as _;
use opentelemetry_otlp::WithExportConfig as _;
use opentelemetry_sdk::propagation::TraceContextPropagator;
use opentelemetry_sdk::trace::{SdkTracerProvider, Tracer};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use crate::error::{ServiceError, ServiceResult};
const OTLP_TRACES_PATH: &str = "/v1/traces";
#[derive(Debug, Clone)]
pub struct TelemetryConfig {
pub service_name: String,
pub service_version: Option<String>,
pub otlp_endpoint: Option<String>,
pub sample_ratio: f64,
}
impl Default for TelemetryConfig {
fn default() -> Self {
let otlp_endpoint = std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok();
let sample_ratio = match std::env::var("OTEL_TRACES_SAMPLER_ARG") {
Ok(v) => match v.parse::<f64>() {
Ok(ratio) => ratio,
Err(_) => {
log::warn!("invalid OTEL_TRACES_SAMPLER_ARG {v:?}; defaulting to 1.0");
1.0
}
},
Err(_) => 1.0,
};
let service_name = match std::env::var("OTEL_SERVICE_NAME") {
Ok(name) => name,
Err(_) => "sunbeam-service".to_string(),
};
Self {
service_name,
service_version: None,
otlp_endpoint,
sample_ratio,
}
}
}
impl TelemetryConfig {
fn traces_endpoint(&self) -> Option<String> {
self.otlp_endpoint.as_ref().map(|endpoint| {
let trimmed = endpoint.trim_end_matches('/');
if trimmed.ends_with(OTLP_TRACES_PATH) {
trimmed.to_string()
} else {
format!("{trimmed}{OTLP_TRACES_PATH}")
}
})
}
}
pub struct TelemetryGuard {
provider: SdkTracerProvider,
}
impl Drop for TelemetryGuard {
fn drop(&mut self) {
if let Err(e) = self.provider.shutdown() {
eprintln!("sunbeam-g2v telemetry: failed to flush tracer provider: {e}");
}
}
}
pub fn tracer(config: &TelemetryConfig) -> ServiceResult<(SdkTracerProvider, Tracer)> {
opentelemetry::global::set_text_map_propagator(TraceContextPropagator::new());
let mut resource_builder =
opentelemetry_sdk::Resource::builder().with_service_name(config.service_name.clone());
if let Some(version) = &config.service_version {
resource_builder = resource_builder.with_attribute(opentelemetry::KeyValue::new(
"service.version",
version.clone(),
));
}
let mut provider_builder = SdkTracerProvider::builder()
.with_resource(resource_builder.build())
.with_sampler(opentelemetry_sdk::trace::Sampler::ParentBased(Box::new(
opentelemetry_sdk::trace::Sampler::TraceIdRatioBased(
config.sample_ratio.clamp(0.0, 1.0),
),
)));
if let Some(endpoint) = config.traces_endpoint() {
let exporter = opentelemetry_otlp::SpanExporter::builder()
.with_http()
.with_endpoint(endpoint)
.build()
.map_err(|e| ServiceError::Configuration(format!("invalid OTLP endpoint: {e}")))?;
provider_builder = provider_builder.with_batch_exporter(exporter);
} else {
tracing::debug!(msg = "no OTLP endpoint configured; spans are recorded but not exported");
}
let provider = provider_builder.build();
let tracer = provider.tracer("sunbeam-g2v");
opentelemetry::global::set_tracer_provider(provider.clone());
Ok((provider, tracer))
}
pub fn init(config: TelemetryConfig) -> ServiceResult<TelemetryGuard> {
let (provider, tracer) = tracer(&config)?;
let filter = match tracing_subscriber::EnvFilter::try_from_default_env() {
Ok(filter) => filter,
Err(_) => tracing_subscriber::EnvFilter::new("info"),
};
let registry = tracing_subscriber::registry()
.with(filter)
.with(tracing_opentelemetry::layer().with_tracer(tracer));
#[cfg(feature = "logging")]
let registry = registry.with(tracing_subscriber::fmt::layer());
registry
.try_init()
.map_err(|e| ServiceError::Configuration(format!("tracing subscriber init: {e}")))?;
Ok(TelemetryGuard { provider })
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_traces_endpoint_appends_traces_path() {
let config = TelemetryConfig {
otlp_endpoint: Some("http://localhost:4318".to_string()),
..Default::default()
};
assert_eq!(
config.traces_endpoint().as_deref(),
Some("http://localhost:4318/v1/traces")
);
}
#[test]
fn test_traces_endpoint_keeps_full_path() {
let config = TelemetryConfig {
otlp_endpoint: Some("http://localhost:4318/v1/traces/".to_string()),
..Default::default()
};
assert_eq!(
config.traces_endpoint().as_deref(),
Some("http://localhost:4318/v1/traces")
);
}
#[test]
fn test_traces_endpoint_none_when_unset() {
let config = TelemetryConfig {
otlp_endpoint: None,
..Default::default()
};
assert!(config.traces_endpoint().is_none());
}
#[test]
fn test_tracer_without_exporter_builds() {
let config = TelemetryConfig {
service_name: "test-service".to_string(),
otlp_endpoint: None,
..Default::default()
};
let (provider, _tracer) = tracer(&config).expect("tracer builds without exporter");
provider.shutdown().expect("shutdown succeeds");
}
}