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