Fix instrumentation yet again

worldcoin · Jul 11, 2023 · 18336a1 · 18336a1
1 parent 5e542fa
commit 18336a1
Showing 1 changed file with 58 additions and 54 deletions.
diff --git a/src/server/custom_middleware/logging_layer.rs b/src/server/custom_middleware/logging_layer.rs
@@ -6,7 +6,7 @@ use axum::response::Response;
 use bytes::Bytes;
 use hyper::body::HttpBody;
 use hyper::{Body, Method};
-use tracing::{error, info, info_span};
+use tracing::{error, info, info_span, Instrument};
 
 // 1 MiB
 const MAX_REQUEST_BODY_SIZE: u64 = 1024 * 1024;
@@ -23,68 +23,72 @@ where
     let request_query = parts.uri.query().map(ToString::to_string);
 
     if let Method::GET = request_method {
-        info_span!("request", ?uri_path, ?request_method, ?request_query)
-            .in_scope(|| async {
-                cli_batteries::trace_from_headers(&parts.headers);
-
-                info!(
-                    uri_path,
-                    ?request_method,
-                    ?request_query,
-                    "Processing request"
-                );
-
-                let body = Body::empty();
-                let request = Request::from_parts(parts, body);
-
-                let response = next.run(request).await;
-
-                let mut response = handle_response(
-                    &uri_path,
-                    &request_method,
-                    request_query.as_deref(),
-                    response,
-                )
-                .await?;
-
-                cli_batteries::trace_to_headers(response.headers_mut());
-
-                Ok(response)
-            })
-            .await
+        let span = info_span!("request", ?uri_path, ?request_method, ?request_query);
+
+        async {
+            cli_batteries::trace_from_headers(&parts.headers);
+
+            info!(
+                uri_path,
+                ?request_method,
+                ?request_query,
+                "Processing request"
+            );
+
+            let body = Body::empty();
+            let request = Request::from_parts(parts, body);
+
+            let response = next.run(request).await;
+
+            let mut response = handle_response(
+                &uri_path,
+                &request_method,
+                request_query.as_deref(),
+                response,
+            )
+            .await?;
+
+            cli_batteries::trace_to_headers(response.headers_mut());
+
+            Ok(response)
+        }
+        .instrument(span)
+        .await
     } else {
         let body = body_to_string(body).await?;
 
-        info_span!("request", ?uri_path, ?request_method, ?request_query, ?body)
-            .in_scope(|| async {
-                cli_batteries::trace_from_headers(&parts.headers);
+        let span = info_span!("request", ?uri_path, ?request_method, ?request_query, ?body);
 
-                info!(
-                    ?uri_path,
-                    ?request_method,
-                    ?request_query,
-                    ?body,
-                    "Processing request"
-                );
+        async {
+            cli_batteries::trace_from_headers(&parts.headers);
 
-                let body = Body::from(body);
-                let request = Request::from_parts(parts, body);
+            info!(
+                ?uri_path,
+                ?request_method,
+                ?request_query,
+                ?body,
+                "Processing request"
+            );
+
+            let body = Body::from(body);
+            let request = Request::from_parts(parts, body);
 
-                let response = next.run(request).await;
+            let response = next.run(request).await;
 
-                let mut response = handle_response(
-                    &uri_path,
-                    &request_method,
-                    request_query.as_deref(),
-                    response,
-                )
-                .await?;
+            let mut response = handle_response(
+                &uri_path,
+                &request_method,
+                request_query.as_deref(),
+                response,
+            )
+            .await?;
 
-                cli_batteries::trace_to_headers(response.headers_mut());
+            cli_batteries::trace_to_headers(response.headers_mut());
 
-                Ok(response)
-            })
-            .await
+            Ok(response)
+        }
+        .instrument(span)
+        .await
     }
 }