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("azure_core::policies::transport=off".parse().unwrap());
      22            0 : 
      23            0 :     let fmt_layer = tracing_subscriber::fmt::layer()
      24            0 :         .with_ansi(false)
      25            0 :         .with_writer(std::io::stderr)
      26            0 :         .with_target(false);
      27              : 
      28            0 :     let otlp_layer = tracing_utils::init_tracing("proxy").await;
      29              : 
      30            0 :     tracing_subscriber::registry()
      31            0 :         .with(env_filter)
      32            0 :         .with(otlp_layer)
      33            0 :         .with(fmt_layer)
      34            0 :         .try_init()?;
      35              : 
      36            0 :     Ok(LoggingGuard)
      37            0 : }
      38              : 
      39              : /// Initialize logging for local_proxy with log prefix and no opentelemetry.
      40              : ///
      41              : /// Logging can be configured using `RUST_LOG` environment variable.
      42            0 : pub fn init_local_proxy() -> anyhow::Result<LoggingGuard> {
      43            0 :     let env_filter = EnvFilter::builder()
      44            0 :         .with_default_directive(LevelFilter::INFO.into())
      45            0 :         .from_env_lossy();
      46            0 : 
      47            0 :     let fmt_layer = tracing_subscriber::fmt::layer()
      48            0 :         .with_ansi(false)
      49            0 :         .with_writer(std::io::stderr)
      50            0 :         .event_format(LocalProxyFormatter(Format::default().with_target(false)));
      51            0 : 
      52            0 :     tracing_subscriber::registry()
      53            0 :         .with(env_filter)
      54            0 :         .with(fmt_layer)
      55            0 :         .try_init()?;
      56              : 
      57            0 :     Ok(LoggingGuard)
      58            0 : }
      59              : 
      60              : pub struct LocalProxyFormatter(Format<Full, SystemTime>);
      61              : 
      62              : impl<S, N> FormatEvent<S, N> for LocalProxyFormatter
      63              : where
      64              :     S: Subscriber + for<'a> LookupSpan<'a>,
      65              :     N: for<'a> FormatFields<'a> + 'static,
      66              : {
      67            0 :     fn format_event(
      68            0 :         &self,
      69            0 :         ctx: &tracing_subscriber::fmt::FmtContext<'_, S, N>,
      70            0 :         mut writer: tracing_subscriber::fmt::format::Writer<'_>,
      71            0 :         event: &tracing::Event<'_>,
      72            0 :     ) -> std::fmt::Result {
      73            0 :         writer.write_str("[local_proxy] ")?;
      74            0 :         self.0.format_event(ctx, writer, event)
      75            0 :     }
      76              : }
      77              : 
      78              : pub struct LoggingGuard;
      79              : 
      80              : impl Drop for LoggingGuard {
      81            0 :     fn drop(&mut self) {
      82            0 :         // Shutdown trace pipeline gracefully, so that it has a chance to send any
      83            0 :         // pending traces before we exit.
      84            0 :         tracing::info!("shutting down the tracing machinery");
      85            0 :         tracing_utils::shutdown_tracing();
      86            0 :     }
      87              : }