From 2ef87f3511113fdb31df3e69592d687f21b9b2bb Mon Sep 17 00:00:00 2001
From: nazeh <ar.nazeh@gmail.com>
Date: Wed, 22 Jan 2025 21:54:50 +0300
Subject: [PATCH] feat(homeserver): custom tracing layer exclude noisy /events/
 #64

---
 pubky-homeserver/src/core/extractors.rs   |  2 -
 pubky-homeserver/src/core/layers/mod.rs   |  1 +
 pubky-homeserver/src/core/layers/trace.rs | 60 +++++++++++++++++++++++
 pubky-homeserver/src/core/routes/mod.rs   | 13 +++--
 4 files changed, 70 insertions(+), 6 deletions(-)
 create mode 100644 pubky-homeserver/src/core/layers/trace.rs

diff --git a/pubky-homeserver/src/core/extractors.rs b/pubky-homeserver/src/core/extractors.rs
index 027cb2a..085ae4e 100644
--- a/pubky-homeserver/src/core/extractors.rs
+++ b/pubky-homeserver/src/core/extractors.rs
@@ -39,8 +39,6 @@ where
             ))
             .map_err(|e| e.into_response())?;
 
-        tracing::debug!(pubky_host = ?pubky_host.public_key().to_string());
-
         Ok(pubky_host)
     }
 }
diff --git a/pubky-homeserver/src/core/layers/mod.rs b/pubky-homeserver/src/core/layers/mod.rs
index 585095b..cceacd2 100644
--- a/pubky-homeserver/src/core/layers/mod.rs
+++ b/pubky-homeserver/src/core/layers/mod.rs
@@ -1,2 +1,3 @@
 pub mod authz;
 pub mod pubky_host;
+pub mod trace;
diff --git a/pubky-homeserver/src/core/layers/trace.rs b/pubky-homeserver/src/core/layers/trace.rs
new file mode 100644
index 0000000..c1433f5
--- /dev/null
+++ b/pubky-homeserver/src/core/layers/trace.rs
@@ -0,0 +1,60 @@
+use std::sync::Arc;
+
+use axum::{extract::Request, Router};
+use tower_http::trace::{
+    DefaultMakeSpan, DefaultOnFailure, DefaultOnRequest, DefaultOnResponse, MakeSpan, OnFailure,
+    OnRequest, OnResponse, TraceLayer,
+};
+use tracing::{Level, Span};
+
+pub fn with_trace_layer(router: Router, excluded_paths: &[&str]) -> Router {
+    let excluded_paths = Arc::new(
+        excluded_paths
+            .iter()
+            .map(|s| s.to_string())
+            .collect::<Vec<_>>(),
+    );
+
+    router.layer(
+        TraceLayer::new_for_http()
+            .make_span_with(move |request: &Request| {
+                if excluded_paths.contains(&request.uri().path().to_string()) {
+                    // Skip logging for the noisy endpoint
+                    tracing::span!(Level::INFO, "request", excluded = true)
+                } else {
+                    // Use the default span for other endpoints
+                    DefaultMakeSpan::new().make_span(request)
+                }
+            })
+            .on_request(|request: &Request, span: &Span| {
+                // Skip logging for excluded spans
+                if span.has_field("excluded") {
+                    return;
+                }
+                // Use the default behavior for other spans
+                DefaultOnRequest::new().on_request(request, span);
+            })
+            .on_response(
+                |response: &axum::response::Response, latency: std::time::Duration, span: &Span| {
+                    // Skip logging for excluded spans
+                    if span.has_field("excluded") {
+                        return;
+                    }
+                    // Use the default behavior for other spans
+                    DefaultOnResponse::new().on_response(response, latency, span);
+                },
+            )
+            .on_failure(
+                |error: tower_http::classify::ServerErrorsFailureClass,
+                 latency: std::time::Duration,
+                 span: &Span| {
+                    // Skip logging for excluded spans
+                    if span.has_field("excluded") {
+                        return;
+                    }
+                    // Use the default behavior for other spans
+                    DefaultOnFailure::new().on_failure(error, latency, span);
+                },
+            ),
+    )
+}
diff --git a/pubky-homeserver/src/core/routes/mod.rs b/pubky-homeserver/src/core/routes/mod.rs
index 2d7e5da..59e7d78 100644
--- a/pubky-homeserver/src/core/routes/mod.rs
+++ b/pubky-homeserver/src/core/routes/mod.rs
@@ -5,15 +5,19 @@ use axum::{
     Router,
 };
 use tower_cookies::CookieManagerLayer;
-use tower_http::{cors::CorsLayer, trace::TraceLayer};
+use tower_http::cors::CorsLayer;
 
 use crate::core::AppState;
 
+use super::layers::trace::with_trace_layer;
+
 mod auth;
 mod feed;
 mod root;
 mod tenants;
 
+const TRACING_EXCLUDED_PATHS: [&str; 1] = ["/events/"];
+
 fn base() -> Router<AppState> {
     Router::new()
         .route("/", get(root::handler))
@@ -27,10 +31,11 @@ fn base() -> Router<AppState> {
 }
 
 pub fn create_app(state: AppState) -> Router {
-    base()
+    let app = base()
         .merge(tenants::router(state.clone()))
         .layer(CookieManagerLayer::new())
         .layer(CorsLayer::very_permissive())
-        .layer(TraceLayer::new_for_http())
-        .with_state(state)
+        .with_state(state);
+
+    with_trace_layer(app, &TRACING_EXCLUDED_PATHS)
 }