Add request_id module

0rzech · Feb 21, 2024 · b2a8d5d · b2a8d5d
1 parent f6f18bf
commit b2a8d5d
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 39 deletions.
diff --git 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
@@ -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<B>(&mut self, _: &Request<B>) -> Option<RequestId> {
+        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<Body>) -> 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
@@ -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
@@ -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<Sink>(
     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<B>(&mut self, _: &Request<B>) -> Option<RequestId> {
-        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<Body>) -> 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(),
     )
 }