fix: merge request_tracing & request_id middleware

src/main.rs

@@ -42,9 +42,7 @@ use client_query::context::Context;
 use config::{ApiKeys, BlocklistEntry, ExchangeRateProvider};
 use indexer_client::IndexerClient;
 use indexing_performance::IndexingPerformance;
-use middleware::{
-    legacy_auth_adapter, RequestTracingLayer, RequireAuthorizationLayer, SetRequestIdLayer,
-};
+use middleware::{legacy_auth_adapter, RequestTracingLayer, RequireAuthorizationLayer};
 use network::{indexer_blocklist, subgraph_client::Client as SubgraphClient};
 use prometheus::{self, Encoder as _};
 use receipts::ReceiptSigner;
@@ -209,13 +207,8 @@ async fn main() {
                         .allow_headers(cors::Any)
                         .allow_methods([http::Method::OPTIONS, http::Method::POST]),
                 )
-                // Set up the query tracing span
-                .layer(RequestTracingLayer)
-                // Set the query ID on the request
-                .layer(SetRequestIdLayer::new(format!("{:?}", signer_address)))
-                // Handle legacy in-path auth, and convert it into a header
+                .layer(RequestTracingLayer::new(format!("{:?}", signer_address)))
                 .layer(axum::middleware::from_fn(legacy_auth_adapter))
-                // Require the query to be authorized
                 .layer(RequireAuthorizationLayer::new(auth_service)),
         );
 

src/middleware.rs

@@ -1,9 +1,7 @@
 mod legacy_auth;
-mod request_id;
 mod request_tracing;
 mod require_auth;
 
 pub use legacy_auth::legacy_auth_adapter;
-pub use request_id::{RequestId, SetRequestIdLayer};
-pub use request_tracing::RequestTracingLayer;
+pub use request_tracing::{RequestId, RequestTracingLayer};
 pub use require_auth::RequireAuthorizationLayer;

src/middleware/request_tracing.rs

@@ -1,4 +1,10 @@
-use std::task::{Context, Poll};
+use std::{
+    sync::{
+        atomic::{self, AtomicU64},
+        Arc,
+    },
+    task::{Context, Poll},
+};
 
 use axum::http::{Request, Response};
 use tower::Service;
@@ -7,6 +13,9 @@ use tracing::{
     instrument::{Instrument, Instrumented},
 };
 
+#[derive(Clone)]
+pub struct RequestId(pub String);
+
 /// Middleware that instruments client query request with a tracing span.
 ///
 /// This middleware instruments the request future with a span:
@@ -20,6 +29,8 @@ use tracing::{
 #[derive(Debug, Clone)]
 pub struct RequestTracing<S> {
     inner: S,
+    gateway_id: String,
+    counter: Arc<AtomicU64>,
 }
 
 impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for RequestTracing<S>
@@ -34,25 +45,47 @@ where
         self.inner.poll_ready(cx)
     }
 
-    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
+    fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
+        const CLOUDFLARE_RAY_ID: http::HeaderName = http::HeaderName::from_static("cf-ray");
+        let request_id = if let Some(ray_id) = req.headers().get(&CLOUDFLARE_RAY_ID) {
+            ray_id.to_str().unwrap_or_default().to_string()
+        } else {
+            let request_count = self.counter.fetch_add(1, atomic::Ordering::Relaxed);
+            format!("{}-{:x}", self.gateway_id, request_count)
+        };
+
+        req.extensions_mut().insert(RequestId(request_id.clone()));
         self.inner.call(req).instrument(tracing::info_span!(
             "client_request",
-            request_id = field::Empty,
+            request_id,
             selector = field::Empty,
         ))
     }
 }
 
-/// A layer that applies the [`RequestTracing`] middleware.
-///
-/// See [`RequestTracing`] for more details.
 #[derive(Debug, Clone)]
-pub struct RequestTracingLayer;
+pub struct RequestTracingLayer {
+    gateway_id: String,
+    counter: Arc<AtomicU64>,
+}
+
+impl RequestTracingLayer {
+    pub fn new(gateway_id: String) -> Self {
+        Self {
+            gateway_id,
+            counter: Arc::new(AtomicU64::new(0)),
+        }
+    }
+}
 
 impl<S> tower::layer::Layer<S> for RequestTracingLayer {
     type Service = RequestTracing<S>;
 
     fn layer(&self, inner: S) -> Self::Service {
-        RequestTracing { inner }
+        RequestTracing {
+            inner,
+            gateway_id: self.gateway_id.clone(),
+            counter: self.counter.clone(),
+        }
     }
 }