fix(proxy): Add testodrome query id HTTP header (#11167)

Handle "X-Neon-Query-ID" header to glue data with testodrome queries.

Relates to the #22486

proxy/src/compute.rs

@@ -290,7 +290,7 @@ impl ConnCfg {
             "connected to compute node at {host} ({socket_addr}) sslmode={:?}, latency={}, query_id={}",
             self.0.get_ssl_mode(),
             ctx.get_proxy_latency(),
-            ctx.get_testodrome_id(),
+            ctx.get_testodrome_id().unwrap_or_default(),
         );
 
         // NB: CancelToken is supposed to hold socket_addr, but we use connect_raw.

proxy/src/context/mod.rs

@@ -272,6 +272,13 @@ impl RequestContext {
             .set_user_agent(user_agent);
     }
 
+    pub(crate) fn set_testodrome_id(&self, query_id: String) {
+        self.0
+            .try_lock()
+            .expect("should not deadlock")
+            .set_testodrome_id(query_id);
+    }
+
     pub(crate) fn set_auth_method(&self, auth_method: AuthMethod) {
         let mut this = self.0.try_lock().expect("should not deadlock");
         this.auth_method = Some(auth_method);
@@ -371,13 +378,12 @@ impl RequestContext {
             .accumulated()
     }
 
-    pub(crate) fn get_testodrome_id(&self) -> String {
+    pub(crate) fn get_testodrome_id(&self) -> Option<String> {
         self.0
             .try_lock()
             .expect("should not deadlock")
             .testodrome_query_id
             .clone()
-            .unwrap_or_default()
     }
 
     pub(crate) fn success(&self) {

proxy/src/serverless/backend.rs

@@ -571,6 +571,11 @@ impl ConnectMechanism for TokioMechanism {
             "compute_id",
             tracing::field::display(&node_info.aux.compute_id),
         );
+
+        if let Some(query_id) = ctx.get_testodrome_id() {
+            info!("latency={}, query_id={}", ctx.get_proxy_latency(), query_id);
+        }
+
         Ok(poll_client(
             self.pool.clone(),
             ctx,
@@ -628,6 +633,10 @@ impl ConnectMechanism for HyperMechanism {
             tracing::field::display(&node_info.aux.compute_id),
         );
 
+        if let Some(query_id) = ctx.get_testodrome_id() {
+            info!("latency={}, query_id={}", ctx.get_proxy_latency(), query_id);
+        }
+
         Ok(poll_http2_client(
             self.pool.clone(),
             ctx,

proxy/src/serverless/mod.rs

@@ -446,6 +446,15 @@ async fn request_handler(
                 .map(Into::into),
         );
 
+        let testodrome_id = request
+            .headers()
+            .get("X-Neon-Query-ID")
+            .map(|value| value.to_str().unwrap_or_default().to_string());
+
+        if let Some(query_id) = testodrome_id {
+            ctx.set_testodrome_id(query_id);
+        }
+
         let span = ctx.span();
         info!(parent: &span, "performing websocket upgrade");