use std::sync::{Arc, OnceLock};
use parking_lot::RwLock;
use tracing_subscriber::reload;
use tracing_subscriber::EnvFilter;
use super::config::LogLevel;
use crate::error::Result;
pub type ReloadHandle = reload::Handle<EnvFilter, tracing_subscriber::Registry>;
static GLOBAL_CONTROLLER: OnceLock<LogLevelController> = OnceLock::new();
#[derive(Clone)]
pub struct LogLevelController {
inner: Arc<LogLevelControllerInner>,
}
struct LogLevelControllerInner {
handle: ReloadHandle,
current_level: RwLock<LogLevel>,
module_levels: RwLock<std::collections::HashMap<String, LogLevel>>,
}
impl LogLevelController {
pub fn new(handle: ReloadHandle, initial_level: LogLevel) -> Self {
Self {
inner: Arc::new(LogLevelControllerInner {
handle,
current_level: RwLock::new(initial_level),
module_levels: RwLock::new(std::collections::HashMap::new()),
}),
}
}
pub fn register_global(self) -> bool {
GLOBAL_CONTROLLER.set(self).is_ok()
}
pub fn global() -> Option<&'static LogLevelController> {
GLOBAL_CONTROLLER.get()
}
pub fn current_level(&self) -> LogLevel {
*self.inner.current_level.read()
}
pub fn module_levels(&self) -> std::collections::HashMap<String, LogLevel> {
self.inner.module_levels.read().clone()
}
pub fn set_level(&self, level: LogLevel) -> Result<()> {
{
let mut current = self.inner.current_level.write();
*current = level;
}
self.apply_filter()
}
pub fn set_module_level(&self, module: impl Into<String>, level: LogLevel) -> Result<()> {
{
let mut levels = self.inner.module_levels.write();
levels.insert(module.into(), level);
}
self.apply_filter()
}
pub fn remove_module_level(&self, module: &str) -> Result<()> {
{
let mut levels = self.inner.module_levels.write();
levels.remove(module);
}
self.apply_filter()
}
pub fn clear_module_levels(&self) -> Result<()> {
{
let mut levels = self.inner.module_levels.write();
levels.clear();
}
self.apply_filter()
}
pub fn enable_debug_mode(&self) -> Result<DebugModeGuard> {
let previous_level = self.current_level();
let previous_modules = self.module_levels();
self.set_level(LogLevel::Debug)?;
Ok(DebugModeGuard {
controller: self.clone(),
previous_level,
previous_modules,
})
}
pub fn enable_module_trace(&self, module: impl Into<String>) -> Result<ModuleTraceGuard> {
let module = module.into();
let previous_level = self.inner.module_levels.read().get(&module).copied();
self.set_module_level(module.clone(), LogLevel::Trace)?;
Ok(ModuleTraceGuard {
controller: self.clone(),
module,
previous_level,
})
}
fn apply_filter(&self) -> Result<()> {
let filter = self.build_filter();
self.inner
.handle
.reload(filter)
.map_err(|e| crate::error::Error::Config(format!("Failed to reload log filter: {}", e)))
}
fn build_filter(&self) -> EnvFilter {
let level = *self.inner.current_level.read();
let modules = self.inner.module_levels.read();
let mut filter_str = format!("trap_sim={}", level.as_filter_str());
for (module, module_level) in modules.iter() {
filter_str.push_str(&format!(",{}={}", module, module_level.as_filter_str()));
}
if !modules.contains_key("tokio") {
filter_str.push_str(",tokio=warn");
}
if !modules.contains_key("hyper") {
filter_str.push_str(",hyper=warn");
}
EnvFilter::try_new(&filter_str).unwrap_or_else(|_| EnvFilter::new(level.as_filter_str()))
}
pub fn current_filter_string(&self) -> String {
let level = *self.inner.current_level.read();
let modules = self.inner.module_levels.read();
let mut parts = vec![format!("trap_sim={}", level.as_filter_str())];
for (module, module_level) in modules.iter() {
parts.push(format!("{}={}", module, module_level.as_filter_str()));
}
parts.join(",")
}
}
impl std::fmt::Debug for LogLevelController {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("LogLevelController")
.field("current_level", &self.current_level())
.field("module_levels", &self.module_levels())
.finish()
}
}
#[must_use = "Debug mode will be reverted when the guard is dropped"]
pub struct DebugModeGuard {
controller: LogLevelController,
previous_level: LogLevel,
previous_modules: std::collections::HashMap<String, LogLevel>,
}
impl Drop for DebugModeGuard {
fn drop(&mut self) {
let _ = self.controller.set_level(self.previous_level);
{
let mut levels = self.controller.inner.module_levels.write();
*levels = self.previous_modules.clone();
}
let _ = self.controller.apply_filter();
}
}
#[must_use = "Module trace will be disabled when the guard is dropped"]
pub struct ModuleTraceGuard {
controller: LogLevelController,
module: String,
previous_level: Option<LogLevel>,
}
impl Drop for ModuleTraceGuard {
fn drop(&mut self) {
let _ = match self.previous_level {
Some(level) => self.controller.set_module_level(&self.module, level),
None => self.controller.remove_module_level(&self.module),
};
}
}
#[derive(Debug, Clone)]
pub struct LevelChangeEvent {
pub change_type: LevelChangeType,
pub previous_level: Option<LogLevel>,
pub new_level: LogLevel,
pub timestamp: std::time::SystemTime,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LevelChangeType {
Global,
Module(String),
}
#[derive(Debug, Clone, Default)]
pub struct LevelChangeStats {
pub total_changes: u64,
pub global_changes: u64,
pub module_changes: u64,
pub last_change: Option<std::time::SystemTime>,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_level_change_event() {
let event = LevelChangeEvent {
change_type: LevelChangeType::Global,
previous_level: Some(LogLevel::Info),
new_level: LogLevel::Debug,
timestamp: std::time::SystemTime::now(),
};
assert_eq!(event.change_type, LevelChangeType::Global);
assert_eq!(event.previous_level, Some(LogLevel::Info));
assert_eq!(event.new_level, LogLevel::Debug);
}
#[test]
fn test_level_change_type() {
let global = LevelChangeType::Global;
let module = LevelChangeType::Module("test::module".to_string());
assert_ne!(global, module);
}
#[test]
fn test_level_change_stats() {
let mut stats = LevelChangeStats::default();
stats.total_changes = 5;
stats.global_changes = 2;
stats.module_changes = 3;
assert_eq!(stats.total_changes, 5);
assert_eq!(stats.global_changes + stats.module_changes, 5);
}
}