From b2a8d5d44c42516b6c82dcd640477b790fe95e57 Mon Sep 17 00:00:00 2001 From: Piotr Orzechowski Date: Wed, 21 Feb 2024 05:55:04 +0100 Subject: [PATCH] Add request_id module --- src/lib.rs | 1 + src/request_id.rs | 34 ++++++++++++++++++++++++++++++++++ src/startup.rs | 16 +++++++++++++--- src/telemetry.rs | 41 +++++------------------------------------ 4 files changed, 53 insertions(+), 39 deletions(-) create mode 100644 src/request_id.rs diff --git a/src/lib.rs b/src/lib.rs index 9079ef9..24ea30d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,5 @@ pub mod configuration; +pub mod request_id; pub mod routes; pub mod startup; pub mod telemetry; diff --git a/src/request_id.rs b/src/request_id.rs new file mode 100644 index 0000000..13c93c4 --- /dev/null +++ b/src/request_id.rs @@ -0,0 +1,34 @@ +use axum::{ + body::Body, + http::{HeaderName, HeaderValue, Request}, +}; +use tower_http::request_id::{MakeRequestId, RequestId}; +use uuid::Uuid; + +#[derive(Clone)] +pub struct RequestUuid; + +impl MakeRequestId for RequestUuid { + fn make_request_id(&mut self, _: &Request) -> Option { + match HeaderValue::from_str(&Uuid::new_v4().to_string()) { + Ok(value) => Some(RequestId::new(value)), + Err(e) => { + tracing::warn!("Failed to create request id header value: {e:?}"); + None + } + } + } +} + +pub fn from_x_request_id(request: &Request) -> Option<&str> { + request + .headers() + .get(HeaderName::from_static("x-request-id")) + .and_then(|value| match value.to_str() { + Ok(value) => Some(value), + Err(e) => { + tracing::warn!("Failed to convert x-request-id to str: {e:?}"); + None + } + }) +} diff --git a/src/startup.rs b/src/startup.rs index 9ccfb40..e6c95c5 100644 --- a/src/startup.rs +++ b/src/startup.rs @@ -1,12 +1,17 @@ use crate::{ + request_id::RequestUuid, routes::{health_check, subscriptions}, - telemetry::{request_span, RequestUuid}, + telemetry::request_span, }; use axum::Router; use sqlx::PgPool; use tokio::net::TcpListener; use tower::ServiceBuilder; -use tower_http::{trace::TraceLayer, ServiceBuilderExt}; +use tower_http::{ + trace::{DefaultOnRequest, DefaultOnResponse, TraceLayer}, + ServiceBuilderExt, +}; +use tracing::Level; pub async fn run(listener: TcpListener, db_pool: PgPool) -> Result<(), std::io::Error> { let app = Router::new() @@ -16,7 +21,12 @@ pub async fn run(listener: TcpListener, db_pool: PgPool) -> Result<(), std::io:: .layer( ServiceBuilder::new() .set_x_request_id(RequestUuid) - .layer(TraceLayer::new_for_http().make_span_with(request_span)) + .layer( + TraceLayer::new_for_http() + .make_span_with(request_span) + .on_request(DefaultOnRequest::new().level(Level::INFO)) + .on_response(DefaultOnResponse::new().level(Level::INFO)), + ) .propagate_x_request_id(), ); diff --git a/src/telemetry.rs b/src/telemetry.rs index de8f0c3..57adc10 100644 --- a/src/telemetry.rs +++ b/src/telemetry.rs @@ -1,13 +1,9 @@ -use axum::{ - body::Body, - http::{HeaderName, HeaderValue, Request}, -}; -use tower_http::request_id::{MakeRequestId, RequestId}; +use crate::request_id::from_x_request_id; +use axum::{body::Body, http::Request}; use tracing::{subscriber::set_global_default, Span, Subscriber}; use tracing_bunyan_formatter::{BunyanFormattingLayer, JsonStorageLayer}; use tracing_log::LogTracer; use tracing_subscriber::{fmt::MakeWriter, layer::SubscriberExt, Registry}; -use uuid::Uuid; pub fn get_subscriber( name: String, @@ -31,38 +27,11 @@ pub fn init_subscriber(subscriber: impl Subscriber + Send + Sync) { set_global_default(subscriber).expect("Failed to set subscriber"); } -#[derive(Clone)] -pub struct RequestUuid; - -impl MakeRequestId for RequestUuid { - fn make_request_id(&mut self, _: &Request) -> Option { - match HeaderValue::from_str(&Uuid::new_v4().to_string()) { - Ok(value) => Some(RequestId::new(value)), - Err(e) => { - tracing::warn!("Failed to create request id header value: {e:?}"); - None - } - } - } -} - pub fn request_span(request: &Request) -> Span { - let request_id = request - .headers() - .get(HeaderName::from_static("x-request-id")) - .and_then(|value| match value.to_str() { - Ok(value) => Some(value), - Err(e) => { - tracing::warn!("Failed to convert x-request-id to str: {e:?}"); - None - } - }); - tracing::info_span!( "Request", - request_id = request_id, - method = request.method().to_string(), - path = request.uri().path(), - query = request.uri().query() + request_id = from_x_request_id(request), + method = %request.method(), + uri = %request.uri(), ) }