tracing_kickstart/

trace.rs

use crate::conf::TracingConfig;

use secrecy::{ExposeSecret, SecretString};
use tracing_error::ErrorLayer;
use std::collections::HashMap;
use std::fmt;
use std::fs::OpenOptions;
use std::time::Duration;
use tracing_subscriber::{EnvFilter, Layer as _};
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::util::SubscriberInitExt;

// opentelemetry - base
use opentelemetry::KeyValue;
use opentelemetry_sdk::resource::Resource;
#[cfg(feature = "detector_telemetry")]
use opentelemetry_sdk::resource::TelemetryResourceDetector;
use opentelemetry_otlp::{Protocol, WithExportConfig, WithHttpConfig};
#[cfg(feature = "detector_hostresource")]
use opentelemetry_resource_detectors::HostResourceDetector;
#[cfg(feature = "detector_os")]
use opentelemetry_resource_detectors::OsResourceDetector;
#[cfg(feature = "detector_process")]
use opentelemetry_resource_detectors::ProcessResourceDetector;

use opentelemetry_semantic_conventions::attribute;

// opentelemetry - traces
use opentelemetry_otlp::SpanExporter;
use opentelemetry_sdk::trace::{Sampler, SdkTracerProvider};
#[cfg(not(feature = "tokio_console"))]
use opentelemetry::trace::TracerProvider as _; // for tracer trait
#[cfg(not(feature = "tokio_console"))]
use tracing_opentelemetry::OpenTelemetryLayer;

// opentelemetry - metrics
use opentelemetry_otlp::MetricExporter;
use opentelemetry_sdk::metrics::{PeriodicReader, SdkMeterProvider};
use tracing_opentelemetry::MetricsLayer;
#[cfg(feature = "exponential_histograms")]
use opentelemetry_sdk::metrics::{Aggregation, InstrumentKind, Stream};

// opentelemetry - logs
use opentelemetry_otlp::LogExporter;
use opentelemetry_sdk::logs::SdkLoggerProvider;
#[cfg(not(feature = "tokio_console"))]
use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge;

pub use opentelemetry_otlp::ExporterBuildError;

// -- custom attributes + attribute helpers

pub mod custom_attribute {
    pub const DEPLOYMENT_BUILD_TYPE: &str = "deployment.build_type";
    #[cfg(feature = "attrs_crate_name")]
    pub const SERVICE_CRATE_NAME: &str = "service.crate_name";
    #[cfg(feature = "attrs_version_expanded")]
    pub const SERVICE_VERSION_MAJOR: &str = "service.version.major";
    #[cfg(feature = "attrs_version_expanded")]
    pub const SERVICE_VERSION_MINOR: &str = "service.version.minor";
    #[cfg(feature = "attrs_version_expanded")]
    pub const SERVICE_VERSION_PATCH: &str = "service.version.patch";

    #[cfg(feature = "attrs_origin")]
    pub const SERVICE_ORIGIN_PACKAGE_NAME: &str = "service.origin.package_name";
    #[cfg(feature = "attrs_origin")]
    pub const SERVICE_ORIGIN_CRATE_NAME: &str = "service.origin.crate_name";
}
#[rustfmt::skip]
pub fn get_build_env() -> &'static str {
    #[cfg(debug_assertions)]
    { "debug" }
    #[cfg(not(debug_assertions))]
    { "release" }
}
pub fn get_origin_package_name() -> Option<&'static str> {
    let package_name = env!("CARGO_PKG_NAME");
    if package_name.is_empty() {
        None
    } else {
        Some(package_name)
    }
}
pub fn get_origin_crate_name() -> Option<&'static str> {
    let package_name = env!("CARGO_CRATE_NAME");
    if package_name.is_empty() {
        None
    } else {
        Some(package_name)
    }
}

fn build_otel_resource(service_attrs: &ServiceAttributeStore, deployment_env: Option<String>, custom_attrs: Vec<KeyValue>) -> Resource {
    // root/primary service name + package name
    let mut builder = Resource::builder_empty()
    .with_attribute(KeyValue::new(attribute::SERVICE_NAME, service_attrs.pkg_name));

    #[cfg(feature = "attrs_crate_name")]
    {
        builder = builder.with_attribute(KeyValue::new(custom_attribute::SERVICE_CRATE_NAME, service_attrs.crate_name));
    }

    // version
    builder = builder.with_attribute(KeyValue::new(attribute::SERVICE_VERSION, service_attrs.version));

    #[cfg(feature = "attrs_version_expanded")] {
        builder = builder
        .with_attribute(KeyValue::new(custom_attribute::SERVICE_VERSION_MAJOR, service_attrs.version_major))
        .with_attribute(KeyValue::new(custom_attribute::SERVICE_VERSION_MINOR, service_attrs.version_minor))
        .with_attribute(KeyValue::new(custom_attribute::SERVICE_VERSION_PATCH, service_attrs.version_patch));
    }

    #[cfg(feature = "attrs_origin")]
    {
        // returns the name of the package that contains the associated tracing call
        if let Some(origin_package_name) = get_origin_package_name() {
            builder = builder.with_attribute(KeyValue::new(custom_attribute::SERVICE_ORIGIN_PACKAGE_NAME, origin_package_name));
        }
        if let Some(origin_crate_name) = get_origin_crate_name() {
            builder = builder.with_attribute(KeyValue::new(custom_attribute::SERVICE_ORIGIN_CRATE_NAME, origin_crate_name));
        }
    }

    // build mode: release/debug
    builder = builder.with_attribute(KeyValue::new(custom_attribute::DEPLOYMENT_BUILD_TYPE, get_build_env()));

    // deployment env set from config/runtime env
    if let Some(env) = deployment_env {
        builder = builder.with_attribute(KeyValue::new(attribute::DEPLOYMENT_ENVIRONMENT_NAME, env));
    }

    // custom resource attrs
    for attr in custom_attrs {
        builder = builder.with_attribute(attr);
    }

    #[cfg(feature = "detector_telemetry")]
    {
        // telemetry sdk stack attrs
        builder = builder.with_detector(Box::new(TelemetryResourceDetector));
    }
    #[cfg(feature = "detector_hostresource")]
    {
        // host id, host arch
        builder = builder.with_detector(Box::new(HostResourceDetector::default()));
    }
    #[cfg(feature = "detector_process")]
    {
        // process args, pid
        builder = builder.with_detector(Box::new(ProcessResourceDetector));
    }
    #[cfg(feature = "detector_os")]
    {
        // os
        builder = builder.with_detector(Box::new(OsResourceDetector));
    }

    builder.build()
}

fn build_otel_headers(auth_header_val: &Option<SecretString>) -> HashMap<String, String> {
    let mut headers: HashMap<String, String> = HashMap::new();

    // add auth headers if provided
    if let Some(auth_header) = auth_header_val.as_ref() {
        headers.insert("Authorization".into(), auth_header.expose_secret().into());
    }

    headers
}

// Construct TracerProvider for OpenTelemetryLayer
fn init_otel_traces_provider(
    collector_endpoint: &str,
    headers: HashMap<String, String>,
    resource: Resource,
) -> Result<SdkTracerProvider, ExporterBuildError> {
    let exporter = SpanExporter::builder()
        .with_http()
        .with_headers(headers)
        .with_endpoint(format!("{collector_endpoint}/v1/traces"))
        .with_protocol(Protocol::HttpBinary)
        // .with_timeout(std::time::Duration::from_secs(3))
        .build()?;

    let provider = SdkTracerProvider::builder()
        // Customize sampling strategy
        .with_sampler(Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(1.0))))
        .with_resource(resource)
        .with_batch_exporter(exporter)
        .build();

    Ok(provider)
}
fn init_otel_logs_provider(
    collector_endpoint: &str,
    headers: HashMap<String, String>,
    resource: Resource,
) -> Result<SdkLoggerProvider, ExporterBuildError> {
    let exporter = LogExporter::builder()
        .with_http()
        .with_headers(headers)
        .with_endpoint(format!("{collector_endpoint}/v1/logs"))
        .with_protocol(Protocol::HttpBinary)
        // .with_timeout(std::time::Duration::from_secs(3))
        .build()?;

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

    Ok(provider)
}
fn init_otel_metrics_provider(
    collector_endpoint: &str,
    headers: HashMap<String, String>,
    resource: Resource,
    interval: Option<Duration>,
) -> Result<SdkMeterProvider, ExporterBuildError> {
    let exporter = MetricExporter::builder()
        .with_http()
        .with_headers(headers)
        .with_endpoint(format!("{collector_endpoint}/v1/metrics"))
        .with_protocol(Protocol::HttpBinary)
        // .with_timeout(std::time::Duration::from_secs(3))
        .build()?;

    let mut periodic = PeriodicReader::builder(exporter);
    if let Some(duration) = interval {
        periodic = periodic.with_interval(duration);
    }
    let mut builder = SdkMeterProvider::builder();
    builder = builder
        .with_resource(resource)
        .with_reader(periodic.build());
    #[cfg(feature = "exponential_histograms")]
    {
        builder = builder.with_view(|inst| {
            if let InstrumentKind::Histogram = inst.kind() {
                let s = Stream::builder()
                    .with_aggregation(Aggregation::Base2ExponentialHistogram {
                        max_size: 160,
                        max_scale: 20,
                        record_min_max: true,
                    })
                    .build()
                    .unwrap();
                Some(s)
            } else {
                None
            }
        });
    }
    let provider = builder.build();

    Ok(provider)
}

/// Compile-time attributes to be provided by the owning application/service.
///
/// Used as a set of parameters required by [`init()`]
///
/// Service attributes can be generated using the provided [`build_attrs!`] macro , e.g.:
#[derive(Debug, Clone)]
pub struct ServiceAttributeStore {
    pub crate_name: &'static str,
    pub pkg_name: &'static str,
    pub version: &'static str,
    pub version_major: &'static str,
    pub version_minor: &'static str,
    pub version_patch: &'static str,
}
impl ServiceAttributeStore {
    pub fn dump(&self) {
        let service_name = self.pkg_name;
        let crate_name = self.crate_name;
        let service_version = self.version;
        let service_version_major = self.version_major;
        let service_version_minor = self.version_minor;
        let service_version_patch = self.version_patch;
        let origin_pkg_name = get_origin_package_name().unwrap_or("- unset -");
        let origin_crate_name = get_origin_crate_name().unwrap_or("- unset -");
        let build_env = get_build_env();

        println!();
        println!("Resolved tracing attributes");
        println!("--------------------");
        println!("service_name (pkg_name): {service_name}");
        println!("service_crate_name:      {crate_name}");
        println!("service_version:         {service_version}");
        println!("service_version_major:   {service_version_major}");
        println!("service_version_minor:   {service_version_minor}");
        println!("service_version_patch:   {service_version_patch}");
        println!("origin_pkg_name:         {origin_pkg_name}");
        println!("origin_crate_name:       {origin_crate_name}");
        println!("build_env:               {build_env}");
        println!();
    }
}

/// Generates service attributes using env! calls.
///
/// This is done using a macro to allow for the `env!(..)` calls to be scoped from the
/// parent package/crate, rather than from `tracing-kickstart`.
#[macro_export]
macro_rules! build_attrs {
    // This macro takes an argument of designator `ident` and
    // creates a function named `$func_name`.
    // The `ident` designator is used for variable/function names.
    () => (
        tracing_kickstart::ServiceAttributeStore {
            crate_name: env!("CARGO_CRATE_NAME"),
            pkg_name: env!("CARGO_PKG_NAME"),
            version: env!("CARGO_PKG_VERSION"),
            version_major: env!("CARGO_PKG_VERSION_MAJOR"),
            version_minor: env!("CARGO_PKG_VERSION_MINOR"),
            version_patch: env!("CARGO_PKG_VERSION_PATCH"),
        }
    )
}
fn validate_non_empty_filter_str(filter: &str, source_name: &'static str) -> bool {
    if filter.is_empty() {
        println!("Ignoring empty filter string sourced from {source_name}");
        false
    } else {
        true
    }
}
fn validate_non_empty_filter(filter: &EnvFilter, source_name: &'static str) -> bool {
    let filter_str = filter.to_string();
    validate_non_empty_filter_str(&filter_str, source_name)
}


/// Initialize tracing
///
/// Note: service attributes can be generated and passed in using the `build_attrs` macro, e.g.:
///
/// ---
///
/// # Examples
///
/// Basic usage
///
/// ```no_run
/// use config::{Config, Environment};
/// use dotenvy::dotenv;
/// use serde::Deserialize;
/// use tracing_kickstart::{TracingConfig, TracingConfigOverride};
///
/// #[derive(Debug, Clone, Deserialize)]
/// pub struct Conf {
///     // pulled in automatically from env vars
///     trace: TracingConfigOverride,
/// }
///
/// fn main() {
///     // load config
///     dotenv().ok(); // load vars from .env file
///     let settings = Config::builder()
///         .add_source(Environment::with_prefix("APP").separator("__").try_parsing(true))
///         .build()
///         .unwrap();
///     let conf = settings.try_deserialize::<Conf>().unwrap(); // deserialize into Conf struct
///
///     // collect attributes for this crate
///     let attrs = tracing_kickstart::build_attrs!();
///     attrs.dump(); // log attributes to stdout
///
///     // set an optional env filter to replace the default set by tracing_kickstart: `info,{crate}=debug`
///     // if set, this will override any filters set using `RUST_LOG`, but can be completely rewritten at runtime
///     // using the `TracingConfigOverride::filter` var.
///     //
///     // or, to retain this base filter but make adjustments at runtime, use `TracingConfigOverride::filter` to
///     // avoid having to copy and paste your base filter into your env vars / .env file.
///     let custom_base_filter = Some(format!("info,{}=trace", attrs.crate_name)); // default: `info,{crate}=debug`
///
///     // optionally add custom resource attributes
///     let custom_resource_attrs = Some(vec![("region".into(), "canada".into())]);
///
///     // add your 'base' configuration here. Most options can be overriden once again at runtime
///     // via `TracingConfigOverride`, which is the type of the required parameter passed in
///     //
///     // the call to `build()` will ensure that any present override values are applied
///     let tracing_config = TracingConfig::builder(&conf.trace)
///         .filter(custom_base_filter) // optionally add a base filter
///         .ansi_output(true)
///         .ansi_sanitization(false) // retain ansi escape codes in `tracing` output
///                                   // (NOTE: do not disable if logs contain untrusted content)
///         .custom_resource_attrs(custom_resource_attrs) // (NOTE: cannot be set with .env override)
///         .build();
///
///     // to experiment with the override resolution, you can print the tracing config to view the output
///     println!("Resolved tracing config:\n{tracing_config:#?}");
///
///     // init tracing, receive a handle for the tracing providers
///     let tracing_providers = tracing_kickstart::init(attrs, tracing_config).unwrap();
///     tracing_providers.register_globally(); // optionally register all configured providers globally
///     tracing::info!("Tracing initialized");
///
///     // graceful shutdown of various tracing providers (logs, metrics, traces) using the provided handle
///     tracing::debug!("Shutting down tracing providers: {tracing_providers:?}");
///     tracing_providers.shutdown();
/// }
/// ```
///
/// ## `EnvFilter`
///
/// The EnvFilter is resolved using the first available from:
/// - `TracingConfigOverride::filter` (typically set from app config env var, e.g. `APP__TRACING__FILTER=app=warn`)
/// - `TracingConfig::filter` (typically set from init code)
/// - `RUST_LOG` env var
/// - The `default_env_filter` parameter in this function (used to overide the default fallback)
/// - default fallback (library defined, set to `"info,{crate_name}=debug`)
/// ---
/// Regardless of how the `EnvFilter` is resolved, all required filters for `console_subscriber` will be added
/// **if the console_subscriber** feature flag is enabled.
///
/// To add an adjustment at runtime - *without copying over the base env filter configured in your init code* -
/// use [`TracingConfigOverride::filter_append`]
// if tracing config is none, otel providers won't be handled
pub fn init(
    service_attrs: ServiceAttributeStore,
    config: TracingConfig,
) -> Result<TraceProviders, ExporterBuildError> {
    // resolve the env filter in the following priority
    let mut base_filter: EnvFilter = {
        // config env filter
        if let Some(filter_str) = &config.filter && validate_non_empty_filter_str(filter_str, "provided config") {
            println!("Resolved tracing EnvFilter from provided config: {filter_str:?}");
            filter_str.into()
        }
        // `RUST LOG`
        else if let Ok(filter) = EnvFilter::try_from_default_env() && validate_non_empty_filter(&filter, "RUST_LOG") {
            println!("Resolved tracing EnvFilter from `RUST_LOG`: {:?}", filter.to_string());
            filter
        }
        // library-defined fallback env filter
        else {
            let filter_str = format!(
                "info,{}=debug",
                service_attrs.crate_name
            );
            println!("Using tracing-kickstart fallback EnvFilter: {filter_str:?}");
            filter_str.into()
        }
    };

    // add from filter_append
    if let Some(append) = config.filter_append.as_ref() {
        println!("Will merge EnvFilter with provided `filter_append`: {append:#?}");

        for dir in append.split(",").filter(|s| !s.is_empty()) {
            base_filter = base_filter.add_directive(dir.parse().expect("valid filter_append syntax"));
        }
    }
    // add env filters for tokio console subscriber (controlled by feature flag)
    #[cfg(feature = "tokio_console")]
    let registry_filter = EnvFilter::new(format!("{base_filter},tokio=trace,runtime=trace"));
    #[cfg(not(feature = "tokio_console"))]
    let registry_filter = base_filter.clone();

    // print the resolved env filter
    println!("Using base tracing filters: {base_filter}");
    if registry_filter.to_string() != base_filter.to_string() {
        println!("Registry tracing filter: {registry_filter}");
    }

    // build base layers
    let layer = tracing_subscriber::registry()
        .with(registry_filter)
        .with(ErrorLayer::default());

    // stdout layer
    let layer = layer.with(
        tracing_subscriber::fmt::layer()
            .with_ansi(config.ansi_output)
            .with_ansi_sanitization(config.ansi_sanitization)
            .with_filter(base_filter.clone()) // use less permissive filter for stdout/logs
    );

    // conditionally add log file layer if path is provided in config
    let file_logging_layer = config.log_file_path.as_ref().map(|path| {
        let file = OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(path)
            .expect("Log file should be writable");

        tracing_subscriber::fmt::layer()
            .with_ansi(false)
            .with_writer(file)
            .with_filter(base_filter)
    });
    let layer = layer.with(file_logging_layer);

    // conditionally add tokio console layer
    #[cfg(feature = "tokio_console")]
    let layer = layer.with(console_subscriber::spawn());

    // default has all 3 provider field options set to None
    let mut providers_handle = TraceProviders::default();

    // init open telemetry providers
    if let Some(otel_config) = &config.otel_config {
        println!("Initializing OTEL config");
        let endpoint = &otel_config.collector_url;
        let headers = build_otel_headers(&otel_config.collector_auth_header);
        let custom_attrs = config.custom_resource_attrs
            .clone()
            .unwrap_or_default()
            .into_iter()
            .map(|(k,v)| KeyValue::new(k, v))
            .collect();
        let resource = build_otel_resource(&service_attrs, config.deployment_env.clone(), custom_attrs);

        // traces
        let traces_provider = init_otel_traces_provider(endpoint, headers.clone(), resource.clone())?;
        // - add tracing layer for tracing/span -> otel/trace
        // skipped when tokio console is enabled to prevent flooding
        #[cfg(not(feature = "tokio_console"))]
        let layer = layer.with(OpenTelemetryLayer::new(traces_provider.tracer(service_attrs.crate_name)).with_level(true));
        providers_handle.traces = Some(traces_provider);

        // logs
        let logs_provider = init_otel_logs_provider(endpoint, headers.clone(), resource.clone())?;
        // - add tracing layer for tracing -> otel/logs
        // skipped when tokio console is enabled to prevent flooding
        #[cfg(not(feature = "tokio_console"))]
        let layer = {
            #[allow(unused_mut)]
            let mut bridge_builder = OpenTelemetryTracingBridge::builder(&logs_provider);
            #[cfg(feature = "otel_span_attributes")]
            {
                use opentelemetry_appender_tracing::layer::TracingSpanAttributes;
                bridge_builder = bridge_builder.with_tracing_span_attributes(TracingSpanAttributes::all());
            }
            layer.with(bridge_builder.build())
        };
        providers_handle.logs = Some(logs_provider);

        // metrics
        let metrics_provider = init_otel_metrics_provider(endpoint, headers, resource, config.metrics_interval_duration())?;
        // - add layer for tracing events -> otel/metrics
        let layer = layer.with(MetricsLayer::new(metrics_provider.clone()));
        providers_handle.metrics = Some(metrics_provider);

        layer.init();
        println!("{:-<1$}", "-", 30);
        tracing::info!("OTEL tracing configured");
    } else {
        layer.init();
        println!("{:-<1$}", "-", 30);
        tracing::warn!("OTEL tracing disabled");
    }

    Ok(providers_handle)
}

// ---- Struct for containing otel providers

// TODO: alternatively use `opentelemetry::global::set_x_provider()` fns
#[derive(Default, Clone)]
pub struct TraceProviders {
    pub traces: Option<SdkTracerProvider>,
    pub logs: Option<SdkLoggerProvider>,
    pub metrics: Option<SdkMeterProvider>,
}
impl TraceProviders {
    /// Calls `opentelemetry::global::set_x_provider(..); for all configured providers, where applicable`
    pub fn register_globally(&self) {
        // register traces
        if let Some(provider) = &self.traces {
            tracing::info!("Traces provider registered globally");
            opentelemetry::global::set_tracer_provider(provider.clone());
        }
        // register metrics
        if let Some(provider) = &self.metrics {
            tracing::info!("Metrics provider registered globally");
            opentelemetry::global::set_meter_provider(provider.clone());
        }
    }

    /// Triggers shutdown for each provider that has been set
    pub fn shutdown(self) {
        // shutdown traces
        if let Some(provider) = self.traces && let Err(error) = provider.shutdown() {
            println!("error shutting down traces provider: {error}");
        }
        // shutdown logs
        if let Some(provider) = self.logs && let Err(error) = provider.shutdown() {
            println!("error shutting down logs provider: {error}");
        }
        // shutdown metrics
        if let Some(provider) = self.metrics && let Err(error) = provider.shutdown() {
            println!("error shutting down metrics provider: {error}");
        }
    }
}
impl fmt::Debug for TraceProviders {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "TraceProviders (")?;
        let mut i = 0;
        if self.traces.is_some() {
            write!(f, "SdkTracerProvider")?;
            i += 1;
        }
        if self.logs.is_some() {
            if i > 0 {
                write!(f, ", ")?;
            }
            write!(f, "SdkLoggerProvider")?;
            i += 1;
        }
        if self.metrics.is_some() {
            if i > 0 {
                write!(f, ", ")?;
            }
            write!(f, "SdkMeterProvider")?;
        }
        write!(f, ")")
    }
}