From 296aa97bd842a1fa74eba352f2502ad3a5e9bf72 Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@rubenfiszel.com>
Date: Sat, 12 Oct 2024 15:19:43 +0200
Subject: [PATCH] add trace_id to all request spans

---
 backend/windmill-api/src/jobs.rs         | 21 ++++++++++-----------
 backend/windmill-api/src/tracing_init.rs | 12 ++++++++++++
 2 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/backend/windmill-api/src/jobs.rs b/backend/windmill-api/src/jobs.rs
index b7d02b040b..286709e535 100644
--- a/backend/windmill-api/src/jobs.rs
+++ b/backend/windmill-api/src/jobs.rs
@@ -11,7 +11,6 @@ use axum::http::HeaderValue;
 use quick_cache::sync::Cache;
 use serde_json::value::RawValue;
 use sqlx::Pool;
-use windmill_common::utils::rd_string;
 use std::collections::HashMap;
 #[cfg(feature = "prometheus")]
 use std::sync::atomic::Ordering;
@@ -3054,6 +3053,7 @@ pub async fn run_script_by_path_inner(
 pub struct WorkflowAsCodeQuery {
     pub skip_update: Option<bool>,
 }
+
 pub async fn run_workflow_as_code(
     authed: ApiAuthed,
     Extension(db): Extension<DB>,
@@ -3065,11 +3065,10 @@ pub async fn run_workflow_as_code(
     Json(task): Json<WorkflowTask>,
 ) -> error::Result<(StatusCode, String)> {
 
-    let rd_string = rd_string(3);
     let mut i = 1;
 
     if *CLOUD_HOSTED {
-        tracing::info!("workflow_as_code_tracing 1 id: {i} {rd_string}");
+        tracing::info!("workflow_as_code_tracing id {i} ");
         i += 1;
     }
 
@@ -3079,7 +3078,7 @@ pub async fn run_workflow_as_code(
 
 
     if *CLOUD_HOSTED {
-        tracing::info!("workflow_as_code_tracing 1 id: {i} {rd_string}");
+        tracing::info!("workflow_as_code_tracing id {i} ");
         i += 1;
     }
 
@@ -3089,7 +3088,7 @@ pub async fn run_workflow_as_code(
 
 
     if *CLOUD_HOSTED {
-        tracing::info!("workflow_as_code_tracing 1 id: {i} {rd_string}");
+        tracing::info!("workflow_as_code_tracing id {i} ");
         i += 1;
     }
 
@@ -3129,7 +3128,7 @@ pub async fn run_workflow_as_code(
 
 
     if *CLOUD_HOSTED {
-        tracing::info!("workflow_as_code_tracing 1 id: {i} {rd_string}");
+        tracing::info!("workflow_as_code_tracing id {i} ");
         i += 1;
     }
 
@@ -3143,7 +3142,7 @@ pub async fn run_workflow_as_code(
 
 
     if *CLOUD_HOSTED {
-        tracing::info!("workflow_as_code_tracing 1 id: {i} {rd_string}");
+        tracing::info!("workflow_as_code_tracing id {i} ");
         i += 1;
     }
 
@@ -3153,7 +3152,7 @@ pub async fn run_workflow_as_code(
 
 
     if *CLOUD_HOSTED {
-        tracing::info!("workflow_as_code_tracing 1 id: {i} {rd_string}");
+        tracing::info!("workflow_as_code_tracing id {i} ");
         i += 1;
     }
 
@@ -3186,7 +3185,7 @@ pub async fn run_workflow_as_code(
 
 
     if *CLOUD_HOSTED {
-        tracing::info!("workflow_as_code_tracing 1 id: {i} {rd_string}");
+        tracing::info!("workflow_as_code_tracing id {i} ");
         i += 1;
     }
 
@@ -3204,7 +3203,7 @@ pub async fn run_workflow_as_code(
 
 
     if *CLOUD_HOSTED {
-        tracing::info!("workflow_as_code_tracing 1 id: {i} {rd_string}");
+        tracing::info!("workflow_as_code_tracing id {i} ");
         i += 1;
     }
 
@@ -3213,7 +3212,7 @@ pub async fn run_workflow_as_code(
 
 
     if *CLOUD_HOSTED {
-        tracing::info!("workflow_as_code_tracing 1 id: {i} {rd_string}");
+        tracing::info!("workflow_as_code_tracing id {i} ");
     }
 
     Ok((StatusCode::CREATED, uuid.to_string()))
diff --git a/backend/windmill-api/src/tracing_init.rs b/backend/windmill-api/src/tracing_init.rs
index 1f93b941cb..60dab73810 100644
--- a/backend/windmill-api/src/tracing_init.rs
+++ b/backend/windmill-api/src/tracing_init.rs
@@ -9,6 +9,7 @@
 use ::tracing::{field, Span};
 use hyper::Response;
 use tower_http::trace::{MakeSpan, OnFailure, OnResponse};
+use uuid::Uuid;
 
 lazy_static::lazy_static! {
     static ref LOG_REQUESTS: bool = std::env::var("LOG_REQUESTS")
@@ -45,17 +46,28 @@ impl<B> OnFailure<B> for MyOnFailure {
         // tracing::error!(latency = latency.as_millis(), "response")
     }
 }
+
+lazy_static::lazy_static! {
+    static ref TRACING_HEADER: String = std::env::var("TRACING_HEADER")
+    .ok().unwrap_or_else(|| "x-tracing-id".to_string());
+}
 #[derive(Clone)]
 pub struct MyMakeSpan {}
 
 impl<B> MakeSpan<B> for MyMakeSpan {
     fn make_span(&mut self, request: &hyper::Request<B>) -> Span {
+        let tracing_id = request
+            .headers()
+            .get(TRACING_HEADER.as_str())
+            .and_then(|x| x.to_str().map(|x| x.to_string()).ok())
+            .unwrap_or(Uuid::new_v4().to_string());
         tracing::info_span!(
             "request",
             method = %request.method(),
             uri = %request.uri(),
             username = field::Empty,
             workspace_id = field::Empty,
+            trace_id = tracing_id,
             email = field::Empty,
         )
     }