use std::{borrow::Cow, collections::BTreeSet};
use http::Request;
use tower_http::trace::MakeSpan;
use tracing::{Level, Span};
use tracing_subscriber::{
EnvFilter, Layer, Registry,
fmt::{self, MakeWriter},
layer::SubscriberExt,
util::SubscriberInitExt,
};
use crate::context::header_to_string;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum LoggingFormat {
Json,
Pretty,
}
#[derive(Clone, Debug)]
pub struct LoggingConfig {
service_name: String,
version: Option<String>,
environment: Option<String>,
format: LoggingFormat,
level_filter: String,
redacted_headers: BTreeSet<String>,
}
impl LoggingConfig {
pub fn production(service_name: impl Into<String>) -> Self {
Self {
service_name: service_name.into(),
version: None,
environment: None,
format: LoggingFormat::Json,
level_filter: "info".to_owned(),
redacted_headers: BTreeSet::new(),
}
}
pub fn development(service_name: impl Into<String>) -> Self {
Self::production(service_name).with_format(LoggingFormat::Pretty)
}
pub fn service_name(&self) -> &str {
&self.service_name
}
pub fn version(mut self, version: impl Into<String>) -> Self {
self.version = Some(version.into());
self
}
pub fn environment(mut self, environment: impl Into<String>) -> Self {
self.environment = Some(environment.into());
self
}
pub fn with_format(mut self, format: LoggingFormat) -> Self {
self.format = format;
self
}
pub fn level_filter(mut self, level_filter: impl Into<String>) -> Self {
self.level_filter = level_filter.into();
self
}
pub fn redact_header(mut self, header: impl AsRef<str>) -> Self {
self.redacted_headers
.insert(header.as_ref().to_ascii_lowercase());
self
}
pub fn redacts_header(&self, header: impl AsRef<str>) -> bool {
self.redacted_headers
.contains(&header.as_ref().to_ascii_lowercase())
}
pub const fn output_format(&self) -> LoggingFormat {
self.format
}
pub const fn format(&self) -> LoggingFormat {
self.format
}
pub fn service_span(&self) -> Span {
tracing::info_span!(
"service",
service.name = %self.service_name,
service.version = %self.version.as_deref().unwrap_or(""),
deployment.environment = %self.environment.as_deref().unwrap_or("")
)
}
pub fn init(&self) -> Result<(), tracing_subscriber::util::TryInitError> {
match self.format {
LoggingFormat::Json => self.subscriber_with_writer(std::io::stderr).try_init(),
LoggingFormat::Pretty => self
.pretty_subscriber_with_writer(std::io::stderr)
.try_init(),
}
}
pub fn subscriber_with_writer<W>(
&self,
writer: W,
) -> impl tracing::Subscriber + Send + Sync + 'static
where
W: for<'writer> MakeWriter<'writer> + Clone + Send + Sync + 'static,
{
let filter =
EnvFilter::try_new(&self.level_filter).unwrap_or_else(|_| EnvFilter::new("info"));
let layer = fmt::layer()
.json()
.flatten_event(true)
.with_current_span(true)
.with_span_list(false)
.with_ansi(false)
.with_target(false)
.with_writer(writer)
.with_filter(filter);
Registry::default().with(layer)
}
pub fn pretty_subscriber_with_writer<W>(
&self,
writer: W,
) -> impl tracing::Subscriber + Send + Sync + 'static
where
W: for<'writer> MakeWriter<'writer> + Clone + Send + Sync + 'static,
{
let filter =
EnvFilter::try_new(&self.level_filter).unwrap_or_else(|_| EnvFilter::new("info"));
let layer = fmt::layer()
.pretty()
.with_ansi(false)
.with_target(false)
.with_writer(writer)
.with_filter(filter);
Registry::default().with(layer)
}
}
#[derive(Clone, Debug)]
pub struct StructuredMakeSpan {
config: LoggingConfig,
route: Option<Cow<'static, str>>,
}
impl StructuredMakeSpan {
pub fn new(config: LoggingConfig) -> Self {
Self {
config,
route: None,
}
}
pub fn route(mut self, route: impl Into<Cow<'static, str>>) -> Self {
self.route = Some(route.into());
self
}
}
impl<B> MakeSpan<B> for StructuredMakeSpan {
fn make_span(&mut self, request: &Request<B>) -> Span {
let request_id = header_to_string(request.headers(), "x-request-id").unwrap_or_default();
let trace_id = header_to_string(request.headers(), "traceparent")
.and_then(|value| value.split('-').nth(1).map(str::to_owned))
.unwrap_or_default();
let route = self
.route
.as_deref()
.map(str::to_owned)
.or_else(|| {
request
.extensions()
.get::<axum::extract::MatchedPath>()
.map(|path| path.as_str().to_owned())
})
.unwrap_or_else(|| "<unknown>".to_owned());
tracing::span!(
Level::INFO,
"http.request",
service.name = %self.config.service_name,
service.version = %self.config.version.as_deref().unwrap_or(""),
deployment.environment = %self.config.environment.as_deref().unwrap_or(""),
request.id = %request_id,
trace.id = %trace_id,
http.method = %request.method(),
http.route = %route,
http.target = %request.uri(),
)
}
}