apollo_infra: cross component tracing lite

victorkstarkware · victorkstarkware · commit a1763bfe3dfa · 2026-01-13T09:18:38.000+02:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/apollo_infra/Cargo.toml b/crates/apollo_infra/Cargo.toml
@@ -18,8 +18,10 @@ apollo_config.workspace = true
 apollo_infra_utils.workspace = true
 apollo_metrics.workspace = true
 async-trait.workspace = true
+derive_more.workspace = true
 hyper = { workspace = true, features = ["client", "http2", "server", "tcp"] }
 metrics-exporter-prometheus.workspace = true
+rand.workspace = true
 rstest.workspace = true
 serde = { workspace = true, features = ["derive"] }
 serde_json.workspace = true
diff --git a/crates/apollo_infra/src/component_client/local_component_client.rs b/crates/apollo_infra/src/component_client/local_component_client.rs
@@ -3,9 +3,11 @@ use serde::de::DeserializeOwned;
 use serde::Serialize;
 use tokio::sync::mpsc::{channel, Sender};
 use tokio::time::Instant;
+use tracing::field::{display, Empty};
+use tracing::instrument;
 
 use crate::component_client::ClientResult;
-use crate::component_definitions::{ComponentClient, RequestWrapper};
+use crate::component_definitions::{ComponentClient, RequestId, RequestWrapper};
 use crate::metrics::LocalClientMetrics;
 use crate::requests::LabeledRequest;
 
@@ -40,10 +42,13 @@ where
     Request: Send + Serialize + DeserializeOwned + LabeledRequest,
     Response: Send + Serialize + DeserializeOwned,
 {
+    #[instrument(skip_all, fields(request_id = Empty))]
     async fn send(&self, request: Request) -> ClientResult<Response> {
+        let request_id = RequestId::generate();
+        tracing::Span::current().record("request_id", display(&request_id));
         let request_label = request.request_label();
         let (res_tx, mut res_rx) = channel::<Response>(1);
-        let request_wrapper = RequestWrapper::new(request, res_tx);
+        let request_wrapper = RequestWrapper::new(request, res_tx, request_id);
         let start = Instant::now();
         self.tx.send(request_wrapper).await.expect("Outbound connection should be open.");
         let response = res_rx.recv().await.expect("Inbound connection should be open.");
diff --git a/crates/apollo_infra/src/component_client/remote_component_client.rs b/crates/apollo_infra/src/component_client/remote_component_client.rs
@@ -14,11 +14,18 @@ use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 use tokio::sync::Mutex;
 use tokio::time::Instant;
-use tracing::{debug, trace, warn};
+use tracing::field::{display, Empty};
+use tracing::{debug, instrument, trace, warn};
 use validator::Validate;
 
 use super::definitions::{ClientError, ClientResult};
-use crate::component_definitions::{ComponentClient, ServerError, APPLICATION_OCTET_STREAM};
+use crate::component_definitions::{
+    ComponentClient,
+    RequestId,
+    ServerError,
+    APPLICATION_OCTET_STREAM,
+    REQUEST_ID_HEADER,
+};
 use crate::metrics::RemoteClientMetrics;
 use crate::requests::LabeledRequest;
 use crate::serde_utils::SerdeWrapper;
@@ -163,10 +170,15 @@ where
         Self { uri, client, config, metrics, _req: PhantomData, _res: PhantomData }
     }
 
-    fn construct_http_request(&self, serialized_request: Bytes) -> HyperRequest<Body> {
+    fn construct_http_request(
+        &self,
+        serialized_request: Bytes,
+        request_id: &RequestId,
+    ) -> HyperRequest<Body> {
         trace!("Constructing remote request");
         HyperRequest::post(self.uri.clone())
             .header(CONTENT_TYPE, APPLICATION_OCTET_STREAM)
+            .header(REQUEST_ID_HEADER, request_id.to_string())
             .body(Body::from(serialized_request))
             .expect("Request building should succeed")
     }
@@ -217,7 +229,10 @@ where
     Request: Send + Serialize + DeserializeOwned + Debug + AsRef<str> + LabeledRequest,
     Response: Send + Serialize + DeserializeOwned + Debug,
 {
+    #[instrument(skip_all, fields(request_id = Empty))]
     async fn send(&self, component_request: Request) -> ClientResult<Response> {
+        let request_id = RequestId::generate();
+        tracing::Span::current().record("request_id", display(&request_id));
         let log_message = format!("{} to {}", component_request.as_ref(), self.uri);
         let request_label = component_request.request_label();
 
@@ -239,7 +254,8 @@ where
         let mut retry_interval_ms = self.config.initial_retry_delay_ms;
         for attempt in 1..max_attempts + 1 {
             trace!("Request {log_message} attempt {attempt} of {max_attempts}");
-            let http_request = self.construct_http_request(serialized_request_bytes.clone());
+            let http_request =
+                self.construct_http_request(serialized_request_bytes.clone(), &request_id);
             let start = Instant::now();
             let res = self.try_send(http_request).await;
             let elapsed = start.elapsed();
diff --git a/crates/apollo_infra/src/component_definitions.rs b/crates/apollo_infra/src/component_definitions.rs
@@ -2,6 +2,8 @@ use std::fmt::{Debug, Formatter, Result};
 
 use apollo_infra_utils::type_name::short_type_name;
 use async_trait::async_trait;
+use derive_more::{Display, FromStr};
+use rand::random;
 use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
@@ -88,15 +90,28 @@ where
     pub request: Request,
     pub tx: Sender<Response>,
     pub creation_time: Instant,
+    pub request_id: RequestId,
 }
 
+#[derive(Clone, Display, FromStr)]
+pub struct RequestId(pub u64);
+
+impl RequestId {
+    pub fn generate() -> Self {
+        Self(random::<u64>())
+    }
+}
+
+/// Header name for Request ID header.
+pub const REQUEST_ID_HEADER: &str = "request-id";
+
 impl<Request, Response> RequestWrapper<Request, Response>
 where
     Request: Send,
     Response: Send,
 {
-    pub fn new(request: Request, tx: Sender<Response>) -> Self {
-        Self { request, tx, creation_time: Instant::now() }
+    pub fn new(request: Request, tx: Sender<Response>, request_id: RequestId) -> Self {
+        Self { request, tx, creation_time: Instant::now(), request_id }
     }
 }
 
diff --git a/crates/apollo_infra/src/component_server/local_component_server.rs b/crates/apollo_infra/src/component_server/local_component_server.rs
@@ -10,13 +10,14 @@ use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc::{channel, Receiver, Sender};
 use tokio::sync::Semaphore;
 use tokio::time::Instant;
-use tracing::{error, info, trace, warn};
+use tracing::{error, info, instrument, trace, warn};
 use validator::Validate;
 
 use crate::component_definitions::{
     ComponentRequestHandler,
     ComponentStarter,
     PrioritizedRequest,
+    RequestId,
     RequestPriority,
     RequestWrapper,
 };
@@ -210,7 +211,7 @@ where
 
         tokio::spawn(async move {
             loop {
-                let (request, tx) = get_next_request_for_processing(
+                let (request, tx, request_id) = get_next_request_for_processing(
                     &mut high_rx,
                     &mut normal_rx,
                     &component_name,
@@ -221,6 +222,7 @@ where
                 process_request(
                     &mut component,
                     request,
+                    request_id,
                     tx,
                     metrics,
                     processing_time_warning_threshold_ms,
@@ -311,7 +313,7 @@ where
         tokio::spawn(async move {
             loop {
                 // TODO(Tsabary): add a test for the queueing time metric.
-                let (request, tx) = get_next_request_for_processing(
+                let (request, tx, request_id) = get_next_request_for_processing(
                     &mut high_rx,
                     &mut normal_rx,
                     &component_name,
@@ -328,6 +330,7 @@ where
                     process_request(
                         &mut cloned_component,
                         request,
+                        request_id,
                         tx,
                         metrics,
                         processing_time_warning_threshold_ms,
@@ -371,9 +374,11 @@ where
     }
 }
 
+#[instrument(skip_all,fields(request_id = %request_id))]
 async fn process_request<Request, Response, Component>(
     component: &mut Component,
     request: Request,
+    request_id: RequestId,
     tx: Sender<Response>,
     metrics: &'static LocalServerMetrics,
     processing_time_warning_threshold_ms: u128,
@@ -430,7 +435,7 @@ async fn get_next_request_for_processing<Request, Response>(
     normal_rx: &mut Receiver<RequestWrapper<Request, Response>>,
     component_name: &str,
     metrics: &'static LocalServerMetrics,
-) -> (Request, Sender<Response>)
+) -> (Request, Sender<Response>, RequestId)
 where
     Request: Send + Debug + LabeledRequest,
     Response: Send,
@@ -451,12 +456,13 @@ where
     let request = request_wrapper.request;
     let tx = request_wrapper.tx;
     let creation_time = request_wrapper.creation_time;
+    let request_id = request_wrapper.request_id;
 
     trace!(
         "Component {component_name} received request {request:?} that was created at \
          {creation_time:?}",
     );
     metrics.record_queueing_time(creation_time.elapsed().as_secs_f64(), request.request_label());
 
-    (request, tx)
+    (request, tx, request_id)
 }
diff --git a/crates/apollo_infra/src/component_server/remote_component_server.rs b/crates/apollo_infra/src/component_server/remote_component_server.rs
@@ -17,15 +17,17 @@ use serde::de::DeserializeOwned;
 use serde::{Deserialize, Serialize};
 use tokio::sync::{OwnedSemaphorePermit, Semaphore};
 use tower::{service_fn, Service, ServiceExt};
-use tracing::{debug, error, trace, warn};
+use tracing::{debug, error, instrument, trace, warn};
 use validator::Validate;
 
 use crate::component_client::{ClientError, LocalComponentClient};
 use crate::component_definitions::{
     ComponentClient,
+    RequestId,
     ServerError,
     APPLICATION_OCTET_STREAM,
     BUSY_PREVIOUS_REQUESTS_MSG,
+    REQUEST_ID_HEADER,
 };
 use crate::component_server::ComponentServerStarter;
 use crate::metrics::RemoteServerMetrics;
@@ -108,8 +110,10 @@ where
         Self { local_client, config: remote_server_config, port, max_concurrency, metrics }
     }
 
+    #[instrument(skip_all,fields(request_id = %request_id))]
     async fn remote_component_server_handler(
         http_request: HyperRequest<Body>,
+        request_id: RequestId,
         local_client: LocalComponentClient<Request, Response>,
         metrics: &'static RemoteServerMetrics,
     ) -> Result<HyperResponse<Body>, hyper::Error> {
@@ -128,6 +132,7 @@ where
 
                 // Wrap the send operation in a tokio::spawn as it is NOT a cancel-safe operation.
                 // Even if the current task is cancelled, the inner task will continue to run.
+                // Note: this creates a new request ID for the local client.
                 let response = tokio::spawn(async move { local_client.send(request).await })
                     .await
                     .expect("Should be able to extract value from the task");
@@ -194,10 +199,17 @@ where
                     Ok(permit) => {
                         metrics.increment_number_of_connections();
                         trace!("Acquired semaphore permit for connection");
-                        let handle_request_service = service_fn(move |req| {
+                        let handle_request_service = service_fn(move |req: HyperRequest<Body>| {
                             trace!("Received request: {:?}", req);
+                            let request_id = req
+                                .headers()
+                                .get(REQUEST_ID_HEADER)
+                                .and_then(|header| header.to_str().ok())
+                                .and_then(|s| s.parse::<RequestId>().ok())
+                                .expect("Request ID should be present in the request headers");
                             Self::remote_component_server_handler(
                                 req,
+                                request_id,
                                 local_client.clone(),
                                 metrics,
                             )
diff --git a/crates/apollo_infra/src/tests/remote_component_client_server_test.rs b/crates/apollo_infra/src/tests/remote_component_client_server_test.rs
@@ -28,10 +28,12 @@ use crate::component_client::{
 };
 use crate::component_definitions::{
     ComponentClient,
+    RequestId,
     RequestWrapper,
     ServerError,
     APPLICATION_OCTET_STREAM,
     BUSY_PREVIOUS_REQUESTS_MSG,
+    REQUEST_ID_HEADER,
 };
 use crate::component_server::{
     ComponentServerStarter,
@@ -447,6 +449,7 @@ async fn faulty_client_setup() {
                 format!("http://[{}]:{}/", self.socket.ip(), self.socket.port()).parse().unwrap();
             let http_request = Request::post(uri)
                 .header(CONTENT_TYPE, APPLICATION_OCTET_STREAM)
+                .header(REQUEST_ID_HEADER, RequestId::generate().to_string())
                 .body(Body::from(SerdeWrapper::new(component_request).wrapper_serialize().unwrap()))
                 .unwrap();
             let http_response = Client::new().request(http_request).await.unwrap();
