Line data    Source code

       1              : use tracing::Subscriber;
       2              : use tracing_subscriber::filter::{EnvFilter, LevelFilter};
       3              : use tracing_subscriber::fmt::format::{Format, Full};
       4              : use tracing_subscriber::fmt::time::SystemTime;
       5              : use tracing_subscriber::fmt::{FormatEvent, FormatFields};
       6              : use tracing_subscriber::prelude::*;
       7              : use tracing_subscriber::registry::LookupSpan;
       8              : 
       9              : /// Initialize logging and OpenTelemetry tracing and exporter.
      10              : ///
      11              : /// Logging can be configured using `RUST_LOG` environment variable.
      12              : ///
      13              : /// OpenTelemetry is configured with OTLP/HTTP exporter. It picks up
      14              : /// configuration from environment variables. For example, to change the
      15              : /// destination, set `OTEL_EXPORTER_OTLP_ENDPOINT=http://jaeger:4318`.
      16              : /// See <https://opentelemetry.io/docs/reference/specification/sdk-environment-variables>
      17            0 : pub async fn init() -> anyhow::Result<LoggingGuard> {
      18            0 :     let env_filter = EnvFilter::builder()
      19            0 :         .with_default_directive(LevelFilter::INFO.into())
      20            0 :         .from_env_lossy()
      21            0 :         .add_directive("aws_config=info".parse().unwrap())
      22            0 :         .add_directive("azure_core::policies::transport=off".parse().unwrap());
      23            0 : 
      24            0 :     let fmt_layer = tracing_subscriber::fmt::layer()
      25            0 :         .with_ansi(false)
      26            0 :         .with_writer(std::io::stderr)
      27            0 :         .with_target(false);
      28              : 
      29            0 :     let otlp_layer = tracing_utils::init_tracing("proxy").await;
      30              : 
      31            0 :     tracing_subscriber::registry()
      32            0 :         .with(env_filter)
      33            0 :         .with(otlp_layer)
      34            0 :         .with(fmt_layer)
      35            0 :         .try_init()?;
      36              : 
      37            0 :     Ok(LoggingGuard)
      38            0 : }
      39              : 
      40              : /// Initialize logging for local_proxy with log prefix and no opentelemetry.
      41              : ///
      42              : /// Logging can be configured using `RUST_LOG` environment variable.
      43            0 : pub fn init_local_proxy() -> anyhow::Result<LoggingGuard> {
      44            0 :     let env_filter = EnvFilter::builder()
      45            0 :         .with_default_directive(LevelFilter::INFO.into())
      46            0 :         .from_env_lossy();
      47            0 : 
      48            0 :     let fmt_layer = tracing_subscriber::fmt::layer()
      49            0 :         .with_ansi(false)
      50            0 :         .with_writer(std::io::stderr)
      51            0 :         .event_format(LocalProxyFormatter(Format::default().with_target(false)));
      52            0 : 
      53            0 :     tracing_subscriber::registry()
      54            0 :         .with(env_filter)
      55            0 :         .with(fmt_layer)
      56            0 :         .try_init()?;
      57              : 
      58            0 :     Ok(LoggingGuard)
      59            0 : }
      60              : 
      61              : pub struct LocalProxyFormatter(Format<Full, SystemTime>);
      62              : 
      63              : impl<S, N> FormatEvent<S, N> for LocalProxyFormatter
      64              : where
      65              :     S: Subscriber + for<'a> LookupSpan<'a>,
      66              :     N: for<'a> FormatFields<'a> + 'static,
      67              : {
      68            0 :     fn format_event(
      69            0 :         &self,
      70            0 :         ctx: &tracing_subscriber::fmt::FmtContext<'_, S, N>,
      71            0 :         mut writer: tracing_subscriber::fmt::format::Writer<'_>,
      72            0 :         event: &tracing::Event<'_>,
      73            0 :     ) -> std::fmt::Result {
      74            0 :         writer.write_str("[local_proxy] ")?;
      75            0 :         self.0.format_event(ctx, writer, event)
      76            0 :     }
      77              : }
      78              : 
      79              : pub struct LoggingGuard;
      80              : 
      81              : impl Drop for LoggingGuard {
      82            0 :     fn drop(&mut self) {
      83            0 :         // Shutdown trace pipeline gracefully, so that it has a chance to send any
      84            0 :         // pending traces before we exit.
      85            0 :         tracing::info!("shutting down the tracing machinery");
      86            0 :         tracing_utils::shutdown_tracing();
      87            0 :     }
      88              : }