use super::{
tracing::{export, Config},
Context,
};
use crate::{Metrics, Spawner};
use axum::{
body::Body,
http::{header, Response, StatusCode},
routing::get,
serve, Extension, Router,
};
use cfg_if::cfg_if;
use std::net::SocketAddr;
use tokio::net::TcpListener;
use tracing::Level;
use tracing_subscriber::{layer::SubscriberExt, Layer, Registry};
pub struct Logging {
pub level: Level,
pub json: bool,
}
pub fn init(
context: Context,
logging: Logging,
metrics: Option<SocketAddr>,
traces: Option<Config>,
) {
let filter = tracing_subscriber::EnvFilter::new(logging.level.to_string());
let log_layer = tracing_subscriber::fmt::layer()
.with_line_number(true)
.with_thread_ids(true)
.with_file(true)
.with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE);
let log_layer = if logging.json {
log_layer.json().boxed()
} else {
log_layer.compact().boxed()
};
let trace_layer = traces.map(|cfg| {
let tracer = export(cfg).expect("Failed to initialize tracer");
tracing_opentelemetry::layer().with_tracer(tracer)
});
cfg_if! {
if #[cfg(feature = "tokio-console")] {
let console_layer = console_subscriber::spawn();
let registry = Registry::default()
.with(filter)
.with(log_layer)
.with(trace_layer)
.with(console_layer);
} else {
let registry = Registry::default()
.with(filter)
.with(log_layer)
.with(trace_layer);
}
}
tracing::subscriber::set_global_default(registry).expect("Failed to set subscriber");
if let Some(cfg) = metrics {
context
.with_label("metrics")
.spawn(move |context| async move {
let listener = TcpListener::bind(cfg)
.await
.expect("Failed to bind metrics server");
let app = Router::new()
.route(
"/metrics",
get(|Extension(ctx): Extension<Context>| async move {
Response::builder()
.status(StatusCode::OK)
.header(header::CONTENT_TYPE, "text/plain; version=0.0.4")
.body(Body::from(ctx.encode()))
.expect("Failed to create response")
}),
)
.layer(Extension(context));
serve(listener, app.into_make_service())
.await
.expect("Could not serve metrics");
});
}
}