use std::{
fmt::Debug,
fs,
future::Future,
path::{Path, PathBuf},
sync::Arc,
};
use anyhow::{Context, Result};
use chrono::Utc;
use fastrace::collector::{Config, ConsoleReporter};
use serde::{Deserialize, Serialize};
use simd_json::{json, value::owned::Object as JsonObject};
use tokio::{self, fs::File, io::AsyncWriteExt};
use tracing::{Event, Subscriber, span};
use tracing_appender::{
non_blocking::WorkerGuard,
rolling::{RollingFileAppender, Rotation},
};
use tracing_subscriber::{
EnvFilter, Registry,
fmt::{
self, FmtContext, FormatEvent, FormatFields,
format::{JsonFields, Writer},
writer::BoxMakeWriter,
},
layer::{Layer, SubscriberExt},
registry::LookupSpan,
};
#[derive(Debug, Deserialize, Clone)]
struct LoggerConfig {
logger: LogConfig,
}
#[derive(Debug, Deserialize, Clone)]
#[serde(rename_all = "lowercase")]
enum Output {
Stdout,
Stderr,
File,
}
#[derive(Debug, Deserialize, Clone)]
#[serde(rename_all = "lowercase")]
enum RotationFreq {
Minutely,
Hourly,
Daily,
Never,
}
#[derive(Debug, Deserialize, Clone)]
struct LogFileConfig {
path: String,
#[serde(default)]
rotation_frequency: Option<RotationFreq>,
}
#[derive(Debug, Deserialize, Clone)]
struct LogConfig {
level: String,
output: Output,
is_show_line: bool,
is_show_module_path: bool,
is_show_target: bool,
file: Option<LogFileConfig>,
}
#[derive(Default, Debug)]
struct SpanFields(pub JsonObject);
struct CaptureSpanFieldsLayer;
impl<S> Layer<S> for CaptureSpanFieldsLayer
where S: Subscriber + for<'a> LookupSpan<'a>
{
fn on_new_span(
&self,
attrs: &span::Attributes<'_>,
id: &span::Id,
ctx: tracing_subscriber::layer::Context<'_, S>,
) {
if let Some(span) = ctx.span(id) {
let mut map = JsonObject::with_capacity(8);
struct V<'a>(&'a mut JsonObject);
impl<'a> tracing::field::Visit for V<'a> {
fn record_debug(&mut self, f: &tracing::field::Field, v: &dyn Debug) {
self.0
.insert(f.name().to_string(), json!(format!("{:?}", v)));
}
fn record_i64(&mut self, f: &tracing::field::Field, v: i64) {
self.0.insert(f.name().to_string(), json!(v));
}
fn record_u64(&mut self, f: &tracing::field::Field, v: u64) {
self.0.insert(f.name().to_string(), json!(v));
}
fn record_bool(&mut self, f: &tracing::field::Field, v: bool) {
self.0.insert(f.name().to_string(), json!(v));
}
fn record_str(&mut self, f: &tracing::field::Field, v: &str) {
self.0.insert(f.name().to_string(), json!(v));
}
}
let mut vis = V(&mut map);
attrs.record(&mut vis);
span.extensions_mut().insert(SpanFields(map));
}
}
fn on_record(
&self,
id: &span::Id,
values: &span::Record<'_>,
ctx: tracing_subscriber::layer::Context<'_, S>,
) {
if let Some(span) = ctx.span(id)
&& let Some(fields) = span.extensions_mut().get_mut::<SpanFields>()
{
struct V<'a>(&'a mut JsonObject);
impl<'a> tracing::field::Visit for V<'a> {
fn record_debug(&mut self, f: &tracing::field::Field, v: &dyn Debug) {
self.0
.insert(f.name().to_string(), json!(format!("{:?}", v)));
}
fn record_i64(&mut self, f: &tracing::field::Field, v: i64) {
self.0.insert(f.name().to_string(), json!(v));
}
fn record_u64(&mut self, f: &tracing::field::Field, v: u64) {
self.0.insert(f.name().to_string(), json!(v));
}
fn record_bool(&mut self, f: &tracing::field::Field, v: bool) {
self.0.insert(f.name().to_string(), json!(v));
}
fn record_str(&mut self, f: &tracing::field::Field, v: &str) {
self.0.insert(f.name().to_string(), json!(v));
}
}
let mut vis = V(&mut fields.0);
values.record(&mut vis);
}
}
}
struct JsonFormatter {
config: Arc<LogConfig>,
}
impl JsonFormatter {
fn new(config: Arc<LogConfig>) -> Self {
Self { config }
}
}
#[derive(Serialize)]
struct LogEntry {
timestamp: String,
level: String,
#[serde(skip_serializing_if = "Option::is_none")]
target: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
module_path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
line: Option<u32>,
fields: JsonObject,
}
impl<S, N> FormatEvent<S, N> for JsonFormatter
where
S: Subscriber + for<'a> LookupSpan<'a>,
N: for<'a> FormatFields<'a> + 'static,
{
fn format_event(
&self,
ctx: &FmtContext<'_, S, N>,
mut writer: Writer<'_>,
event: &Event<'_>,
) -> std::fmt::Result {
let mut visitor = JsonVisitor::default();
event.record(&mut visitor);
let mut fields = visitor.fields;
if let Some(scope) = ctx.event_scope() {
let mut span_names = Vec::with_capacity(8);
for span in scope.from_root() {
span_names.push(span.name().to_string());
if let Some(ext) = span.extensions().get::<SpanFields>() {
for (k, v) in &ext.0 {
fields.entry(k.clone()).or_insert(v.clone());
}
}
}
fields.insert("span_names".to_string(), json!(span_names));
}
let log_entry = LogEntry {
timestamp: Utc::now().to_rfc3339(),
level: event.metadata().level().to_string(),
target: if self.config.is_show_target {
Some(event.metadata().target().to_string())
} else {
None
},
module_path: if self.config.is_show_module_path {
Some(event.metadata().module_path().unwrap_or("").to_string())
} else {
None
},
line: if self.config.is_show_line {
event.metadata().line()
} else {
None
},
fields,
};
writeln!(
writer,
"{}",
simd_json::to_string(&log_entry).map_err(|_| std::fmt::Error)?
)
}
}
#[derive(Default)]
struct JsonVisitor {
fields: JsonObject,
}
impl tracing::field::Visit for JsonVisitor {
fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn Debug) {
self.fields
.insert(field.name().to_string(), json!(format!("{:?}", value)));
}
fn record_i64(&mut self, field: &tracing::field::Field, value: i64) {
self.fields.insert(field.name().to_string(), json!(value));
}
fn record_u64(&mut self, field: &tracing::field::Field, value: u64) {
self.fields.insert(field.name().to_string(), json!(value));
}
fn record_bool(&mut self, field: &tracing::field::Field, value: bool) {
self.fields.insert(field.name().to_string(), json!(value));
}
fn record_str(&mut self, field: &tracing::field::Field, value: &str) {
self.fields.insert(field.name().to_string(), json!(value));
}
}
pub fn init_logger(config_path: &str) -> anyhow::Result<WorkerGuard> {
let config_content = fs::read_to_string(config_path)
.with_context(|| format!("Failed to read config file: {config_path}"))?;
let config: LoggerConfig = serde_yaml::from_str(&config_content)
.with_context(|| format!("Failed to parse config file: {config_path}"))?;
let (writer, guard) = make_writer(&config.logger)?;
fastrace::set_reporter(ConsoleReporter, Config::default());
let compat_layer = fastrace_tracing::FastraceCompatLayer::new();
let env_filter = EnvFilter::try_new(&config.logger.level)
.or_else(|_| EnvFilter::try_from_default_env())
.context("Failed to parse log level from config or env")?;
let json_layer = fmt::layer()
.with_writer(writer)
.with_ansi(false)
.json()
.event_format(JsonFormatter::new(Arc::new(config.logger)))
.fmt_fields(JsonFields::default());
let subscriber = Registry::default()
.with(env_filter)
.with(compat_layer)
.with(CaptureSpanFieldsLayer)
.with(json_layer);
tracing::subscriber::set_global_default(subscriber)
.context("Failed to set global default subscriber")?;
Ok(guard)
}
fn make_writer(cfg: &LogConfig) -> anyhow::Result<(BoxMakeWriter, WorkerGuard)> {
Ok(match cfg.output {
Output::Stdout => {
let (w, g) = tracing_appender::non_blocking(std::io::stdout());
(BoxMakeWriter::new(w), g)
},
Output::Stderr => {
let (w, g) = tracing_appender::non_blocking(std::io::stderr());
(BoxMakeWriter::new(w), g)
},
Output::File => {
let fcfg = cfg
.file
.clone()
.context("log.file is required for output=file")?;
let path = PathBuf::from(&fcfg.path);
let dir = path.parent().unwrap_or_else(|| Path::new(""));
let rotation = match fcfg.rotation_frequency.unwrap_or(RotationFreq::Never) {
RotationFreq::Minutely => Rotation::MINUTELY,
RotationFreq::Hourly => Rotation::HOURLY,
RotationFreq::Daily => Rotation::DAILY,
RotationFreq::Never => Rotation::NEVER,
};
let file_appender = RollingFileAppender::new(
rotation,
dir,
path.file_name().unwrap_or_default(),
);
let (w, g) = tracing_appender::non_blocking(file_appender);
(BoxMakeWriter::new(w), g)
},
})
}
pub trait LoggableToFile {
fn get_name() -> &'static str {
"unknown"
}
fn save_to_file(
file_name: &str,
content: &str,
) -> impl Future<Output = Result<()>> + Send {
perform_save_to_file(file_name, content)
}
}
pub async fn perform_save_to_file<P: AsRef<Path>, C: AsRef<[u8]>>(
file_name: P,
content: C,
) -> Result<()> {
if let Some(parent_dir) = file_name.as_ref().parent() {
tokio::fs::create_dir_all(parent_dir)
.await
.context("Failed to create directory for the file")?;
}
let mut file = File::create(&file_name)
.await
.context("Failed to create file")?;
file.write_all(content.as_ref())
.await
.context("Failed to write content to file")?;
Ok(())
}