feat(gax-internal): add traces to raw HTTP requests

src/gax-internal/src/http.rs

@@ -27,9 +27,7 @@ pub mod reqwest;
 use crate::as_inner::as_inner;
 use crate::attempt_info::AttemptInfo;
 #[cfg(google_cloud_unstable_tracing)]
-use crate::observability::{
-    create_http_attempt_span, record_http_response_attributes, record_intermediate_client_request,
-};
+use crate::observability::create_http_attempt_span;
 use google_cloud_auth::credentials::{
     Builder as CredentialsBuilder, CacheableResource, Credentials,
 };
@@ -249,6 +247,37 @@ impl ReqwestClient {
         let request = self
             .request(builder, &options, attempt_info.remaining_time)
             .await?;
+        #[cfg(google_cloud_unstable_tracing)]
+        if self._tracing_enabled {
+            return self
+                .execute_http_traced(request, options, attempt_info)
+                .await;
+        }
+        self.execute_http_inner(request).await
+    }
+
+    #[cfg(google_cloud_unstable_tracing)]
+    async fn execute_http_traced(
+        &self,
+        request: reqwest::Request,
+        options: RequestOptions,
+        attempt_info: AttemptInfo,
+    ) -> Result<reqwest::Response> {
+        use crate::observability::HttpResultExt;
+        let span = create_http_attempt_span(
+            &request,
+            &options,
+            self.instrumentation,
+            attempt_info.attempt_count,
+        );
+        let (method, url) = (request.method().clone(), request.url().clone());
+        self.execute_http_inner(request)
+            .instrument(span.clone())
+            .await
+            .record_http(&span, attempt_info.attempt_count, method, url)
+    }
+
+    async fn execute_http_inner(&self, request: reqwest::Request) -> Result<reqwest::Response> {
         self.inner.execute(request).await.map_err(map_send_error)
     }
 
@@ -368,45 +397,36 @@ impl ReqwestClient {
     ) -> Result<reqwest::Response> {
         let request = self.request(builder, options, remaining_time).await?;
         #[cfg(google_cloud_unstable_tracing)]
-        let method = request.method().clone();
-        #[cfg(google_cloud_unstable_tracing)]
-        let url = request.url().clone();
-
-        #[cfg(google_cloud_unstable_tracing)]
-        let (reqwest_result, span) = if self._tracing_enabled {
-            let span =
-                create_http_attempt_span(&request, options, self.instrumentation, _attempt_count);
-            // The instrument call ensures the span is entered/exited as the execute future is polled.
-            let result = self.inner.execute(request).instrument(span.clone()).await;
-            (result, Some(span))
-        } else {
-            (self.inner.execute(request).await, None)
-        };
-        #[cfg(not(google_cloud_unstable_tracing))]
-        let reqwest_result = self.inner.execute(request).await;
-
-        let intermediate_result = reqwest_result.map_err(map_send_error);
-        let intermediate_result = match intermediate_result {
-            Ok(response) if !response.status().is_success() => self::to_http_error(response).await,
-            other => other,
-        };
-
-        // Record span before parsing result, after parsing error.
-        #[cfg(google_cloud_unstable_tracing)]
         if self._tracing_enabled {
-            if let Some(s) = span {
-                record_http_response_attributes(&s, intermediate_result.as_ref());
-            }
-            // Record to the client request span after s exits.
-            record_intermediate_client_request(
-                intermediate_result.as_ref(),
-                _attempt_count,
-                &method,
-                &url,
-            );
+            return self
+                .request_attempt_traced(request, options, _attempt_count)
+                .await;
         }
+        self.request_attempt_inner(request).await
+    }
+
+    #[cfg(google_cloud_unstable_tracing)]
+    async fn request_attempt_traced(
+        &self,
+        request: reqwest::Request,
+        options: &RequestOptions,
+        attempt_count: u32,
+    ) -> Result<reqwest::Response> {
+        use crate::observability::HttpResultExt;
+        let span = create_http_attempt_span(&request, options, self.instrumentation, attempt_count);
+        let (method, url) = (request.method().clone(), request.url().clone());
+        self.request_attempt_inner(request)
+            .instrument(span.clone())
+            .await
+            .record_http(&span, attempt_count, method, url)
+    }
 
-        intermediate_result
+    async fn request_attempt_inner(&self, request: reqwest::Request) -> Result<reqwest::Response> {
+        let response = self.inner.execute(request).await.map_err(map_send_error)?;
+        if !response.status().is_success() {
+            return self::to_http_error(response).await;
+        }
+        Ok(response)
     }
 
     fn get_retry_policy(&self, options: &RequestOptions) -> Arc<dyn RetryPolicy> {
@@ -758,8 +778,8 @@ mod tests {
         let t3_span = client_request_span!("Service", "test_method", &TEST_INSTRUMENTATION_INFO);
 
         // Simulate T4 span scope ending before calling to_http_response
+        let t4_span = tracing::info_span!("t4_span");
         {
-            let t4_span = tracing::info_span!("t4_span");
             let _t4_enter = t4_span.enter();
             // T4 work happens here
         } // T4 exit
@@ -768,10 +788,14 @@ mod tests {
         let url = "https://example.com".parse().unwrap();
 
         // Manually call the enrichment function, mimicking request_attempt
-        {
+        let response = {
+            use crate::observability::HttpResultExt;
+
             let _enter = t3_span.enter();
-            record_intermediate_client_request(Ok(&response), 1, &Method::GET, &url);
-        }
+            Ok(response)
+                .record_http(&t4_span, 1, Method::GET, url)
+                .unwrap()
+        };
 
         let _ = super::to_http_response::<wkt::Empty>(response)
             .instrument(t3_span.clone())

src/gax-internal/src/observability.rs

@@ -30,9 +30,7 @@ mod errors;
 mod http_tracing;
 
 #[cfg(all(google_cloud_unstable_tracing, feature = "_internal-http-client"))]
-pub(crate) use http_tracing::{
-    create_http_attempt_span, record_http_response_attributes, record_intermediate_client_request,
-};
+pub(crate) use http_tracing::{ResultExt as HttpResultExt, create_http_attempt_span};
 
 #[cfg(all(google_cloud_unstable_tracing, feature = "_internal-grpc-client"))]
 pub mod grpc_tracing;