revoke-trace 0.3.0

Distributed tracing with OpenTelemetry for Revoke framework
Documentation
use opentelemetry_sdk::trace::TraceError;
use std::time::Duration;

/// 导出器配置
#[derive(Debug, Clone)]
pub struct ExporterConfig {
    /// 导出端点
    pub endpoint: String,
    /// 导出超时
    pub timeout: Duration,
    /// 是否使用 TLS
    pub use_tls: bool,
    /// 自定义头部
    pub headers: Vec<(String, String)>,
}

impl Default for ExporterConfig {
    fn default() -> Self {
        Self {
            endpoint: "http://localhost:4317".to_string(),
            timeout: Duration::from_secs(10),
            use_tls: false,
            headers: Vec::new(),
        }
    }
}

/// 追踪导出器
pub enum TraceExporter {
    #[cfg(feature = "otlp")]
    Otlp,
    /// 控制台导出器(用于调试)
    Console,
    /// 无操作导出器
    Noop,
}

impl TraceExporter {
    /// 创建 OTLP 导出器
    #[cfg(feature = "otlp")]
    pub fn otlp(_config: ExporterConfig) -> Self {
        Self::Otlp
    }


    /// 创建控制台导出器
    pub fn console() -> Self {
        Self::Console
    }

    /// 创建无操作导出器
    pub fn noop() -> Self {
        Self::Noop
    }

    /// 构建导出器
    pub async fn build(
        self,
        service_name: &str,
    ) -> Result<opentelemetry_sdk::trace::Tracer, TraceError> {
        use opentelemetry::KeyValue;
        use opentelemetry::trace::TracerProvider;
        use opentelemetry_sdk::Resource;

        let resource = Resource::builder()
            .with_service_name(service_name.to_string())
            .with_attribute(KeyValue::new("service.version", env!("CARGO_PKG_VERSION")))
            .build();

        match self {
            #[cfg(feature = "otlp")]
            Self::Otlp => {
                use opentelemetry_otlp::WithExportConfig;

                opentelemetry_otlp::SpanExporter::builder()
                    .with_tonic()
                    .with_endpoint("http://localhost:4317")
                    .build()
                    .map_err(|e| TraceError::Other(Box::new(e)))
                    .and_then(|exporter| {
                        let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
                            .with_batch_exporter(exporter)
                            .with_resource(resource)
                            .build();
                        Ok(provider.tracer(service_name.to_string()))
                    })
            }
            Self::Console => {
                // 使用简单的 TracerProvider
                let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
                    .with_resource(resource)
                    .build();

                Ok(provider.tracer(service_name.to_string()))
            }
            Self::Noop => {
                // 创建一个空的 provider
                let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
                    .with_resource(resource)
                    .build();
                Ok(provider.tracer(service_name.to_string()))
            }
        }
    }
}

/// 导出器构建器
pub struct ExporterBuilder {
    config: ExporterConfig,
}

impl ExporterBuilder {
    /// 创建新的构建器
    pub fn new() -> Self {
        Self {
            config: ExporterConfig::default(),
        }
    }

    /// 设置端点
    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
        self.config.endpoint = endpoint.into();
        self
    }

    /// 设置超时
    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.config.timeout = timeout;
        self
    }

    /// 启用 TLS
    pub fn with_tls(mut self, use_tls: bool) -> Self {
        self.config.use_tls = use_tls;
        self
    }

    /// 添加自定义头部
    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.config.headers.push((key.into(), value.into()));
        self
    }

    /// 构建 OTLP 导出器
    #[cfg(feature = "otlp")]
    pub fn build_otlp(self) -> TraceExporter {
        TraceExporter::otlp(self.config)
    }

    /// 构建控制台导出器
    pub fn build_console(self) -> TraceExporter {
        TraceExporter::console()
    }
}

impl Default for ExporterBuilder {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_exporter_config() {
        let config = ExporterConfig {
            endpoint: "http://localhost:4318".to_string(),
            timeout: Duration::from_secs(30),
            use_tls: true,
            ..Default::default()
        };

        assert_eq!(config.endpoint, "http://localhost:4318");
        assert_eq!(config.timeout, Duration::from_secs(30));
        assert!(config.use_tls);
    }

    #[test]
    fn test_exporter_builder() {
        let exporter = ExporterBuilder::new()
            .with_endpoint("http://custom:4317")
            .with_timeout(Duration::from_secs(20))
            .with_tls(true)
            .with_header("Authorization", "Bearer token")
            .build_console();

        match exporter {
            TraceExporter::Console => {
                // 正确构建了控制台导出器
            }
            _ => panic!("Expected console exporter"),
        }
    }
}