rocket/trace/subscriber/
dynamic.rs

Help
1use std::sync::OnceLock;
2
3use tracing::{Dispatch, Event, Metadata};
4use tracing::subscriber::{Subscriber, Interest};
5use tracing::span::{Attributes, Id, Record};
6
7use tracing_subscriber::reload;
8use tracing_subscriber::registry::{Registry, LookupSpan};
9use tracing_subscriber::layer::{Context, Layer, Layered, SubscriberExt};
10use tracing_subscriber::util::SubscriberInitExt;
11
12use crate::config::Config;
13use crate::trace::subscriber::{Compact, Pretty, RequestId, RequestIdLayer, RocketFmt};
14use crate::trace::TraceFormat;
15
16/// A subscriber that is either a [`Pretty`] or [`Compact`] [`RocketFmt`].
17pub struct RocketDynFmt {
18    inner: either::Either<RocketFmt<Compact>, RocketFmt<Pretty>>,
19}
20
21impl From<RocketFmt<Compact>> for RocketDynFmt {
22    fn from(value: RocketFmt<Compact>) -> Self {
23        RocketDynFmt { inner: either::Either::Left(value) }
24    }
25}
26
27impl From<RocketFmt<Pretty>> for RocketDynFmt {
28    fn from(value: RocketFmt<Pretty>) -> Self {
29        RocketDynFmt { inner: either::Either::Right(value) }
30    }
31}
32
33impl RocketDynFmt {
34    /// Creates a new `RocketDynFmt` subscriber given a `Config`.
35    ///
36    /// [`Config::log_format`] determines which `RocketFmt` subscriber (either
37    /// [`Pretty`] or [`Compact`]) is used.
38    ///
39    /// If `config` is `None`, [`Config::debug_default()`] is used, which uses
40    /// the [`Pretty`] subscriber by default.
41    pub fn new(config: Option<&Config>) -> Self {
42        let default = Config::debug_default();
43        let workers = config.map_or(default.workers, |c| c.workers);
44        let colors = config.map_or(default.cli_colors, |c| c.cli_colors);
45        let level = config.map_or(default.log_level, |c| c.log_level);
46        let format = config.map_or(default.log_format, |c| c.log_format);
47
48        match format {
49            TraceFormat::Pretty => Self::from(RocketFmt::<Pretty>::new(workers, colors, level)),
50            TraceFormat::Compact => Self::from(RocketFmt::<Compact>::new(workers, colors, level)),
51        }
52    }
53
54    pub(crate) fn init(config: Option<&Config>) {
55        type Handle = reload::Handle<RocketDynFmt, Layered<RequestIdLayer, Registry>>;
56
57        static HANDLE: OnceLock<Handle> = OnceLock::new();
58
59        // Do nothing if there's no config and we've already initialized.
60        if config.is_none() && HANDLE.get().is_some() {
61            return;
62        }
63
64        let formatter = Self::new(config);
65        if let Some(handle) = HANDLE.get() {
66            return assert!(handle.modify(|layer| *layer = formatter).is_ok());
67        }
68
69        let (layer, reload_handle) = reload::Layer::new(formatter);
70        let result = tracing_subscriber::registry()
71            .with(RequestId::layer())
72            .with(layer)
73            .try_init();
74
75        if result.is_ok() {
76            assert!(HANDLE.set(reload_handle).is_ok());
77        }
78    }
79}
80
81macro_rules! forward {
82    ($T:ident => $(& $r:tt)? $method:ident ( $($p:ident : $t:ty),* ) $(-> $R:ty)?) => {
83        #[inline(always)]
84        fn $method(& $($r)? self $(, $p : $t)*) $(-> $R)? {
85            match & $($r)* self.inner {
86                either::Either::Left(layer) => Layer::<$T>::$method(layer, $($p),*),
87                either::Either::Right(layer) => Layer::<$T>::$method(layer, $($p),*),
88            }
89        }
90    };
91}
92
93impl<S: Subscriber + for<'a> LookupSpan<'a>> Layer<S> for RocketDynFmt {
94    forward!(S => on_register_dispatch(subscriber: &Dispatch));
95    forward!(S => &mut on_layer(subscriber: &mut S));
96    forward!(S => register_callsite(metadata: &'static Metadata<'static>) -> Interest);
97    forward!(S => enabled(metadata: &Metadata<'_>, ctx: Context<'_, S>) -> bool);
98    forward!(S => on_new_span(attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>));
99    forward!(S => on_record(_span: &Id, _values: &Record<'_>, _ctx: Context<'_, S>));
100    forward!(S => on_follows_from(_span: &Id, _follows: &Id, _ctx: Context<'_, S>));
101    forward!(S => event_enabled(_event: &Event<'_>, _ctx: Context<'_, S>) -> bool);
102    forward!(S => on_event(_event: &Event<'_>, _ctx: Context<'_, S>));
103    forward!(S => on_enter(_id: &Id, _ctx: Context<'_, S>));
104    forward!(S => on_exit(_id: &Id, _ctx: Context<'_, S>));
105    forward!(S => on_close(_id: Id, _ctx: Context<'_, S>));
106    forward!(S => on_id_change(_old: &Id, _new: &Id, _ctx: Context<'_, S>));
107}
rocket/trace/subscriber/dynamic.rs

rocket/trace/subscriber/
dynamic.rs
