avantis-utils 0.7.2

Swiss army knife utilities for all kind of Avantis rust projects
Documentation
use gethostname::gethostname;
use opentelemetry::global::set_text_map_propagator;
use opentelemetry::sdk::propagation::TraceContextPropagator;
use opentelemetry::sdk::trace;
use opentelemetry::sdk::Resource;
use opentelemetry::KeyValue;
use opentelemetry_otlp::WithExportConfig;
use serde::Deserialize;
use std::collections::HashSet;
use thiserror::Error;
use tracing::info;
use tracing::subscriber::set_global_default;
use tracing::Subscriber;
use tracing_bunyan_formatter::{BunyanFormattingLayer, JsonStorageLayer};
use tracing_log::LogTracer;
use tracing_subscriber::filter::FilterFn;
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::Layer;
use tracing_subscriber::{EnvFilter, Registry};

#[derive(Deserialize, Clone, PartialEq, Eq, Debug)]
pub struct TelemetrySetting {
    pub otel_collector_endpoint: String,
    pub disabled_targets: HashSet<String>,
    pub log_level: String,
}

impl TelemetrySetting {
    fn log_level_filter<S>(&self) -> impl Layer<S>
    where
        S: Subscriber,
    {
        EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(self.log_level.clone()))
    }

    fn bunyan_formatter<S>(&self, service_name: &'static str) -> impl Layer<S>
    where
        S: Subscriber + for<'span> tracing_subscriber::registry::LookupSpan<'span>,
    {
        BunyanFormattingLayer::new(service_name.to_string(), std::io::stdout)
    }

    fn disable_targets_filter<S>(&self) -> impl Layer<S>
    where
        S: Subscriber,
    {
        let disabled_targets = self.disabled_targets.clone();
        FilterFn::new(move |metadata| !disabled_targets.contains(metadata.target()))
    }

    fn tracer<S>(&self, service_name: &'static str) -> impl Layer<S>
    where
        S: Subscriber + for<'span> tracing_subscriber::registry::LookupSpan<'span>,
    {
        let tracer = opentelemetry_otlp::new_pipeline()
            .tracing()
            .with_trace_config(trace::config().with_resource(Resource::new(vec![
                KeyValue::new("service.name", service_name),
                KeyValue::new("host.name", gethostname().into_string().unwrap()),
            ])))
            .with_exporter(
                opentelemetry_otlp::new_exporter()
                    .tonic()
                    .with_endpoint(self.otel_collector_endpoint.clone()),
            )
            .install_batch(opentelemetry::runtime::Tokio)
            .unwrap();

        tracing_opentelemetry::layer().with_tracer(tracer)
    }

    fn subscriber(&self, service_name: &'static str) -> impl Subscriber {
        Registry::default()
            .with(self.log_level_filter())
            .with(self.disable_targets_filter())
            .with(JsonStorageLayer)
            .with(self.bunyan_formatter(service_name))
            .with(self.tracer(service_name))
    }

    pub fn init_telemetry(&self, service_name: &'static str) -> Result<(), Error> {
        LogTracer::init().map_err(|_| Error::TelemetryAlreadyInit)?;
        set_text_map_propagator(TraceContextPropagator::new());
        set_global_default(self.subscriber(service_name))
            .map_err(|_| Error::TelemetryAlreadyInit)?;

        info!(
            "initializing telemetry with log level [{}]: Done",
            self.log_level
        );

        Ok(())
    }
}

#[derive(Error, Debug)]
pub enum Error {
    #[error("telemetry already initialized")]
    TelemetryAlreadyInit,
}