From 8cbace8605956df1263770924d51bfa2267be3c6 Mon Sep 17 00:00:00 2001
From: xfy <i@rua.plus>
Date: Thu, 19 Sep 2024 22:56:29 +0800
Subject: [PATCH] feat(axum): update log format

---
 web/Rust/axum/src/main.rs            |  1 -
 web/Rust/axum/src/middlewares/mod.rs | 68 +++++++++++++++++-----------
 2 files changed, 41 insertions(+), 28 deletions(-)

diff --git a/web/Rust/axum/src/main.rs b/web/Rust/axum/src/main.rs
index 4f76f6e..0b13256 100644
--- a/web/Rust/axum/src/main.rs
+++ b/web/Rust/axum/src/main.rs
@@ -20,7 +20,6 @@ type Result<T> = std::result::Result<T, Box<dyn Error>>;
 async fn main() -> Result<()> {
     dotenv().ok();
     init_logger();
-    info!("Hello, world!");
 
     let port = env::var("PHTHONUS_PORT")
         .map(|port| port.parse::<u16>().unwrap_or(DEFAULT_PORT))
diff --git a/web/Rust/axum/src/middlewares/mod.rs b/web/Rust/axum/src/middlewares/mod.rs
index 3880d6a..8405f6d 100644
--- a/web/Rust/axum/src/middlewares/mod.rs
+++ b/web/Rust/axum/src/middlewares/mod.rs
@@ -1,4 +1,4 @@
-use std::time::Duration;
+use std::{fmt::Display, time::Duration};
 
 use axum::{
     body::Bytes,
@@ -40,30 +40,44 @@ pub async fn add_version(
 /// This middleware will calculate each request latency
 /// and add request's information to each info_span.
 pub fn logging_route(router: Router) -> Router {
-    router.layer(
-        TraceLayer::new_for_http()
-            .make_span_with(|req: &Request<_>| {
-                let unknown = &HeaderValue::from_static("Unknown");
-                let empty = &HeaderValue::from_static("");
-                let headers = req.headers();
-                let ua = headers
-                    .get("User-Agent")
-                    .unwrap_or(unknown)
-                    .to_str()
-                    .unwrap_or("Unknown");
-                let host = headers.get("Host").unwrap_or(empty).to_str().unwrap_or("");
-                info_span!("HTTP", method = ?req.method(), host, uri = ?req.uri(), ua)
-            })
-            .on_request(|_req: &Request<_>, _span: &Span| {})
-            .on_response(|res: &Response, latency: Duration, _span: &Span| {
-                info!("{} {}μs", res.status(), latency.as_micros());
-            })
-            .on_body_chunk(|_chunk: &Bytes, _latency: Duration, _span: &Span| {})
-            .on_eos(|_trailers: Option<&HeaderMap>, _stream_duration: Duration, _span: &Span| {})
-            .on_failure(
-                |error: ServerErrorsFailureClass, _latency: Duration, _span: &Span| {
-                    error!("{}", error);
-                },
-            ),
-    )
+    let make_span = |req: &Request<_>| {
+        let unknown = &HeaderValue::from_static("Unknown");
+        let empty = &HeaderValue::from_static("");
+        let headers = req.headers();
+        let ua = headers
+            .get("User-Agent")
+            .unwrap_or(unknown)
+            .to_str()
+            .unwrap_or("Unknown");
+        let host = headers.get("Host").unwrap_or(empty).to_str().unwrap_or("");
+        info_span!("HTTP", method = ?req.method(), host, uri = ?req.uri(), ua)
+    };
+
+    let trace_layer = TraceLayer::new_for_http()
+        .make_span_with(make_span)
+        .on_request(|_req: &Request<_>, _span: &Span| {})
+        .on_response(|res: &Response, latency: Duration, _span: &Span| {
+            info!("{}", format_latency(latency, res.status()));
+        })
+        .on_body_chunk(|_chunk: &Bytes, _latency: Duration, _span: &Span| {})
+        .on_eos(|_trailers: Option<&HeaderMap>, _stream_duration: Duration, _span: &Span| {})
+        .on_failure(
+            |error: ServerErrorsFailureClass, latency: Duration, _span: &Span| {
+                error!("{}", format_latency(latency, error));
+            },
+        );
+
+    router.layer(trace_layer)
+}
+
+/// Format request latency and status message
+/// return a string
+fn format_latency(latency: Duration, status: impl Display) -> String {
+    let micros = latency.as_micros();
+    let millis = latency.as_millis();
+    if micros >= 1000 {
+        format!("{} {}ms", status, millis)
+    } else {
+        format!("{} {}μs", status, micros)
+    }
 }