pi_logger 0.8.0

use opentelemetry::{
    global::{self, BoxedTracer},
    propagation::TextMapCompositePropagator,
    trace::{TraceContextExt, Tracer, TracerProvider},
    InstrumentationScope, KeyValue,
};
use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge;
use opentelemetry_otlp::{
    tonic_types::metadata::MetadataMap, LogExporter, MetricExporter, Protocol, SpanExporter,
};
use opentelemetry_otlp::{WithExportConfig, WithTonicConfig};
use opentelemetry_sdk::{
    logs::SdkLoggerProvider,
    metrics::SdkMeterProvider,
    propagation::{BaggagePropagator, TraceContextPropagator},
    trace::{Config, SdkTracerProvider},
};
use opentelemetry_sdk::{metrics::Temporality, Resource};
use parking_lot::RwLock;
use std::{
    error::Error,
    str::FromStr,
    sync::{Arc, OnceLock},
    time::Duration,
};
use tracing::{info, Subscriber};
use tracing_opentelemetry::OpenTelemetryLayer;
use tracing_subscriber::EnvFilter;
use tracing_subscriber::{
    filter::{LevelFilter, Targets},
    reload,
};
use tracing_subscriber::{prelude::*, Registry};

lazy_static! {
    // 动态调整trace日志等级
    static ref TRACE_RELOAD_HANDLE: Arc<RwLock<Option<reload::Handle<EnvFilter, tracing_subscriber::Registry>>>> = Arc::new(RwLock::new(None));
}

pub struct OpenTelemetryMgr {
    pub tracer: Option<SdkTracerProvider>,
    pub meter_provider: Option<SdkMeterProvider>,
    pub logger_provider: Option<SdkLoggerProvider>,
    pub server_name: String,
    pub version: String,
    pub endpoint: String,
}

impl OpenTelemetryMgr {
    pub fn new(server_name: &str, version: &str, endpoint: &str) -> Self {
        OpenTelemetryMgr {
            tracer: None,
            meter_provider: None,
            logger_provider: None,
            server_name: server_name.to_string(),
            version: version.to_string(),
            endpoint: endpoint.to_string(),
        }
    }

    pub fn init_logs(&mut self) {
        let exporter = LogExporter::builder()
            .with_http()
            .with_protocol(Protocol::HttpBinary)
            .with_endpoint(self.endpoint.clone() + "/v1/logs")
            .build()
            .expect("Failed to create log exporter");

        let logger_provider = SdkLoggerProvider::builder()
            .with_batch_exporter(exporter)
            .with_resource(get_resource(&self.server_name, &self.version))
            .build();
        self.logger_provider = Some(logger_provider);
    }

    pub fn init_tracing(&mut self) {
        // 创建一个Baggage传播器
        let baggage_propagator = BaggagePropagator::new();
        // 创建一个TraceContext传播器
        let trace_context_propagator = TraceContextPropagator::new();
        // 创建一个TextMap复合传播器，包含Baggage和TraceContext传播器
        let composite_propagator = TextMapCompositePropagator::new(vec![
            Box::new(baggage_propagator),
            Box::new(trace_context_propagator),
        ]);

        // 设置全局的TextMap传播器
        global::set_text_map_propagator(composite_propagator);

        // 创建一个SpanExporter，使用HTTP协议
        let exporter = SpanExporter::builder()
            .with_http()
            .with_protocol(Protocol::HttpBinary) //can be changed to `Protocol::HttpJson` to export in JSON format
            .with_endpoint(self.endpoint.clone() + "/v1/traces")
            .with_timeout(Duration::from_secs(3))
            .build()
            .expect("Failed to create trace exporter");

        let provider = SdkTracerProvider::builder()
            .with_batch_exporter(exporter)
            .with_resource(get_resource(&self.server_name, &self.version))
            .build();

        global::set_tracer_provider(provider.clone());
        global::tracer_provider().tracer("pi_logger");
        self.tracer = Some(provider);
    }

    pub fn init_metrics(&mut self) {
        let exporter = MetricExporter::builder()
            .with_http()
            .with_protocol(Protocol::HttpBinary) //can be changed to `Protocol::HttpJson` to export in JSON format
            .with_endpoint(self.endpoint.clone() + "/v1/metrics")
            .with_temporality(Temporality::LowMemory)
            .with_timeout(Duration::from_secs(3))
            .build()
            .expect("Failed to create metric exporter");

        let provider = SdkMeterProvider::builder()
            .with_periodic_exporter(exporter)
            .with_resource(get_resource(&self.server_name, &self.version))
            .build();

        global::set_meter_provider(provider.clone());
        self.meter_provider = Some(provider);
    }

    pub fn subscriber(&mut self) {
        let filter_otel = EnvFilter::new(std::env::var("RUST_LOG").unwrap_or("info".to_string()))
            .add_directive("hyper=off".parse().unwrap())
            .add_directive("tonic=off".parse().unwrap())
            .add_directive("h2=off".parse().unwrap())
            .add_directive("reqwest=off".parse().unwrap());
        let (filter, reload_handle) = reload::Layer::new(filter_otel);
        *TRACE_RELOAD_HANDLE.write() = Some(reload_handle);

        // tracing_subscriber::registry()
        //     .with(filter)
        //     .with(
        //         tracing_opentelemetry::layer()
        //             .with_tracer(self.tracer.clone().unwrap().tracer("pi_logger")),
        //     )
        //     .init();

        let subscriber = Registry::default() // 或者 registry()
            .with(filter)
            .with(
                tracing_opentelemetry::layer()
                    .with_tracer(self.tracer.clone().unwrap().tracer("pi_logger")),
            );
        tracing::subscriber::set_global_default(subscriber);
    }
}

fn get_resource(server_name: &str, version: &str) -> Resource {
    static RESOURCE: OnceLock<Resource> = OnceLock::new();
    RESOURCE
        .get_or_init(|| {
            Resource::builder()
                .with_service_name(server_name.to_string())
                .with_attribute(KeyValue::new("service.version", version.to_string()))
                .build()
        })
        .clone()
}

/// 动态调整trace日志等级
/// log: "info,mod=debug"
pub fn reload_trace_level(log: &str) -> Result<(), Box<dyn std::error::Error>> {
    // 打印日志信息
    println!("reload_trace_level:{:?}", log);
    // 从环境变量中获取日志级别
    let filter = EnvFilter::try_new(log)?;
    // 如果TRACE_RELOAD_HANDLE存在，则修改日志级别
    if let Some(handle) = TRACE_RELOAD_HANDLE.write().as_ref() {
        handle.modify(|filter_tmp| {
            *filter_tmp = filter;
        });
    };
    // 返回成功
    Ok(())
}

// 初始化OpenTelemetryMgr，传入服务器名称和版本号
pub fn init(server_name: &str, version: &str, endpoint: &str) -> Option<OpenTelemetryMgr> {
    // 通过TRACE_RELOAD_HANDLE判断是否初始化过
    if TRACE_RELOAD_HANDLE.read().is_none() {
        // 创建一个新的OpenTelemetryMgr实例
        let mut mgr = OpenTelemetryMgr::new(server_name, version, endpoint);
        // 初始化日志
        mgr.init_logs();
        // 初始化追踪
        mgr.init_tracing();
        // 初始化度量
        mgr.init_metrics();
        // 订阅
        mgr.subscriber();
        // 返回OpenTelemetryMgr实例
        Some(mgr)
    } else {
        None
    }
}

// 判断是否初始化过
pub fn is_init() -> bool {
    TRACE_RELOAD_HANDLE.read().is_some()
}

// 创建一个TextMapCompositePropagator，它包含一个BaggagePropagator和一个TraceContextPropagator
pub fn create_baggage() -> TextMapCompositePropagator {
    // 创建一个BaggagePropagator
    let baggage_propagator = BaggagePropagator::new();
    // 创建一个TraceContextPropagator
    let trace_context_propagator = TraceContextPropagator::new();

    // 创建一个TextMapCompositePropagator，并将BaggagePropagator和TraceContextPropagator添加到其中
    TextMapCompositePropagator::new(vec![
        Box::new(baggage_propagator),
        Box::new(trace_context_propagator),
    ])
}

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

    #[test]
    fn test_trace() {
        // 设置环境变量
        std::env::set_var("RUST_LOG", "info");
        let mgr = init("test", "1.0.0", "http://127.0.0.1:4318").unwrap();
        {
            let main_span = tracing::info_span!("Main operation");
            let _enter = main_span.enter();
            tracing::event!(name: "Nice operation!", tracing::Level::INFO, some.key = 100);
            info!(target: "my-target", "hello from {}. My price is {}. I am also inside a Span!", "banana", 2.99);

            let sub_span = tracing::info_span!("Sub span event");
            let _enter = sub_span.enter();
        }
        mgr.logger_provider.clone().unwrap().shutdown();
        mgr.meter_provider.clone().unwrap().shutdown();
        mgr.tracer.clone().unwrap().shutdown();
    }
}