imks/telemetry/logs.rs

//! Log export: JSON console output + OpenTelemetry log bridge (OTLP).

use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge;
use opentelemetry_otlp::{LogExporter, Protocol, WithExportConfig};
use opentelemetry_sdk::logs::SdkLoggerProvider;
use opentelemetry_sdk::Resource;
use tracing_subscriber::fmt::format::FmtSpan;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::EnvFilter;
use tracing_subscriber::Registry;

use super::config::{OtlpProtocol, TelemetryConfig};
use crate::ImksResult;

/// Initialize the tracing subscriber.
///
/// Layer order (critical for OpenTelemetry compatibility):
/// 1. Registry
/// 2. OpenTelemetry trace layer (must be first — needs LookupSpan)
/// 3. EnvFilter
/// 4. Console formatting layer (JSON)
/// 5. OpenTelemetry log bridge
///
/// Returns the SdkLoggerProvider for graceful shutdown.
pub fn init_subscriber(
    config: &TelemetryConfig,
    resource: Option<&Resource>,
    otel_trace_layer: Option<
        tracing_opentelemetry::OpenTelemetryLayer<Registry, opentelemetry_sdk::trace::Tracer>,
    >,
) -> ImksResult<SdkLoggerProvider> {
    let env_filter =
        EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(&config.log_level));

    let (logger_provider, log_bridge_layer) = if config.logs_enabled {
        let exporter = build_log_exporter(config)?;

        let resource = resource.cloned().unwrap_or_else(|| Resource::builder().build());

        let provider = SdkLoggerProvider::builder()
            .with_resource(resource)
            .with_batch_exporter(exporter)
            .build();

        let bridge = OpenTelemetryTracingBridge::new(&provider);
        (Some(provider), Some(bridge))
    } else {
        (None, None)
    };

    match (otel_trace_layer, log_bridge_layer) {
        (Some(trace_layer), Some(log_layer)) => {
            let subscriber = Registry::default()
                .with(trace_layer)
                .with(env_filter)
                .with(make_json_fmt())
                .with(log_layer);
            set_subscriber(subscriber);
        }
        (Some(trace_layer), None) => {
            let subscriber = Registry::default()
                .with(trace_layer)
                .with(env_filter)
                .with(make_json_fmt());
            set_subscriber(subscriber);
        }
        (None, Some(log_layer)) => {
            let subscriber = Registry::default()
                .with(env_filter)
                .with(make_json_fmt())
                .with(log_layer);
            set_subscriber(subscriber);
        }
        (None, None) => {
            let subscriber = Registry::default()
                .with(env_filter)
                .with(make_json_fmt());
            set_subscriber(subscriber);
        }
    }

    let logger_provider = logger_provider.unwrap_or_else(|| SdkLoggerProvider::builder().build());

    Ok(logger_provider)
}

/// Create the JSON fmt layer with span context.
fn make_json_fmt<S>() -> tracing_subscriber::fmt::Layer<
    S,
    tracing_subscriber::fmt::format::JsonFields,
    tracing_subscriber::fmt::format::Format<tracing_subscriber::fmt::format::Json>,
>
where
    S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
{
    tracing_subscriber::fmt::layer()
        .json()
        .with_span_events(FmtSpan::CLOSE)
        .with_current_span(true)
        .with_span_list(true)
}

fn set_subscriber<S>(subscriber: S)
where
    S: tracing::Subscriber + Send + Sync + 'static,
{
    match tracing::subscriber::set_global_default(subscriber) {
        Ok(()) => {}
        Err(e) => {
            tracing::warn!("Could not set global tracing subscriber: {e}");
        }
    }
}

fn build_log_exporter(config: &TelemetryConfig) -> ImksResult<LogExporter> {
    match config.otlp_protocol {
        OtlpProtocol::Grpc => LogExporter::builder()
            .with_tonic()
            .with_endpoint(&config.otlp_endpoint)
            .build()
            .map_err(|e| crate::ImksError::Internal(format!("OTLP gRPC log exporter: {e}"))),
        OtlpProtocol::HttpProtobuf => LogExporter::builder()
            .with_http()
            .with_protocol(Protocol::HttpBinary)
            .with_endpoint(&config.otlp_endpoint)
            .build()
            .map_err(|e| crate::ImksError::Internal(format!("OTLP HTTP log exporter: {e}"))),
    }
}