use std::fmt;
use std::time::Duration;
use opentelemetry::global;
use opentelemetry::trace::TracerProvider as _;
use opentelemetry::KeyValue;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::propagation::TraceContextPropagator;
use opentelemetry_sdk::trace::{RandomIdGenerator, Sampler};
use opentelemetry_sdk::Resource;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;
use tracing_subscriber::EnvFilter;
#[derive(Debug, Clone)]
pub struct TelemetryConfig {
pub service_name: String,
pub service_version: String,
pub environment: String,
pub otlp_endpoint: Option<String>,
pub sampling_ratio: f64,
pub enable_export: bool,
}
impl Default for TelemetryConfig {
fn default() -> Self {
Self {
service_name: "rs3gw".to_string(),
service_version: env!("CARGO_PKG_VERSION").to_string(),
environment: std::env::var("ENVIRONMENT").unwrap_or_else(|_| "development".to_string()),
otlp_endpoint: std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").ok(),
sampling_ratio: std::env::var("OTEL_TRACES_SAMPLER_ARG")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(1.0),
enable_export: std::env::var("OTEL_TRACES_EXPORTER")
.map(|e| e != "none")
.unwrap_or(false),
}
}
}
#[derive(Debug)]
pub struct TelemetryError(String);
impl fmt::Display for TelemetryError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "Telemetry error: {}", self.0)
}
}
impl std::error::Error for TelemetryError {}
pub fn init_telemetry(config: TelemetryConfig) -> Result<(), TelemetryError> {
global::set_text_map_propagator(TraceContextPropagator::new());
if config.enable_export {
let resource = Resource::builder_empty()
.with_service_name(config.service_name.clone())
.with_attribute(KeyValue::new(
opentelemetry_semantic_conventions::resource::SERVICE_VERSION,
config.service_version.clone(),
))
.with_attribute(KeyValue::new(
"deployment.environment",
config.environment.clone(),
))
.build();
let sampler = if config.sampling_ratio >= 1.0 {
Sampler::AlwaysOn
} else if config.sampling_ratio <= 0.0 {
Sampler::AlwaysOff
} else {
Sampler::TraceIdRatioBased(config.sampling_ratio)
};
let endpoint = config
.otlp_endpoint
.as_deref()
.unwrap_or("http://localhost:4317");
let exporter = opentelemetry_otlp::SpanExporter::builder()
.with_tonic()
.with_endpoint(endpoint)
.with_timeout(Duration::from_secs(10))
.build()
.map_err(|e| TelemetryError(format!("Failed to create OTLP exporter: {}", e)))?;
let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
.with_batch_exporter(exporter)
.with_resource(resource)
.with_sampler(sampler)
.with_id_generator(RandomIdGenerator::default())
.build();
global::set_tracer_provider(provider.clone());
let telemetry_layer = tracing_opentelemetry::layer()
.with_tracer(provider.tracer(config.service_name.clone()));
let env_filter = EnvFilter::try_from_default_env()
.or_else(|_| EnvFilter::try_new("info"))
.map_err(|e| TelemetryError(format!("Failed to create env filter: {}", e)))?;
tracing_subscriber::registry()
.with(env_filter)
.with(telemetry_layer)
.with(tracing_subscriber::fmt::layer().with_target(true))
.try_init()
.map_err(|e| TelemetryError(format!("Failed to initialize subscriber: {}", e)))?;
tracing::info!(
service_name = %config.service_name,
service_version = %config.service_version,
environment = %config.environment,
otlp_endpoint = %endpoint,
sampling_ratio = %config.sampling_ratio,
"OpenTelemetry distributed tracing initialized"
);
} else {
let env_filter = EnvFilter::try_from_default_env()
.or_else(|_| EnvFilter::try_new("info"))
.map_err(|e| TelemetryError(format!("Failed to create env filter: {}", e)))?;
tracing_subscriber::registry()
.with(env_filter)
.with(tracing_subscriber::fmt::layer().with_target(true))
.try_init()
.map_err(|e| TelemetryError(format!("Failed to initialize subscriber: {}", e)))?;
tracing::info!(
service_name = %config.service_name,
service_version = %config.service_version,
environment = %config.environment,
"Tracing initialized without OpenTelemetry export"
);
}
Ok(())
}
pub fn shutdown_telemetry() {
tracing::info!("Shutting down OpenTelemetry tracing");
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_default_config() {
let config = TelemetryConfig::default();
assert_eq!(config.service_name, "rs3gw");
assert_eq!(config.sampling_ratio, 1.0);
}
#[test]
fn test_custom_config() {
let config = TelemetryConfig {
service_name: "test-service".to_string(),
service_version: "1.0.0".to_string(),
environment: "test".to_string(),
otlp_endpoint: Some("http://localhost:4317".to_string()),
sampling_ratio: 0.5,
enable_export: true,
};
assert_eq!(config.service_name, "test-service");
assert_eq!(config.service_version, "1.0.0");
assert_eq!(config.environment, "test");
assert_eq!(config.sampling_ratio, 0.5);
assert!(config.enable_export);
}
#[test]
fn test_init_telemetry() {
let config = TelemetryConfig::default();
let result = init_telemetry(config);
assert!(result.is_ok());
}
}