[server] Add logging #766

agdb_server/Cargo.toml

@@ -15,8 +15,13 @@ agdb = { version = "0.5.1", path = "../agdb" }
 anyhow = "1.0.75"
 axum = "0.6.20"
 hyper = "0.14.27"
+serde = { version = "1.0.192", features = ["derive"] }
+serde_json = "1.0.108"
 tokio = { version = "1.34.0", features = ["full"] }
 tower = "0.4.13"
+tower-http = { version = "0.4.4", features = ["map-request-body"] }
+tracing = "0.1.40"
+tracing-subscriber = "0.3.18"
 
 [dev-dependencies]
 assert_cmd = "2.0.12"

agdb_server/src/app.rs

@@ -1,20 +1,35 @@
+use crate::logger;
+use axum::body;
 use axum::extract::State;
+use axum::middleware;
 use axum::routing;
 use axum::Router;
 use hyper::StatusCode;
 use tokio::sync::broadcast::Sender;
+use tower::ServiceBuilder;
+use tower_http::map_request_body::MapRequestBodyLayer;
 
 pub(crate) fn app(shutdown_sender: Sender<()>) -> Router {
-    Router::new().route("/", routing::get(root)).route(
-        "/shutdown",
-        routing::get(shutdown).with_state(shutdown_sender),
-    )
+    let logger = ServiceBuilder::new()
+        .layer(MapRequestBodyLayer::new(body::boxed))
+        .layer(middleware::from_fn(logger::logger));
+
+    Router::new()
+        .route("/", routing::get(root))
+        .route("/error", routing::get(error))
+        .route("/shutdown", routing::get(shutdown))
+        .layer(logger)
+        .with_state(shutdown_sender)
 }
 
 async fn root() -> &'static str {
     "Hello, World!"
 }
 
+async fn error() -> StatusCode {
+    StatusCode::INTERNAL_SERVER_ERROR
+}
+
 async fn shutdown(State(shutdown_sender): State<Sender<()>>) -> StatusCode {
     match shutdown_sender.send(()) {
         Ok(_) => StatusCode::OK,
@@ -43,6 +58,16 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn error() -> anyhow::Result<()> {
+        let app = app(Sender::<()>::new(1));
+        let request = Request::builder().uri("/error").body(Body::empty())?;
+        let response = app.oneshot(request).await?;
+
+        assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
+        Ok(())
+    }
+
     #[tokio::test]
     async fn shutdown() -> anyhow::Result<()> {
         let (shutdown_sender, _shutdown_receiver) = tokio::sync::broadcast::channel::<()>(1);
@@ -51,7 +76,6 @@ mod tests {
         let response = app.oneshot(request).await?;
 
         assert_eq!(response.status(), StatusCode::OK);
-
         Ok(())
     }
 
@@ -62,7 +86,6 @@ mod tests {
         let response = app.oneshot(request).await?;
 
         assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
-
         Ok(())
     }
 }

agdb_server/src/logger.rs

@@ -0,0 +1,96 @@
+use axum::body::BoxBody;
+use axum::body::Full;
+use axum::middleware::Next;
+use axum::response::IntoResponse;
+use axum::response::Response;
+use hyper::Request;
+use hyper::StatusCode;
+use serde::Serialize;
+use std::collections::HashMap;
+use std::time::Instant;
+
+#[derive(Default, Serialize)]
+struct LogRecord {
+    method: String,
+    uri: String,
+    version: String,
+    request_headers: HashMap<String, String>,
+    request_body: String,
+    status: u16,
+    time: u128,
+    response_headers: HashMap<String, String>,
+    response: String,
+}
+
+pub(crate) async fn logger(
+    request: Request<BoxBody>,
+    next: Next<BoxBody>,
+) -> Result<impl IntoResponse, Response> {
+    let mut log_record = LogRecord::default();
+    let request = request_log(request, &mut log_record).await?;
+    let now = Instant::now();
+    let response = next.run(request).await;
+    log_record.time = now.elapsed().as_micros();
+    let response = response_log(response, &mut log_record).await;
+    print_log(log_record);
+    response
+}
+
+fn print_log(log_record: LogRecord) {
+    let message = serde_json::to_string(&log_record).unwrap_or_default();
+
+    if log_record.status < 500 {
+        tracing::info!(message);
+    } else {
+        tracing::error!(message);
+    }
+}
+
+async fn request_log(
+    request: Request<BoxBody>,
+    log_record: &mut LogRecord,
+) -> Result<Request<BoxBody>, Response> {
+    log_record.method = request.method().to_string();
+    log_record.uri = request.uri().to_string();
+    log_record.version = format!("{:?}", request.version());
+    log_record.request_headers = request
+        .headers()
+        .iter()
+        .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string()))
+        .collect();
+
+    let (parts, body) = request.into_parts();
+    let bytes = hyper::body::to_bytes(body).await;
+    let response = bytes
+        .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()).into_response())?;
+
+    log_record.request_body = String::from_utf8_lossy(&response).to_string();
+
+    Ok(Request::from_parts(
+        parts,
+        axum::body::boxed(Full::from(response)),
+    ))
+}
+
+async fn response_log(
+    response: Response<BoxBody>,
+    log_record: &mut LogRecord,
+) -> Result<impl IntoResponse, Response> {
+    log_record.status = response.status().as_u16();
+    log_record.response_headers = response
+        .headers()
+        .iter()
+        .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string()))
+        .collect();
+
+    let (parts, body) = response.into_parts();
+    let bytes = hyper::body::to_bytes(body).await;
+    let resposne = bytes
+        .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()).into_response())?;
+    log_record.response = String::from_utf8_lossy(&resposne).to_string();
+
+    Ok(Response::from_parts(
+        parts,
+        axum::body::boxed(Full::from(resposne)),
+    ))
+}

agdb_server/src/main.rs

@@ -1,9 +1,11 @@
 mod app;
+mod logger;
 
 use axum::Server;
 use std::net::SocketAddr;
 use tokio::signal;
 use tokio::sync::broadcast::Receiver;
+use tracing::Level;
 
 async fn shutdown_signal(mut shutdown_shutdown: Receiver<()>) {
     tokio::select! {
@@ -14,8 +16,10 @@ async fn shutdown_signal(mut shutdown_shutdown: Receiver<()>) {
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
+    tracing_subscriber::fmt().with_max_level(Level::INFO).init();
+
     let (shutdown_sender, shutdown_receiver) = tokio::sync::broadcast::channel::<()>(1);
-    let app = app::app(shutdown_sender.clone());
+    let app = app::app(shutdown_sender);
     let addr = SocketAddr::from(([127, 0, 0, 1], 3000));
     let shutdown = shutdown_signal(shutdown_receiver);