use crate::{
exporter::{ExporterConfig, TraceExporter},
sampler::TraceSampler,
};
use opentelemetry::trace::TracerProvider;
use opentelemetry_sdk::trace::{SdkTracerProvider, TraceError};
use tracing_subscriber::{EnvFilter, Registry, layer::SubscriberExt, util::SubscriberInitExt};
#[derive(Debug, Clone)]
pub struct TracerConfig {
pub service_name: String,
pub service_version: String,
pub environment: String,
pub sampler: TraceSampler,
pub exporter_config: ExporterConfig,
pub enable_logging: bool,
pub log_filter: String,
}
impl Default for TracerConfig {
fn default() -> Self {
Self {
service_name: "revoke-service".to_string(),
service_version: env!("CARGO_PKG_VERSION").to_string(),
environment: "development".to_string(),
sampler: TraceSampler::always_on(),
exporter_config: ExporterConfig::default(),
enable_logging: true,
log_filter: "info".to_string(),
}
}
}
pub struct RevokeTracer {
provider: SdkTracerProvider,
config: TracerConfig,
}
impl RevokeTracer {
pub async fn new(config: TracerConfig) -> Result<Self, TraceError> {
let exporter = {
#[cfg(feature = "otlp")]
{
TraceExporter::otlp(config.exporter_config.clone())
}
#[cfg(not(feature = "otlp"))]
{
TraceExporter::console()
}
};
let _tracer = exporter.build(&config.service_name).await?;
use opentelemetry::KeyValue;
use opentelemetry_sdk::Resource;
let resource = Resource::builder()
.with_service_name(config.service_name.clone())
.with_attribute(KeyValue::new(
"service.version",
config.service_version.clone(),
))
.build();
let provider = SdkTracerProvider::builder().with_resource(resource).build();
Ok(Self { provider, config })
}
pub fn install_global(self) -> Result<(), TraceError> {
opentelemetry::global::set_tracer_provider(self.provider.clone());
if self.config.enable_logging {
self.setup_tracing_subscriber()?;
}
Ok(())
}
fn setup_tracing_subscriber(&self) -> Result<(), TraceError> {
let tracer = self.provider.tracer("revoke-trace");
let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer);
let fmt_layer = tracing_subscriber::fmt::layer()
.with_target(true)
.with_thread_ids(true)
.with_level(true);
let filter = EnvFilter::try_from_default_env()
.unwrap_or_else(|_| EnvFilter::new(&self.config.log_filter));
Registry::default()
.with(filter)
.with(fmt_layer)
.with(otel_layer)
.try_init()
.map_err(|e| TraceError::Other(Box::new(e)))?;
Ok(())
}
pub fn tracer(&self, name: &str) -> opentelemetry_sdk::trace::Tracer {
self.provider.tracer(name.to_string())
}
pub fn force_flush(&self) -> Result<(), TraceError> {
let _ = self.provider.force_flush();
Ok(())
}
pub fn shutdown(self) -> Result<(), TraceError> {
drop(self.provider);
Ok(())
}
}
pub struct TracerBuilder {
config: TracerConfig,
}
impl TracerBuilder {
pub fn new(service_name: impl Into<String>) -> Self {
let mut config = TracerConfig::default();
config.service_name = service_name.into();
Self { config }
}
pub fn with_version(mut self, version: impl Into<String>) -> Self {
self.config.service_version = version.into();
self
}
pub fn with_environment(mut self, env: impl Into<String>) -> Self {
self.config.environment = env.into();
self
}
pub fn with_sampler(mut self, sampler: TraceSampler) -> Self {
self.config.sampler = sampler;
self
}
pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
self.config.exporter_config.endpoint = endpoint.into();
self
}
pub fn with_logging(mut self, enable: bool) -> Self {
self.config.enable_logging = enable;
self
}
pub fn with_log_filter(mut self, filter: impl Into<String>) -> Self {
self.config.log_filter = filter.into();
self
}
pub async fn build(self) -> Result<RevokeTracer, TraceError> {
RevokeTracer::new(self.config).await
}
}
#[macro_export]
macro_rules! trace_span {
($name:expr) => {
$crate::span::SpanBuilder::new($name).start()
};
($name:expr, $($key:ident = $value:expr),*) => {
$crate::span::SpanBuilder::new($name)
$(.with_attribute(stringify!($key), $value))*
.start()
};
}
#[macro_export]
macro_rules! trace_error {
($span:expr, $error:expr) => {{
use $crate::span::SpanExt;
$span.record_exception(&$error);
$span.set_status($crate::span::SpanStatus::error($error.to_string()));
}};
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_tracer_builder() {
let tracer = TracerBuilder::new("test-service")
.with_version("1.0.0")
.with_environment("test")
.with_sampler(TraceSampler::always_off())
.with_endpoint("http://localhost:4317")
.with_logging(false)
.build()
.await;
assert!(tracer.is_ok());
let tracer = tracer.unwrap();
assert_eq!(tracer.config.service_name, "test-service");
assert_eq!(tracer.config.service_version, "1.0.0");
assert_eq!(tracer.config.environment, "test");
}
#[test]
fn test_trace_span_macro() {
let _span = tracing::info_span!("test_operation");
let _span_with_attrs =
tracing::info_span!("test_operation", method = "GET", path = "/api/test");
}
}