qrt_log_utils/

lib.rs

use opentelemetry::{global, KeyValue};
use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge;
use opentelemetry_otlp::{LogExporter, MetricExporter, SpanExporter, WithExportConfig};
use opentelemetry_sdk::logs::SdkLoggerProvider;
use opentelemetry_sdk::metrics::{PeriodicReader, SdkMeterProvider};
use opentelemetry_sdk::trace::SdkTracerProvider;
use opentelemetry_sdk::Resource;
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use tracing_subscriber::filter::filter_fn;
use tracing_subscriber::prelude::*;
use tracing_subscriber::EnvFilter;

// pub use for client
pub use opentelemetry;
use opentelemetry::trace::TracerProvider;
pub use tracing;

const DEFAULT_ENDPOINT: &'static str = "http://localhost:4317";

pub struct LoggerContext {
    pub logger_provider: SdkLoggerProvider,
    pub tracer_provider: SdkTracerProvider,
    pub meter_provider: SdkMeterProvider,
}

impl LoggerContext {
    pub fn shudown(self) {
        let _ = self.meter_provider.shutdown();
        let _ = self.tracer_provider.shutdown();
        let _ = self.logger_provider.shutdown();
    }
}

fn get_resource(service_name: &'static str) -> Resource {
    static RESOURCE: OnceLock<Resource> = OnceLock::new();
    RESOURCE
        .get_or_init(|| {
            let mut builder = Resource::builder().with_service_name(service_name);

            #[cfg(feature = "detect-host")]
            {
                builder =
                    builder.with_attribute(gen_host_info().expect("failed to build host info"));
            }

            builder.build()
        })
        .clone()
}
fn init_traces(endpoint: &str, service_name: &'static str) -> SdkTracerProvider {
    let exporter = SpanExporter::builder()
        .with_tonic()
        .with_endpoint(endpoint)
        .build()
        .expect("Failed to create span exporter");
    SdkTracerProvider::builder()
        .with_resource(get_resource(service_name))
        .with_batch_exporter(exporter)
        .build()
}

fn init_metrics(endpoint: &str, service_name: &'static str) -> SdkMeterProvider {
    let exporter = MetricExporter::builder()
        .with_tonic()
        .with_endpoint(endpoint)
        .build()
        .expect("Failed to create metric exporter");

    let reader = PeriodicReader::builder(exporter)
        .with_interval(Duration::from_secs(1))
        .build();
    SdkMeterProvider::builder()
        .with_reader(reader)
        .with_resource(get_resource(service_name))
        .build()
}

fn init_logs(endpoint: &str, service_name: &'static str) -> SdkLoggerProvider {
    let exporter = LogExporter::builder()
        .with_tonic()
        .with_endpoint(endpoint)
        .build()
        .expect("Failed to create log exporter");

    SdkLoggerProvider::builder()
        .with_resource(get_resource(service_name))
        .with_batch_exporter(exporter)
        .build()
}

fn target_in_whitelist(target: &str, whitelist: &[String]) -> bool {
    whitelist.iter().any(|crate_name| {
        if target == crate_name {
            true
        } else if target.starts_with(crate_name) {
            target[crate_name.len()..].starts_with("::")
        } else {
            false
        }
    })
}

fn init_tracer(
    application_name: &'static str,
    endpoint: &str,
    provider: SdkTracerProvider,
    crate_whitelist: Vec<String>,
) -> SdkLoggerProvider {
    let logger_provider = init_logs(endpoint, application_name);

    // Create a new OpenTelemetryTracingBridge using the above LoggerProvider.
    let otel_layer = OpenTelemetryTracingBridge::new(&logger_provider);

    // For the OpenTelemetry layer, add a tracing filter to filter events from
    // OpenTelemetry and its dependent crates (opentelemetry-otlp uses crates
    // like reqwest/tonic etc.) from being sent back to OTel itself, thus
    // preventing infinite telemetry generation. The filter levels are set as
    // follows:
    // - Allow `info` level and above by default.
    // - Restrict `opentelemetry`, `hyper`, `tonic`, and `reqwest` completely.
    // Note: This will also drop events from crates like `tonic` etc. even when
    // they are used outside the OTLP Exporter. For more details, see:
    // https://github.com/open-telemetry/opentelemetry-rust/issues/761
    let filter_otel = EnvFilter::new("info")
        .add_directive("hyper=off".parse().unwrap())
        .add_directive("opentelemetry=off".parse().unwrap())
        .add_directive("tonic=off".parse().unwrap())
        .add_directive("h2=off".parse().unwrap());
    let otel_layer = otel_layer.with_filter(filter_otel);

    // Create a new tracing::Fmt layer to print the logs to stdout. It has a
    // default filter of `info` level and above, and `debug` and above for logs
    // from OpenTelemetry crates. The filter levels can be customized as needed.
    let filter = EnvFilter::from_default_env()
        .add_directive(format!("{}=info", &application_name).parse().unwrap());
    println!(
        "filter={} crate_whitelist={}",
        &filter,
        &crate_whitelist.join(",")
    );
    let tracer = provider.tracer("trace_demo");
    let telemetry = tracing_opentelemetry::layer().with_tracer(tracer);
    let fmt_layer = tracing_subscriber::fmt::layer()
        .with_thread_names(true)
        .with_filter(filter);

    let crate_whitelist: Vec<String> = crate_whitelist
        .into_iter()
        .filter_map(|entry| {
            let trimmed = entry.trim();
            if trimmed.is_empty() {
                None
            } else {
                Some(trimmed.to_string())
            }
        })
        .collect();
    let whitelist = Arc::new(crate_whitelist);
    let whitelist_active = !whitelist.is_empty();
    let whitelist_filter = filter_fn({
        let whitelist = Arc::clone(&whitelist);
        move |metadata| {
            if whitelist_active {
                target_in_whitelist(metadata.target(), whitelist.as_slice())
            } else {
                true
            }
        }
    });
    let otel_layer = otel_layer.with_filter(whitelist_filter.clone());
    let telemetry = telemetry.with_filter(whitelist_filter.clone());
    let fmt_layer = fmt_layer.with_filter(whitelist_filter);

    // Initialize the tracing subscriber with the OpenTelemetry layer and the
    // Fmt layer.
    tracing_subscriber::registry()
        .with(otel_layer)
        .with(telemetry)
        .with(fmt_layer)
        .init();
    logger_provider
}

#[derive(Debug, Default)]
pub struct LoggerConfig {
    pub endpoint: Option<String>,
    /// When non-empty, only logs/traces originating from the listed crates are emitted.
    pub crate_whitelist: Vec<String>,
}

impl LoggerConfig {
    /// Start building a logger configuration with fluent setters.
    pub fn builder() -> LoggerConfigBuilder {
        LoggerConfigBuilder::default()
    }
}

#[derive(Debug, Default)]
pub struct LoggerConfigBuilder {
    endpoint: Option<String>,
    crate_whitelist: Vec<String>,
}

impl LoggerConfigBuilder {
    /// Override the OTLP endpoint (defaults to `http://localhost:4317`).
    #[must_use]
    pub fn endpoint(mut self, endpoint: impl Into<String>) -> Self {
        self.endpoint = Some(endpoint.into());
        self
    }

    /// Add a crate to the whitelist (ignored if blank).
    #[must_use]
    pub fn add_whitelist_crate(mut self, crate_name: impl Into<String>) -> Self {
        let name = crate_name.into();
        if !name.trim().is_empty() {
            self.crate_whitelist.push(name);
        }
        self
    }

    /// Extend the whitelist with multiple crate names.
    #[must_use]
    pub fn add_whitelist_crates<I, S>(mut self, crates: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        for entry in crates {
            self = self.add_whitelist_crate(entry);
        }
        self
    }

    /// Build the final configuration.
    pub fn build(self) -> LoggerConfig {
        LoggerConfig {
            endpoint: self.endpoint,
            crate_whitelist: self.crate_whitelist,
        }
    }
}

pub fn init_logger(service_name: &'static str, log_config: LoggerConfig) -> LoggerContext {
    let LoggerConfig {
        endpoint,
        crate_whitelist,
    } = log_config;
    let endpoint = endpoint
        .as_ref()
        .map(|v| v.as_str())
        .unwrap_or(DEFAULT_ENDPOINT);
    // default to allow log from current application
    let tracer_provider = init_traces(endpoint, service_name);
    let logger_provider = init_tracer(
        service_name,
        endpoint,
        tracer_provider.clone(),
        crate_whitelist,
    );
    global::set_tracer_provider(tracer_provider.clone());
    let meter_provider = init_metrics(endpoint, service_name);
    global::set_meter_provider(meter_provider.clone());
    LoggerContext {
        logger_provider,
        tracer_provider,
        meter_provider,
    }
}

#[cfg(feature = "detect-host")]
fn gen_host_info() -> Result<KeyValue, std::io::Error> {
    hostname::get().map(|val| {
        KeyValue::new(
            "host.name".to_string(),
            val.into_string().unwrap_or_else(|_| "unknown".into()),
        )
    })
}

#[cfg(test)]
mod test {
    use std::time::Duration;

    use opentelemetry::{global, KeyValue};
    use tracing::info;
    use tracing::instrument;

    use crate::{init_logger, LoggerConfig};

    #[instrument]
    fn foo() {
        // let trace = register_dist_tracing_root(TraceId::default(), None);
        // println!("trace value: {:?}", trace);
        info!("test");
        bar();
    }

    #[instrument]
    fn bar() {
        // let trace = register_dist_tracing_root(TraceId::default(), None);
        // println!("trace value: {:?}", trace);
        info!("test2");
    }

    #[tokio::test]
    async fn test_logger() {
        init_logger("test_logger", LoggerConfig::default());
        for _ in 0..100 {
            // let span = span!(Level::Info, "my_span");
            // let _guard = span.enter();
            foo();
        }
        tokio::time::sleep(Duration::from_secs(4)).await;
    }

    #[instrument]
    fn add_counter() {
        let counter = global::meter("aaa").f64_counter("testCounterF64").build();
        counter.add(10f64, &[KeyValue::new("rate", "standard")]);
    }

    #[test]
    fn target_whitelist_matching() {
        let whitelist = vec!["crate_a".to_string(), "crate_b".to_string()];
        assert!(super::target_in_whitelist("crate_a", &whitelist));
        assert!(super::target_in_whitelist(
            "crate_a::module::sub",
            &whitelist
        ));
        assert!(super::target_in_whitelist("crate_b::something", &whitelist));
        assert!(!super::target_in_whitelist("crate", &whitelist));
        assert!(!super::target_in_whitelist("crate_c::foo", &whitelist));
    }
}