revoke_trace/

tracer.rs

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(),
        }
    }
}

/// Revoke 追踪器
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();

        // 创建 provider
        let provider = SdkTracerProvider::builder().with_resource(resource).build();

        Ok(Self { provider, config })
    }

    /// 安装为全局追踪器
    pub fn install_global(self) -> Result<(), TraceError> {
        // 设置全局 tracer provider
        opentelemetry::global::set_tracer_provider(self.provider.clone());

        // 如果启用了日志集成，设置 tracing subscriber
        if self.config.enable_logging {
            self.setup_tracing_subscriber()?;
        }

        Ok(())
    }

    /// 设置 tracing subscriber
    fn setup_tracing_subscriber(&self) -> Result<(), TraceError> {
        // 创建 OpenTelemetry 层
        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())
    }

    /// 强制刷新所有待处理的 spans
    pub fn force_flush(&self) -> Result<(), TraceError> {
        // 在新版本中 force_flush 不返回错误
        let _ = self.provider.force_flush();
        Ok(())
    }

    /// 关闭追踪器
    pub fn shutdown(self) -> Result<(), TraceError> {
        // 需要获取 TracerProvider 的所有权
        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
    }
}

/// 便捷的宏，用于创建 span
#[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() {
        // 测试 span 创建
        let _span = tracing::info_span!("test_operation");
        // 测试带属性的 span
        let _span_with_attrs =
            tracing::info_span!("test_operation", method = "GET", path = "/api/test");
        // 如果没有 panic，说明成功创建
    }
}