rocket_community/trace/subscriber/
dynamic.rs

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

rocket_community/trace/subscriber/
dynamic.rs
