fix: make trace_id visible only if in use (#359)

* fix: make trace_id visible only if in use * cr
n0-computer · Oct 18, 2022 · c29bf85 · c29bf85
1 parent c2305fb
commit c29bf85
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 20 deletions.
diff --git a/iroh-gateway/src/error.rs b/iroh-gateway/src/error.rs
@@ -9,6 +9,7 @@ use axum::{
     response::{IntoResponse, Response},
 };
 use http::{HeaderMap, HeaderValue};
+use opentelemetry::trace::TraceId;
 use serde_json::json;
 
 use crate::constants::HEADER_X_TRACE_ID;
@@ -17,7 +18,7 @@ use crate::constants::HEADER_X_TRACE_ID;
 pub struct GatewayError {
     pub status_code: StatusCode,
     pub message: String,
-    pub trace_id: String,
+    pub trace_id: TraceId,
     pub method: Option<http::Method>,
 }
 
@@ -33,10 +34,12 @@ impl GatewayError {
 impl IntoResponse for GatewayError {
     fn into_response(self) -> Response {
         let mut headers = HeaderMap::new();
-        headers.insert(
-            &HEADER_X_TRACE_ID,
-            HeaderValue::from_str(&self.trace_id).unwrap(),
-        );
+        if self.trace_id != TraceId::INVALID {
+            headers.insert(
+                &HEADER_X_TRACE_ID,
+                HeaderValue::from_str(&self.trace_id.to_string()).unwrap(),
+            );
+        }
         match self.method {
             Some(http::Method::HEAD) => {
                 let mut rb = Response::builder().status(self.status_code);
@@ -45,12 +48,20 @@ impl IntoResponse for GatewayError {
                 rb.body(BoxBody::default()).unwrap()
             }
             _ => {
-                let body = axum::Json(json!({
-                    "code": self.status_code.as_u16(),
-                    "success": false,
-                    "message": self.message,
-                    "trace_id": self.trace_id,
-                }));
+                let body = if self.trace_id != TraceId::INVALID {
+                    axum::Json(json!({
+                        "code": self.status_code.as_u16(),
+                        "success": false,
+                        "message": self.message,
+                        "trace_id": self.trace_id.to_string(),
+                    }))
+                } else {
+                    axum::Json(json!({
+                        "code": self.status_code.as_u16(),
+                        "success": false,
+                        "message": self.message,
+                    }))
+                };
                 let mut res = body.into_response();
                 res.headers_mut().extend(headers);
                 let status = res.status_mut();

diff --git a/iroh-gateway/src/handlers.rs b/iroh-gateway/src/handlers.rs
@@ -753,7 +753,7 @@ where
         status_code,
         body: body::boxed(body),
         headers,
-        trace_id: get_current_trace_id().to_string(),
+        trace_id: get_current_trace_id(),
     })
 }
 
@@ -767,7 +767,7 @@ fn error<T: ContentLoader>(
     GatewayError {
         status_code,
         message: message.to_string(),
-        trace_id: get_current_trace_id().to_string(),
+        trace_id: get_current_trace_id(),
         method: None,
     }
 }

diff --git a/iroh-gateway/src/response.rs b/iroh-gateway/src/response.rs
@@ -3,6 +3,8 @@ use axum::{
     http::{header::*, HeaderValue, StatusCode},
     response::{IntoResponse, Redirect, Response},
 };
+use iroh_metrics::get_current_trace_id;
+use opentelemetry::trace::TraceId;
 
 use crate::constants::*;
 
@@ -107,7 +109,7 @@ pub struct GatewayResponse {
     pub status_code: StatusCode,
     pub body: BoxBody,
     pub headers: HeaderMap,
-    pub trace_id: String,
+    pub trace_id: TraceId,
 }
 
 impl IntoResponse for GatewayResponse {
@@ -120,10 +122,12 @@ impl IntoResponse for GatewayResponse {
             }
         }
         let mut rb = Response::builder().status(self.status_code);
-        self.headers.insert(
-            &HEADER_X_TRACE_ID,
-            HeaderValue::from_str(&self.trace_id).unwrap(),
-        );
+        if self.trace_id != TraceId::INVALID {
+            self.headers.insert(
+                &HEADER_X_TRACE_ID,
+                HeaderValue::from_str(&self.trace_id.to_string()).unwrap(),
+            );
+        }
         let rh = rb.headers_mut().unwrap();
         rh.extend(self.headers);
         rb.body(self.body).unwrap()
@@ -138,7 +142,7 @@ impl GatewayResponse {
             status_code,
             body: BoxBody::default(),
             headers: HeaderMap::new(),
-            trace_id: String::new(),
+            trace_id: get_current_trace_id(),
         }
     }
 
@@ -155,7 +159,7 @@ impl GatewayResponse {
             status_code: StatusCode::NOT_MODIFIED,
             body: BoxBody::default(),
             headers: HeaderMap::new(),
-            trace_id: String::new(),
+            trace_id: get_current_trace_id(),
         }
     }
 }