pict-rs/src/init_tracing.rs

use crate::config::{LogFormat, OpenTelemetry, Tracing};
use color_eyre::config::Theme;
use console_subscriber::ConsoleLayer;
use opentelemetry::{trace::TracerProvider, KeyValue};
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::{propagation::TraceContextPropagator, Resource};
use tracing::subscriber::set_global_default;
use tracing_error::ErrorLayer;
use tracing_log::LogTracer;
use tracing_subscriber::{
    fmt::format::FmtSpan, layer::SubscriberExt, registry::LookupSpan, Layer, Registry,
};

pub(super) fn init_tracing(tracing: &Tracing) -> color_eyre::Result<()> {
    let eyre_theme = if tracing.logging.no_ansi {
        Theme::new()
    } else {
        Theme::dark()
    };

    color_eyre::config::HookBuilder::new()
        .theme(eyre_theme)
        .install()?;

    LogTracer::init()?;

    opentelemetry::global::set_text_map_propagator(TraceContextPropagator::new());

    let fmt_span = if tracing.logging.log_spans {
        FmtSpan::NEW | FmtSpan::CLOSE
    } else {
        FmtSpan::NONE
    };

    let format_layer = tracing_subscriber::fmt::layer()
        .with_span_events(fmt_span)
        .with_ansi(!tracing.logging.no_ansi);

    match tracing.logging.format {
        LogFormat::Compact => with_format(format_layer.compact(), tracing),
        LogFormat::Json => with_format(format_layer.json(), tracing),
        LogFormat::Normal => with_format(format_layer, tracing),
        LogFormat::Pretty => with_format(format_layer.pretty(), tracing),
    }
}

fn with_format<F>(format_layer: F, tracing: &Tracing) -> color_eyre::Result<()>
where
    F: Layer<Registry> + Send + Sync,
{
    let format_layer = format_layer.with_filter(tracing.logging.targets.targets.clone());

    let subscriber = Registry::default()
        .with(format_layer)
        .with(ErrorLayer::default());

    if let Some(address) = tracing.console.address {
        let console_layer = ConsoleLayer::builder()
            .event_buffer_capacity(tracing.console.buffer_capacity)
            .server_addr(address)
            .spawn();

        let subscriber = subscriber.with(console_layer);

        with_subscriber(subscriber, &tracing.opentelemetry)?;
        tracing::info!("Starting console on {address}");
        Ok(())
    } else {
        with_subscriber(subscriber, &tracing.opentelemetry)
    }
}

fn with_subscriber<S>(subscriber: S, otel: &OpenTelemetry) -> color_eyre::Result<()>
where
    S: SubscriberExt + Send + Sync,
    for<'a> S: LookupSpan<'a>,
{
    if let Some(url) = otel.url.as_ref() {
        let tracer = opentelemetry_otlp::new_pipeline()
            .tracing()
            .with_trace_config(opentelemetry_sdk::trace::Config::default().with_resource(
                Resource::new(vec![KeyValue::new(
                    "service.name",
                    otel.service_name.clone(),
                )]),
            ))
            .with_exporter(
                opentelemetry_otlp::new_exporter()
                    .tonic()
                    .with_endpoint(url.as_str()),
            )
            .install_batch(opentelemetry_sdk::runtime::Tokio)?;

        let otel_layer = tracing_opentelemetry::layer()
            .with_tracer(tracer.tracer("pict-rs-tracer"))
            .with_filter(otel.targets.as_ref().targets.clone());

        let subscriber = subscriber.with(otel_layer);

        set_global_default(subscriber)?;
    } else {
        set_global_default(subscriber)?;
    }

    Ok(())
}