Skip to content

Instantly share code, notes, and snippets.

@Rexagon
Created August 12, 2024 20:13
Show Gist options
  • Save Rexagon/e28b18b4ef9e3645e03fe87ad39e6bc9 to your computer and use it in GitHub Desktop.
Save Rexagon/e28b18b4ef9e3645e03fe87ad39e6bc9 to your computer and use it in GitHub Desktop.
use std::io::IsTerminal;
use std::num::NonZeroUsize;
use std::path::{Path, PathBuf};
use anyhow::Result;
use serde::de::Visitor;
use serde::{Deserialize, Deserializer, Serialize};
use tracing::Subscriber;
use tracing_appender::rolling::Rotation;
use tracing_subscriber::filter::Directive;
use tracing_subscriber::{fmt, Layer};
pub struct LoggerTargets {
directives: Vec<Directive>,
}
impl LoggerTargets {
pub fn load_from<P: AsRef<Path>>(path: P) -> Result<Self> {
crate::serde_helpers::load_json_from_file(path)
}
pub fn build_subscriber(&self) -> tracing_subscriber::filter::EnvFilter {
let mut builder = tracing_subscriber::filter::EnvFilter::default();
for item in &self.directives {
builder = builder.add_directive(item.clone());
}
builder
}
}
impl<'de> Deserialize<'de> for LoggerTargets {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
struct LoggerVisitor;
impl<'de> Visitor<'de> for LoggerVisitor {
type Value = LoggerTargets;
fn expecting(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("a list of targets")
}
fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
where
A: serde::de::MapAccess<'de>,
{
let mut directives = Vec::new();
while let Some((target, level)) = map.next_entry::<String, String>()? {
let directive = format!("{}={}", target, level)
.parse::<Directive>()
.map_err(serde::de::Error::custom)?;
directives.push(directive);
}
Ok(LoggerTargets { directives })
}
}
deserializer.deserialize_map(LoggerVisitor)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoggerConfig {
pub outputs: Vec<LoggerOutput>,
}
impl Default for LoggerConfig {
fn default() -> Self {
Self {
outputs: vec![LoggerOutput::Stderr(LoggerStderrOutput)],
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum LoggerOutput {
Stderr(LoggerStderrOutput),
File(LoggerFileOutput),
}
impl LoggerOutput {
pub fn as_layer<S>(&self) -> Result<Box<dyn Layer<S> + Send + Sync + 'static>>
where
S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
{
match self {
Self::Stderr(stderr) => Ok(stderr.as_layer()),
Self::File(file) => file.as_layer::<S>(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoggerStderrOutput;
impl LoggerStderrOutput {
#[allow(clippy::unused_self)]
pub fn as_layer<S>(&self) -> Box<dyn Layer<S> + Send + Sync + 'static>
where
S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
{
if is_systemd_child() {
fmt::layer().without_time().with_ansi(false).boxed()
} else if !std::io::stdout().is_terminal() {
fmt::layer().with_ansi(false).boxed()
} else {
fmt::layer().boxed()
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoggerFileOutput {
pub dir: PathBuf,
#[serde(default)]
pub human_readable: bool,
#[serde(default = "log_file_prefix")]
pub file_prefix: String,
#[serde(default = "max_log_files")]
pub max_files: NonZeroUsize,
}
impl LoggerFileOutput {
pub fn as_layer<S>(&self) -> Result<Box<dyn Layer<S> + Send + Sync + 'static>>
where
S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>,
{
let writer = tracing_appender::rolling::Builder::new()
.rotation(Rotation::HOURLY)
.filename_prefix(&self.file_prefix)
.max_log_files(self.max_files.get())
.build(&self.dir)?;
Ok(if self.human_readable {
fmt::layer()
.without_time()
.with_ansi(false)
.with_writer(writer)
.boxed()
} else {
tracing_stackdriver::layer().with_writer(writer).boxed()
})
}
}
fn log_file_prefix() -> String {
"tycho.log".to_owned()
}
fn max_log_files() -> NonZeroUsize {
NonZeroUsize::new(25).unwrap()
}
pub fn is_systemd_child() -> bool {
#[cfg(target_os = "linux")]
unsafe {
libc::getppid() == 1
}
#[cfg(not(target_os = "linux"))]
{
false
}
}
// ...
fn init_logger(config: &LoggerConfig, logger_targets: Option<PathBuf>) -> Result<()> {
use tracing_subscriber::layer::SubscriberExt;
use tracing_subscriber::{reload, EnvFilter};
let try_make_filter = {
let logger_targets = logger_targets.clone();
move || {
Ok::<_, anyhow::Error>(match &logger_targets {
None => EnvFilter::builder()
.with_default_directive(tracing::Level::INFO.into())
.from_env_lossy(),
Some(path) => LoggerTargets::load_from(path)
.wrap_err("failed to load logger config")?
.build_subscriber(),
})
}
};
let (layer, handle) = reload::Layer::new(try_make_filter()?);
let subscriber = tracing_subscriber::registry().with(layer).with(
config
.outputs
.iter()
.map(|o| o.as_layer())
.collect::<Result<Vec<_>>>()?,
);
tracing::subscriber::set_global_default(subscriber).unwrap();
if let Some(logger_config) = logger_targets {
tokio::spawn(async move {
tracing::info!(
logger_config = %logger_config.display(),
"started watching for changes in logger config"
);
let get_metadata = move || {
std::fs::metadata(&logger_config)
.ok()
.and_then(|m| m.modified().ok())
};
let mut last_modified = get_metadata();
let mut interval = tokio::time::interval(Duration::from_secs(10));
loop {
interval.tick().await;
let modified = get_metadata();
if last_modified == modified {
continue;
}
last_modified = modified;
match try_make_filter() {
Ok(filter) => {
if handle.reload(filter).is_err() {
break;
}
tracing::info!("reloaded logger config");
}
Err(e) => tracing::error!(%e, "failed to reload logger config"),
}
}
tracing::info!("stopped watching for changes in logger config");
});
}
Ok(())
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment