use std::sync::atomic::{AtomicU8, Ordering};
use std::sync::OnceLock;
use tracing_core::{LevelFilter, Metadata};
use tracing_subscriber::layer::{Context, Filter};
use tracing_subscriber::registry::LookupSpan;
pub use tracing::{debug, error, info, trace, warn};
pub use tracing::{debug_span, error_span, info_span, trace_span, warn_span};
pub use tracing::Level;
const fn level_to_u8(level: Level) -> u8 {
match level {
Level::TRACE => 0,
Level::DEBUG => 1,
Level::INFO => 2,
Level::WARN => 3,
Level::ERROR => 4,
}
}
const fn u8_to_level(val: u8) -> Level {
match val {
0 => Level::TRACE,
1 => Level::DEBUG,
2 => Level::INFO,
3 => Level::WARN,
_ => Level::ERROR,
}
}
static LEVEL: OnceLock<AtomicU8> = OnceLock::new();
fn global_level() -> &'static AtomicU8 {
LEVEL.get_or_init(|| AtomicU8::new(level_to_u8(Level::WARN)))
}
pub fn set_level(level: Level) {
global_level().store(level_to_u8(level), Ordering::Relaxed);
}
pub fn current_level() -> Level {
u8_to_level(global_level().load(Ordering::Relaxed))
}
pub fn parse_level(s: &str) -> Option<Level> {
match s.to_ascii_lowercase().as_str() {
"trace" | "0" => Some(Level::TRACE),
"debug" | "1" => Some(Level::DEBUG),
"info" | "2" => Some(Level::INFO),
"warn" | "warning" | "3" => Some(Level::WARN),
"error" | "4" => Some(Level::ERROR),
_ => None,
}
}
pub struct AtomicLevelFilter;
impl<S> Filter<S> for AtomicLevelFilter
where
S: tracing::Subscriber + for<'a> LookupSpan<'a>,
{
fn enabled(&self, meta: &Metadata<'_>, _cx: &Context<'_, S>) -> bool {
let threshold = global_level().load(Ordering::Relaxed);
level_to_u8(*meta.level()) >= threshold
}
fn max_level_hint(&self) -> Option<LevelFilter> {
Some(LevelFilter::TRACE)
}
}
pub fn init() {
use tracing_subscriber::fmt;
use tracing_subscriber::prelude::*;
if let Ok(level) = std::env::var("RUST_LOG") {
if let Some(level) = parse_level(&level) {
set_level(level);
}
}
let fmt_layer = fmt::layer()
.with_writer(std::io::stderr)
.with_target(true)
.with_thread_ids(true)
.with_file(true)
.with_line_number(true);
let subscriber = tracing_subscriber::registry().with(fmt_layer.with_filter(AtomicLevelFilter));
let _ = tracing::subscriber::set_global_default(subscriber);
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn level_roundtrip() {
for level in [
Level::TRACE,
Level::DEBUG,
Level::INFO,
Level::WARN,
Level::ERROR,
] {
assert_eq!(u8_to_level(level_to_u8(level)), level);
}
}
#[test]
fn parse_level_cases() {
assert_eq!(parse_level("trace"), Some(Level::TRACE));
assert_eq!(parse_level("DEBUG"), Some(Level::DEBUG));
assert_eq!(parse_level("Info"), Some(Level::INFO));
assert_eq!(parse_level("warning"), Some(Level::WARN));
assert_eq!(parse_level("WARN"), Some(Level::WARN));
assert_eq!(parse_level("error"), Some(Level::ERROR));
assert_eq!(parse_level("0"), Some(Level::TRACE));
assert_eq!(parse_level("4"), Some(Level::ERROR));
assert_eq!(parse_level("garbage"), None);
}
#[test]
fn set_and_get_level() {
set_level(Level::DEBUG);
assert_eq!(current_level(), Level::DEBUG);
set_level(Level::WARN);
assert_eq!(current_level(), Level::WARN);
}
}