pub mod metrics;
pub use metrics::MetricsCollector;
use crate::error::KovaError;
#[derive(Debug, Clone, PartialEq)]
pub enum OtlpProtocol {
Grpc,
Http,
}
#[derive(Debug, Clone, PartialEq)]
pub enum ExporterConfig {
Stdout,
Otlp {
endpoint: String,
protocol: OtlpProtocol,
},
Jaeger { endpoint: String },
}
#[derive(Debug, Clone)]
pub struct TelemetryConfig {
pub log_level: tracing::Level,
pub exporter: ExporterConfig,
pub sampling_rate: f64,
}
impl Default for TelemetryConfig {
fn default() -> Self {
Self {
log_level: tracing::Level::INFO,
exporter: ExporterConfig::Stdout,
sampling_rate: 1.0,
}
}
}
impl TelemetryConfig {
pub fn builder() -> TelemetryConfigBuilder {
TelemetryConfigBuilder::default()
}
pub fn init(&self) -> Result<(), KovaError> {
#[cfg(feature = "telemetry")]
{
self.init_with_otel()
}
#[cfg(not(feature = "telemetry"))]
{
self.init_basic()
}
}
#[cfg(not(feature = "telemetry"))]
fn init_basic(&self) -> Result<(), KovaError> {
use tracing_subscriber::{EnvFilter, fmt};
let filter = EnvFilter::new(self.log_level.to_string());
fmt()
.with_env_filter(filter)
.try_init()
.map_err(|e| KovaError::Build(format!("Failed to init tracing subscriber: {e}")))?;
Ok(())
}
#[cfg(feature = "telemetry")]
fn init_with_otel(&self) -> Result<(), KovaError> {
use opentelemetry::trace::TracerProvider as _;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::trace::{Sampler, SdkTracerProvider};
use tracing_opentelemetry::OpenTelemetryLayer;
use tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt};
let sampler = if (self.sampling_rate - 1.0_f64).abs() < f64::EPSILON {
Sampler::AlwaysOn
} else if self.sampling_rate <= 0.0 {
Sampler::AlwaysOff
} else {
Sampler::TraceIdRatioBased(self.sampling_rate)
};
let provider = match &self.exporter {
ExporterConfig::Stdout => {
let exporter = opentelemetry_stdout::SpanExporter::default();
SdkTracerProvider::builder()
.with_batch_exporter(exporter)
.with_sampler(sampler)
.build()
}
ExporterConfig::Otlp { endpoint, protocol } => {
let exporter = match protocol {
OtlpProtocol::Grpc => opentelemetry_otlp::SpanExporter::builder()
.with_tonic()
.with_endpoint(endpoint)
.build()
.map_err(|e| KovaError::Build(format!("OTLP gRPC exporter error: {e}")))?,
OtlpProtocol::Http => opentelemetry_otlp::SpanExporter::builder()
.with_http()
.with_endpoint(endpoint)
.build()
.map_err(|e| KovaError::Build(format!("OTLP HTTP exporter error: {e}")))?,
};
SdkTracerProvider::builder()
.with_batch_exporter(exporter)
.with_sampler(sampler)
.build()
}
ExporterConfig::Jaeger { endpoint } => {
let exporter = opentelemetry_otlp::SpanExporter::builder()
.with_tonic()
.with_endpoint(endpoint)
.build()
.map_err(|e| KovaError::Build(format!("Jaeger exporter error: {e}")))?;
SdkTracerProvider::builder()
.with_batch_exporter(exporter)
.with_sampler(sampler)
.build()
}
};
let tracer = provider.tracer("kova");
let otel_layer = OpenTelemetryLayer::new(tracer);
let filter = EnvFilter::new(self.log_level.to_string());
tracing_subscriber::registry()
.with(filter)
.with(fmt::layer())
.with(otel_layer)
.try_init()
.map_err(|e| KovaError::Build(format!("Failed to init tracing subscriber: {e}")))?;
Ok(())
}
}
#[derive(Debug, Clone, Default)]
pub struct TelemetryConfigBuilder {
config: TelemetryConfig,
}
impl TelemetryConfigBuilder {
pub fn log_level(mut self, level: tracing::Level) -> Self {
self.config.log_level = level;
self
}
pub fn exporter(mut self, exporter: ExporterConfig) -> Self {
self.config.exporter = exporter;
self
}
pub fn sampling_rate(mut self, rate: f64) -> Self {
self.config.sampling_rate = rate.clamp(0.0, 1.0);
self
}
pub fn build(self) -> TelemetryConfig {
self.config
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_config_values() {
let cfg = TelemetryConfig::default();
assert_eq!(cfg.log_level, tracing::Level::INFO);
assert_eq!(cfg.exporter, ExporterConfig::Stdout);
assert!((cfg.sampling_rate - 1.0).abs() < f64::EPSILON);
}
#[test]
fn builder_sets_all_fields() {
let cfg = TelemetryConfig::builder()
.log_level(tracing::Level::DEBUG)
.exporter(ExporterConfig::Otlp {
endpoint: "http://localhost:4317".into(),
protocol: OtlpProtocol::Grpc,
})
.sampling_rate(0.5)
.build();
assert_eq!(cfg.log_level, tracing::Level::DEBUG);
assert_eq!(
cfg.exporter,
ExporterConfig::Otlp {
endpoint: "http://localhost:4317".into(),
protocol: OtlpProtocol::Grpc,
}
);
assert!((cfg.sampling_rate - 0.5).abs() < f64::EPSILON);
}
#[test]
fn builder_clamps_sampling_rate() {
let cfg = TelemetryConfig::builder().sampling_rate(2.0).build();
assert!((cfg.sampling_rate - 1.0).abs() < f64::EPSILON);
let cfg = TelemetryConfig::builder().sampling_rate(-0.5).build();
assert!(cfg.sampling_rate.abs() < f64::EPSILON);
}
#[test]
fn exporter_config_variants() {
let stdout = ExporterConfig::Stdout;
let otlp_grpc = ExporterConfig::Otlp {
endpoint: "http://localhost:4317".into(),
protocol: OtlpProtocol::Grpc,
};
let otlp_http = ExporterConfig::Otlp {
endpoint: "http://localhost:4318".into(),
protocol: OtlpProtocol::Http,
};
let jaeger = ExporterConfig::Jaeger {
endpoint: "http://localhost:14250".into(),
};
assert!(!format!("{stdout:?}").is_empty());
assert!(!format!("{otlp_grpc:?}").is_empty());
assert!(!format!("{otlp_http:?}").is_empty());
assert!(!format!("{jaeger:?}").is_empty());
}
#[test]
fn builder_jaeger_exporter() {
let cfg = TelemetryConfig::builder()
.log_level(tracing::Level::WARN)
.exporter(ExporterConfig::Jaeger {
endpoint: "http://jaeger:14250".into(),
})
.build();
assert_eq!(cfg.log_level, tracing::Level::WARN);
assert_eq!(
cfg.exporter,
ExporterConfig::Jaeger {
endpoint: "http://jaeger:14250".into(),
}
);
}
#[test]
fn builder_all_log_levels() {
for level in [
tracing::Level::TRACE,
tracing::Level::DEBUG,
tracing::Level::INFO,
tracing::Level::WARN,
tracing::Level::ERROR,
] {
let cfg = TelemetryConfig::builder().log_level(level).build();
assert_eq!(cfg.log_level, level);
}
}
}