feat(o11y): implement trace context propagation for HTTP and gRPC (#5018)

Implement W3C Trace Context propagation (`traceparent` and `tracestate`)
across both `reqwest` and `tonic` transports.

Bump `google-cloud-gax-internal` version from 0.7.10 to 0.7.11.

Author: haphungw

Cargo.lock

@@ -458,7 +458,7 @@ tokio-test        = { default-features = false, version = "0.4" }
 # Local packages used as dependencies.
 google-cloud-auth        = { default-features = false, version = "1.7.0", path = "src/auth" }
 google-cloud-gax         = { default-features = false, version = "1.8.0", path = "src/gax" }
-gaxi                     = { default-features = false, version = "0.7.10", path = "src/gax-internal", package = "google-cloud-gax-internal" }
+gaxi                     = { default-features = false, version = "0.7.11", path = "src/gax-internal", package = "google-cloud-gax-internal" }
 wkt                      = { default-features = false, version = "1", path = "src/wkt", package = "google-cloud-wkt" }
 google-cloud-wkt         = { default-features = false, version = "1", path = "src/wkt", package = "google-cloud-wkt" }
 google-cloud-api         = { default-features = false, version = "1.3.0", path = "src/generated/api/types" }

Cargo.toml

@@ -14,7 +14,7 @@
 
 [package]
 name        = "google-cloud-gax-internal"
-version     = "0.7.10"
+version     = "0.7.11"
 description = "Google Cloud Client Libraries for Rust - Implementation Details"
 build       = "build.rs"
 # Inherit other attributes from the workspace.

src/gax-internal/Cargo.toml

@@ -275,7 +275,13 @@ impl ReqwestClient {
             .record_http(&span, attempt_info.attempt_count, method, url)
     }
 
-    async fn execute_http_inner(&self, request: reqwest::Request) -> Result<reqwest::Response> {
+    #[cfg_attr(not(google_cloud_unstable_tracing), allow(unused_mut))]
+    async fn execute_http_inner(&self, mut request: reqwest::Request) -> Result<reqwest::Response> {
+        #[cfg(google_cloud_unstable_tracing)]
+        crate::observability::propagation::inject_context(
+            &tracing::Span::current(),
+            request.headers_mut(),
+        );
         self.inner.execute(request).await.map_err(map_send_error)
     }
 
@@ -413,7 +419,16 @@ impl ReqwestClient {
             .record_http(&span, attempt_count, method, url)
     }
 
-    async fn request_attempt_inner(&self, request: reqwest::Request) -> Result<reqwest::Response> {
+    #[cfg_attr(not(google_cloud_unstable_tracing), allow(unused_mut))]
+    async fn request_attempt_inner(
+        &self,
+        mut request: reqwest::Request,
+    ) -> Result<reqwest::Response> {
+        #[cfg(google_cloud_unstable_tracing)]
+        crate::observability::propagation::inject_context(
+            &tracing::Span::current(),
+            request.headers_mut(),
+        );
         let response = self.inner.execute(request).await.map_err(map_send_error)?;
         if !response.status().is_success() {
             return self::to_http_error(response).await;

src/gax-internal/src/http.rs

@@ -183,10 +183,11 @@ where
         self.inner.poll_ready(cx)
     }
 
-    fn call(&mut self, req: http::Request<B>) -> Self::Future {
+    fn call(&mut self, mut req: http::Request<B>) -> Self::Future {
         let attempt_count = req.extensions().get::<AttemptCount>().map(|a| a.as_i64());
         let resource_name = req.extensions().get::<ResourceName>().map(|r| r.as_str());
         let span = create_grpc_span(req.uri(), &self.layer.inner, attempt_count, resource_name);
+        crate::observability::propagation::inject_context(&span, req.headers_mut());
         let future = self.inner.call(req);
         ResponseFuture {
             inner: future,
@@ -220,9 +221,9 @@ impl<S> NoTracingTowerService<S> {
     }
 }
 
-impl<S, Req, ResBody> Service<Req> for NoTracingTowerService<S>
+impl<S, B, ResBody> Service<http::Request<B>> for NoTracingTowerService<S>
 where
-    S: Service<Req, Response = http::Response<ResBody>>,
+    S: Service<http::Request<B>, Response = http::Response<ResBody>>,
     S::Future: Send + 'static,
     ResBody: http_body::Body<Data = bytes::Bytes, Error = tonic::Status> + Send + 'static,
 {
@@ -238,7 +239,11 @@ where
         self.inner.poll_ready(cx)
     }
 
-    fn call(&mut self, req: Req) -> Self::Future {
+    fn call(&mut self, mut req: http::Request<B>) -> Self::Future {
+        crate::observability::propagation::inject_context(
+            &tracing::Span::current(),
+            req.headers_mut(),
+        );
         NoTracingFuture {
             inner: self.inner.call(req),
             _phantom: std::marker::PhantomData,

src/gax-internal/src/observability/grpc_tracing.rs

@@ -696,4 +696,57 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+    async fn propagate_trace_context() -> anyhow::Result<()> {
+        let (endpoint, _server) = start_echo_server().await?;
+
+        let mut config = google_cloud_gax_internal::options::ClientConfig::default();
+        config.tracing = true;
+        config.cred = Some(test_credentials());
+        let client = grpc::Client::new(config, &endpoint).await?;
+
+        let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder().build();
+        let tracer = opentelemetry::trace::TracerProvider::tracer(&tracer_provider, "test");
+        let telemetry = tracing_opentelemetry::layer().with_tracer(tracer);
+        use tracing_subscriber::layer::SubscriberExt;
+        let subscriber = tracing_subscriber::registry().with(telemetry);
+        let _guard = tracing::subscriber::set_default(subscriber);
+
+        let extensions = {
+            let mut e = tonic::Extensions::new();
+            e.insert(tonic::GrpcMethod::new(
+                "google.test.v1.EchoServices",
+                "Echo",
+            ));
+            e
+        };
+        let request = google::test::v1::EchoRequest {
+            message: "test message".into(),
+            ..Default::default()
+        };
+
+        use tracing::Instrument;
+        let span = tracing::info_span!("parent_span");
+        let response = client
+            .execute::<_, google::test::v1::EchoResponse>(
+                extensions,
+                http::uri::PathAndQuery::from_static("/google.test.v1.EchoService/Echo"),
+                request,
+                RequestOptions::default(),
+                "test-client",
+                "",
+            )
+            .instrument(span)
+            .await?;
+
+        let inner = response.into_inner();
+        assert!(
+            inner.metadata.contains_key("traceparent"),
+            "Metadata should contain traceparent. Metadata: {:?}",
+            inner.metadata
+        );
+
+        Ok(())
+    }
 }

src/gax-internal/tests/grpc_observability.rs

@@ -26,7 +26,7 @@ mod tests {
     use google_cloud_gax_internal::options::{ClientConfig, InstrumentationClientInfo};
     use google_cloud_test_utils::test_layer::{AttributeValue, TestLayer};
     use http::{Method, StatusCode};
-    use httptest::matchers::request::{body, method, path};
+    use httptest::matchers::request::{body, headers, method, path};
     use httptest::{Expectation, Server, all_of, responders::*};
     use opentelemetry_semantic_conventions::{attribute as otel_attr, trace as otel_trace};
     use serde::Deserialize;
@@ -524,4 +524,42 @@ mod tests {
 
         assert_eq!(got, want);
     }
+
+    #[tokio::test]
+    async fn propagate_trace_context() {
+        let server = Server::run();
+        let server_addr = server.addr();
+        let server_url = format!("http://{}", server_addr);
+        server.expect(
+            Expectation::matching(all_of![
+                method("GET"),
+                path("/test"),
+                headers(httptest::matchers::contains((
+                    "traceparent",
+                    httptest::matchers::any()
+                ))),
+            ])
+            .respond_with(status_code(200).body("{\"hello\": \"world\"}")),
+        );
+
+        let client = create_client(true, server_url.clone()).await;
+
+        let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder().build();
+        let tracer = opentelemetry::trace::TracerProvider::tracer(&tracer_provider, "test");
+        let telemetry = tracing_opentelemetry::layer().with_tracer(tracer);
+        use tracing_subscriber::layer::SubscriberExt;
+        let subscriber = tracing_subscriber::registry().with(telemetry);
+        let _guard = tracing::subscriber::set_default(subscriber);
+
+        let options = RequestOptions::default().insert_extension(PathTemplate("/test"));
+        let request = client.builder(Method::GET, "/test".to_string());
+
+        let span = tracing::info_span!("parent_span");
+        let result: Result<Response<TestResponse>> = client
+            .execute(request, None::<NoBody>, options)
+            .instrument(span)
+            .await;
+
+        assert!(result.is_ok(), "{result:?}");
+    }
 }