feat: Use query id as trace id

Cargo.toml

@@ -328,7 +328,6 @@ hashlink = "0.8"
 headers = "0.4.0"
 hex = "0.4.3"
 hickory-resolver = "0.25"
-highway = "1.1"
 hive_metastore = "0.1.0"
 hostname = "0.3.1"
 http = "1"

src/query/service/Cargo.toml

@@ -126,7 +126,6 @@ futures-util = { workspace = true }
 geozero = { workspace = true }
 headers = { workspace = true }
 hex = { workspace = true }
-highway = { workspace = true }
 http = { workspace = true }
 humantime = { workspace = true }
 indicatif = { workspace = true }

src/query/service/src/servers/flight/v1/actions/init_query_env.rs

@@ -28,7 +28,10 @@ use crate::servers::flight::v1::packets::QueryEnv;
 
 pub static INIT_QUERY_ENV: &str = "/actions/init_query_env";
 
-pub async fn init_query_env(env: QueryEnv) -> Result<()> {
+pub async fn init_query_env(mut env: QueryEnv) -> Result<()> {
+    // Update query id to make sure they are compatible.
+    env.query_id = env.query_id.replace('-', "");
+
     let mut tracking_workload_group = None;
     let mut parent_mem_stat = ParentMemStat::StaticRef(&GLOBAL_MEM_STAT);
 

src/query/service/src/servers/http/middleware/session.rs

@@ -672,8 +672,8 @@ impl<E: Endpoint> Endpoint for HTTPSessionEndpoint<E> {
         let query_id = req
             .headers()
             .get(HEADER_QUERY_ID)
-            .map(|id| id.to_str().unwrap().to_string())
-            .unwrap_or_else(|| Uuid::new_v4().to_string());
+            .map(|id| id.to_str().unwrap().replace('-', ""))
+            .unwrap_or_else(|| Uuid::now_v7().simple().to_string());
 
         let mut login_history = LoginHistory::new();
         login_history.handler = LoginHandler::HTTP;

src/query/service/src/servers/http/v1/http_query_handlers.rs

@@ -34,7 +34,6 @@ use databend_common_metrics::http::metrics_incr_http_response_errors_count;
 use databend_common_version::DATABEND_SEMVER;
 use fastrace::func_path;
 use fastrace::prelude::*;
-use highway::HighwayHash;
 use http::HeaderMap;
 use http::HeaderValue;
 use http::StatusCode;
@@ -56,6 +55,7 @@ use poem::Request;
 use poem::Route;
 use serde::Deserialize;
 use serde::Serialize;
+use uuid::Uuid;
 
 use super::query::ExecuteStateKind;
 use super::query::HttpQuery;
@@ -773,8 +773,8 @@ fn query_id_not_found(query_id: &str, node_id: &str) -> PoemError {
 }
 
 fn query_id_to_trace_id(query_id: &str) -> TraceId {
-    let [hash_high, hash_low] = highway::PortableHash::default().hash128(query_id.as_bytes());
-    TraceId(((hash_high as u128) << 64) + (hash_low as u128))
+    let uuid = Uuid::parse_str(query_id).unwrap_or_else(|_| Uuid::now_v7());
+    TraceId(uuid.as_u128())
 }
 
 /// The HTTP query endpoints are expected to be responses within 60 seconds.

src/query/service/src/servers/mysql/mysql_interactive_worker.rs

@@ -214,15 +214,20 @@ impl<W: AsyncWrite + Send + Sync + Unpin> AsyncMysqlShim<W> for InteractiveWorke
         query: &'a str,
         writer: QueryResultWriter<'a, W>,
     ) -> Result<()> {
-        let query_id = Uuid::new_v4().to_string();
+        let query_id = Uuid::now_v7();
+        // Ensure the query id shares the same representation as trace_id.
+        let query_id_str = query_id.simple().to_string();
+
         let sampled =
             thread_rng().gen_range(0..100) <= self.base.session.get_trace_sample_rate()?;
-        let root = Span::root(func_path!(), SpanContext::random().sampled(sampled))
+        let span_context =
+            SpanContext::new(TraceId(query_id.as_u128()), SpanId::default()).sampled(sampled);
+        let root = Span::root(func_path!(), span_context)
             .with_properties(|| self.base.session.to_fastrace_properties());
 
         let mut tracking_payload = ThreadTracker::new_tracking_payload();
-        tracking_payload.query_id = Some(query_id.clone());
-        tracking_payload.mem_stat = Some(MemStat::create(query_id.clone()));
+        tracking_payload.query_id = Some(query_id_str.clone());
+        tracking_payload.mem_stat = Some(MemStat::create(query_id_str.to_string()));
         let _guard = ThreadTracker::tracking(tracking_payload);
 
         ThreadTracker::tracking_future(async {
@@ -247,7 +252,7 @@ impl<W: AsyncWrite + Send + Sync + Unpin> AsyncMysqlShim<W> for InteractiveWorke
             let instant = Instant::now();
             let query_result = self
                 .base
-                .do_query(query_id, query)
+                .do_query(query_id_str, query)
                 .await
                 .map_err(|err| err.display_with_sql(query));