patagia-control/instrumentation/src/lib.rs

use anyhow::{anyhow, Result};
use once_cell::sync::Lazy;
use opentelemetry::{trace::TracerProvider as _, KeyValue};
use opentelemetry_otlp::{WithExportConfig, WithTonicConfig};
use opentelemetry_sdk::{
    metrics::{MeterProviderBuilder, PeriodicReader, SdkMeterProvider},
    propagation::TraceContextPropagator,
    runtime,
    trace::{RandomIdGenerator, Sampler, TracerProvider},
    Resource,
};
use tracing_opentelemetry::{MetricsLayer, OpenTelemetryLayer};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;

static RESOURCE: Lazy<Resource> = Lazy::new(|| {
    Resource::new(vec![
        KeyValue::new(
            opentelemetry_semantic_conventions::resource::SERVICE_NAME,
            env!("CARGO_PKG_NAME"),
        ),
        KeyValue::new(
            opentelemetry_semantic_conventions::resource::SERVICE_VERSION,
            env!("CARGO_PKG_VERSION"),
        ),
    ])
});

// Construct MeterProvider for MetricsLayer
fn init_meter_provider(otel_endpoint: &String) -> Result<SdkMeterProvider> {
    let exporter = opentelemetry_otlp::MetricExporter::builder()
        .with_tonic()
        .with_endpoint(otel_endpoint)
        .with_tls_config(tonic::transport::ClientTlsConfig::new().with_native_roots())
        .with_compression(opentelemetry_otlp::Compression::Gzip)
        .with_temporality(opentelemetry_sdk::metrics::Temporality::default())
        .build()
        .map_err(|e| anyhow!("Error creating OTLP metric exporter: {:?}", e))?;

    let meter_provider = MeterProviderBuilder::default()
        .with_resource(RESOURCE.clone())
        .with_reader(
            PeriodicReader::builder(exporter, runtime::Tokio)
                .with_interval(std::time::Duration::from_secs(10))
                .build(),
        )
        .build();

    opentelemetry::global::set_meter_provider(meter_provider.clone());

    Ok(meter_provider)
}

// Construct TracerProvider for OpenTelemetryLayer
fn init_tracer_provider(otel_endpoint: &String) -> Result<TracerProvider> {
    let exporter = opentelemetry_otlp::SpanExporter::builder()
        .with_tonic()
        .with_tls_config(tonic::transport::ClientTlsConfig::new().with_native_roots())
        .with_compression(opentelemetry_otlp::Compression::Gzip)
        .with_endpoint(otel_endpoint)
        .build()
        .map_err(|e| anyhow!("Error creating OTLP span exporter: {:?}", e))?;

    let tracer_provider = opentelemetry_sdk::trace::TracerProvider::builder()
        .with_sampler(Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(
            1.0,
        ))))
        .with_resource(RESOURCE.clone())
        .with_id_generator(RandomIdGenerator::default())
        .with_batch_exporter(exporter, opentelemetry_sdk::runtime::Tokio)
        .build();

    Ok(tracer_provider)
}

// Initialize tracing-subscriber and return TracingGuard for opentelemetry-related termination processing
pub fn init_tracing(otel_endpoint: Option<&String>) -> Result<TracingGuard> {
    let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(""));

    let fmt_layer = tracing_subscriber::fmt::layer().with_writer(std::io::stderr);

    let meter_provider = match otel_endpoint {
        Some(endpoint) => Some(init_meter_provider(endpoint)?),
        None => None,
    };

    let metrics_layer = match meter_provider {
        Some(ref p) => Some(MetricsLayer::new(p.to_owned())),
        None => None,
    };

    let tracer_provider = match otel_endpoint {
        Some(endpoint) => Some(init_tracer_provider(endpoint)?),
        None => None,
    };

    let trace_layer = match tracer_provider {
        Some(ref p) => Some(OpenTelemetryLayer::new(p.tracer("tracing-otel-subscriber"))),
        None => None,
    };

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

    tracing_subscriber::registry()
        .with(env_filter)
        .with(fmt_layer)
        .with(metrics_layer)
        .with(trace_layer)
        .init();

    Ok(TracingGuard {
        meter_provider,
        tracer_provider,
    })
}

pub struct TracingGuard {
    meter_provider: Option<SdkMeterProvider>,
    tracer_provider: Option<TracerProvider>,
}

impl Drop for TracingGuard {
    fn drop(&mut self) {
        if let Some(tracer_provider) = &self.tracer_provider {
            for result in tracer_provider.force_flush() {
                if let Err(err) = result {
                    eprintln!("{err:?}");
                }
            }
            if let Err(err) = tracer_provider.shutdown() {
                eprintln!("{err:?}");
            }
        }

        if let Some(meter_provider) = &self.meter_provider {
            if let Err(err) = meter_provider.force_flush() {
                eprintln!("{err:?}");
            }
            if let Err(err) = meter_provider.shutdown() {
                eprintln!("{err:?}");
            }
        }
    }
}