Skip to content

Commit

Permalink
Add 'RocketDynFmt' subscriber constructor.
Browse files Browse the repository at this point in the history
Resolves #2840.
  • Loading branch information
SergioBenitez committed Aug 6, 2024
1 parent 4eb6cdf commit f2ca2ad
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 33 deletions.
9 changes: 1 addition & 8 deletions core/lib/src/config/secret_key.rs
Original file line number Diff line number Diff line change
@@ -1,17 +1,10 @@
use std::fmt;

use cookie::Key;
use serde::{de, ser, Deserialize, Serialize};
use serde::{de, ser, Deserialize};

use crate::request::{Outcome, Request, FromRequest};

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
enum Kind {
Zero,
Generated,
Provided
}

/// A cryptographically secure secret key.
///
/// A `SecretKey` is primarily used by [private cookies]. See the [configuration
Expand Down
5 changes: 0 additions & 5 deletions core/lib/src/trace/subscriber/common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,11 +58,6 @@ impl<K: private::FmtKind> RocketFmt<K> {
}
}

pub fn reset(&mut self, cli_colors: CliColors, level: Option<Level>) {
let state = std::mem::take(&mut self.state);
*self = Self { state, ..Self::new(0, cli_colors, level) };
}

pub fn style(&self, metadata: &Metadata<'_>) -> Style {
match *metadata.level() {
Level::ERROR => self.style.red(),
Expand Down
51 changes: 31 additions & 20 deletions core/lib/src/trace/subscriber/dynamic.rs
Original file line number Diff line number Diff line change
@@ -1,17 +1,19 @@
use std::sync::OnceLock;

use tracing::{Dispatch, Event, Level, Metadata};
use tracing::{Dispatch, Event, Metadata};
use tracing::subscriber::{Subscriber, Interest};
use tracing::span::{Attributes, Id, Record};

use tracing_subscriber::reload;
use tracing_subscriber::registry::{Registry, LookupSpan};
use tracing_subscriber::layer::{Context, Layer, Layered, SubscriberExt};
use tracing_subscriber::reload;
use tracing_subscriber::util::SubscriberInitExt;

use crate::config::{Config, CliColors};
use crate::config::Config;
use crate::trace::subscriber::{Compact, Pretty, RequestId, RequestIdLayer, RocketFmt};
use crate::trace::TraceFormat;

/// A subscriber that is either a [`Pretty`] or [`Compact`] [`RocketFmt`].
pub struct RocketDynFmt {
inner: either::Either<RocketFmt<Compact>, RocketFmt<Pretty>>,
}
Expand All @@ -29,7 +31,27 @@ impl From<RocketFmt<Pretty>> for RocketDynFmt {
}

impl RocketDynFmt {
pub fn init(config: Option<&Config>) {
/// Creates a new `RocketDynFmt` subscriber given a `Config`.
///
/// [`Config::log_format`] determines which `RocketFmt` subscriber (either
/// [`Pretty`] or [`Compact`]) is used.
///
/// If `config` is `None`, [`Config::debug_default()`] is used, which uses
/// the [`Pretty`] subscriber by default.
pub fn new(config: Option<&Config>) -> Self {
let default = Config::debug_default();
let workers = config.map_or(default.workers, |c| c.workers);
let colors = config.map_or(default.cli_colors, |c| c.cli_colors);
let level = config.map_or(default.log_level, |c| c.log_level);
let format = config.map_or(default.log_format, |c| c.log_format);

match format {
TraceFormat::Pretty => Self::from(RocketFmt::<Pretty>::new(workers, colors, level)),
TraceFormat::Compact => Self::from(RocketFmt::<Compact>::new(workers, colors, level)),
}
}

pub(crate) fn init(config: Option<&Config>) {
type Handle = reload::Handle<RocketDynFmt, Layered<RequestIdLayer, Registry>>;

static HANDLE: OnceLock<Handle> = OnceLock::new();
Expand All @@ -39,32 +61,21 @@ impl RocketDynFmt {
return;
}

let workers = config.map_or(num_cpus::get(), |c| c.workers);
let colors = config.map_or(CliColors::Auto, |c| c.cli_colors);
let level = config.map_or(Some(Level::INFO), |c| c.log_level);
let format = config.map_or(TraceFormat::Pretty, |c| c.log_format);

let formatter = |format| match format {
TraceFormat::Pretty => Self::from(RocketFmt::<Pretty>::new(workers, colors, level)),
TraceFormat::Compact => Self::from(RocketFmt::<Compact>::new(workers, colors, level)),
};
let formatter = Self::new(config);
if let Some(handle) = HANDLE.get() {
return assert!(handle.modify(|layer| *layer = formatter).is_ok());
}

let (layer, reload_handle) = reload::Layer::new(formatter(format));
let (layer, reload_handle) = reload::Layer::new(formatter);
let result = tracing_subscriber::registry()
.with(RequestId::layer())
.with(layer)
.try_init();

if result.is_ok() {
assert!(HANDLE.set(reload_handle).is_ok());
} else if let Some(handle) = HANDLE.get() {
assert!(handle.modify(|layer| *layer = formatter(format)).is_ok());
}
}

pub fn reset(&mut self, cli_colors: CliColors, level: Option<Level>) {
either::for_both!(&mut self.inner, f => f.reset(cli_colors, level))
}
}

macro_rules! forward {
Expand Down

0 comments on commit f2ca2ad

Please sign in to comment.