From 4e088b40f69a5ebdf7c5cdc8fd654d2f8ec91cf2 Mon Sep 17 00:00:00 2001
From: David Pedersen <david.pdrsn@gmail.com>
Date: Fri, 10 Sep 2021 19:56:00 +0200
Subject: [PATCH] Show how to configure `TraceLayer` (#314)

This isn't very obvious so makes sense to include in the example
---
 examples/tracing-aka-logging/src/main.rs | 43 +++++++++++++++++++++---
 1 file changed, 39 insertions(+), 4 deletions(-)

diff --git a/examples/tracing-aka-logging/src/main.rs b/examples/tracing-aka-logging/src/main.rs
index 52a19673..0071c419 100644
--- a/examples/tracing-aka-logging/src/main.rs
+++ b/examples/tracing-aka-logging/src/main.rs
@@ -4,9 +4,16 @@
 //! cargo run -p example-tracing-aka-logging
 //! ```
 
-use axum::{handler::get, response::Html, Router};
-use std::net::SocketAddr;
-use tower_http::trace::TraceLayer;
+use axum::{
+    body::Bytes,
+    handler::get,
+    http::{HeaderMap, Request, Response},
+    response::Html,
+    Router,
+};
+use std::{net::SocketAddr, time::Duration};
+use tower_http::{classify::ServerErrorsFailureClass, trace::TraceLayer};
+use tracing::Span;
 
 #[tokio::main]
 async fn main() {
@@ -24,8 +31,36 @@ async fn main() {
         .route("/", get(handler))
         // `TraceLayer` is provided by tower-http so you have to add that as a dependency.
         // It provides good defaults but is also very customizable.
+        //
         // See https://docs.rs/tower-http/0.1.1/tower_http/trace/index.html for more details.
-        .layer(TraceLayer::new_for_http());
+        .layer(TraceLayer::new_for_http())
+        // If you want to customize the behavior using closures here is how
+        //
+        // This is just for demonstration, you don't need to add this middleware twice
+        .layer(
+            TraceLayer::new_for_http()
+                .on_request(|_request: &Request<_>, _span: &Span| {
+                    // ...
+                })
+                .on_response(
+                    |_response: &Response<_>, _latency: Duration, _span: &Span| {
+                        // ...
+                    },
+                )
+                .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| {
+                        // ...
+                    },
+                ),
+        );
 
     // run it
     let addr = SocketAddr::from(([127, 0, 0, 1], 3000));