Skip to main content

espresso_utils/
logging.rs

1use clap::{Parser, ValueEnum};
2use log_panics::BacktraceMode;
3use tracing_subscriber::{
4    EnvFilter, Layer, Registry,
5    fmt::{format::FmtSpan, writer::BoxMakeWriter},
6    layer::SubscriberExt,
7    util::SubscriberInitExt,
8};
9
10/// Controls how backtraces are logged on panic.
11///
12/// The values here match the possible values of `RUST_LOG_FORMAT`, and their corresponding behavior
13/// on backtrace logging is:
14/// * `full`: print a prettified dump of the stack trace and span trace to stdout, optimized for
15///   human readability rather than machine parsing
16/// * `compact`: output the default panic message, with backtraces controlled by `RUST_BACKTRACE`
17/// * `json`: output the panic message and stack trace as a tracing event. This in turn works with
18///   the behavior of the tracing subscriber with `RUST_LOG_FORMAT=json` to output the event in a
19///   machine-parseable, JSON format.
20#[derive(Clone, Copy, Debug, Default, ValueEnum)]
21enum BacktraceLoggingMode {
22    #[default]
23    Full,
24    Compact,
25    Json,
26}
27
28/// Logging configuration.
29#[derive(Clone, Debug, Default, Parser)]
30pub struct Config {
31    #[clap(long, env = "RUST_LOG_FORMAT")]
32    backtrace_mode: Option<BacktraceLoggingMode>,
33}
34
35impl Config {
36    /// Get the logging configuration from the environment.
37    pub fn from_env() -> Self {
38        Self::parse_from(std::iter::empty::<String>())
39    }
40
41    /// Initialize logging and panic handlers based on this configuration.
42    pub fn init(&self) {
43        initialize_logging();
44        self.install_panic_hook();
45    }
46
47    /// Like `init`, but logs to stderr so that stdout carries only the program's own output.
48    pub fn init_on_stderr(&self) {
49        initialize_logging_on_stderr();
50        self.install_panic_hook();
51    }
52
53    /// Like `init`, but also attaches an extra tracing `Layer` (e.g. an OTel bridge).
54    pub fn init_with_otel<L>(&self, otel_layer: Option<L>)
55    where
56        L: Layer<FmtSubscriber> + Send + Sync + 'static,
57    {
58        initialize_logging_with(otel_layer);
59        self.install_panic_hook();
60    }
61
62    fn install_panic_hook(&self) {
63        if let BacktraceLoggingMode::Json = self.backtrace_mode.unwrap_or_default() {
64            log_panics::Config::new()
65                .backtrace_mode(BacktraceMode::Resolved)
66                .install_panic_hook();
67        }
68    }
69}
70
71/// A type-erased fmt layer attached over a `Registry`. The OTel bridge layer
72/// wraps a subscriber of type `Layered<ErasedFmtLayer, Registry>`.
73pub type ErasedFmtLayer = Box<dyn Layer<Registry> + Send + Sync + 'static>;
74
75/// Composed subscriber type after the fmt layer is applied. Callers building
76/// an OTel layer for `initialize_logging_with` should target this subscriber.
77pub type FmtSubscriber = tracing_subscriber::layer::Layered<ErasedFmtLayer, Registry>;
78
79/// Initializes logging on stdout.
80pub fn initialize_logging() {
81    Registry::default().with(fmt_layer(Stream::Stdout)).init();
82}
83
84/// Initializes logging on stderr.
85///
86/// For binaries whose stdout carries data a caller may pipe or capture, rather than diagnostics.
87pub fn initialize_logging_on_stderr() {
88    Registry::default().with(fmt_layer(Stream::Stderr)).init();
89}
90
91/// Initializes logging with an optional extra `Layer` (e.g. an OTel bridge).
92///
93/// The extra layer must implement `Layer<FmtSubscriber>`. A polymorphic layer
94/// like `OpenTelemetryTracingBridge` satisfies this naturally.
95pub fn initialize_logging_with<L>(extra: Option<L>)
96where
97    L: Layer<FmtSubscriber> + Send + Sync + 'static,
98{
99    Registry::default()
100        .with(fmt_layer(Stream::Stdout))
101        .with(extra)
102        .init();
103}
104
105/// Which standard stream the fmt layer writes to.
106#[derive(Clone, Copy, Debug)]
107pub enum Stream {
108    Stdout,
109    Stderr,
110}
111
112fn fmt_layer(stream: Stream) -> ErasedFmtLayer {
113    let span_event_filter = parse_span_filter();
114    let writer = match stream {
115        Stream::Stdout => BoxMakeWriter::new(std::io::stdout),
116        Stream::Stderr => BoxMakeWriter::new(std::io::stderr),
117    };
118    let json_mode = std::env::var("RUST_LOG_FORMAT") == Ok("json".to_string());
119    if json_mode {
120        tracing_subscriber::fmt::layer()
121            .json()
122            .with_writer(writer)
123            .with_span_events(span_event_filter)
124            .with_filter(EnvFilter::from_default_env())
125            .boxed()
126    } else {
127        tracing_subscriber::fmt::layer()
128            .with_writer(writer)
129            .with_span_events(span_event_filter)
130            .with_filter(EnvFilter::from_default_env())
131            .boxed()
132    }
133}
134
135fn parse_span_filter() -> FmtSpan {
136    match std::env::var("RUST_LOG_SPAN_EVENTS") {
137        Ok(val) => val
138            .split(',')
139            .map(|s| match s.trim() {
140                "new" => FmtSpan::NEW,
141                "enter" => FmtSpan::ENTER,
142                "exit" => FmtSpan::EXIT,
143                "close" => FmtSpan::CLOSE,
144                "active" => FmtSpan::ACTIVE,
145                "full" => FmtSpan::FULL,
146                _ => FmtSpan::NONE,
147            })
148            .fold(FmtSpan::NONE, |acc, x| acc | x),
149        Err(_) => FmtSpan::NONE,
150    }
151}