feat: Implement trace context extraction and injection with integration test

Author: westarle

java-storage-nio/google-cloud-nio-bom/pom.xml

@@ -53,7 +53,7 @@
         <dependency>
           <groupId>com.google.cloud</groupId>
           <artifactId>google-cloud-nio</artifactId>
-          <version>0.128.14</version><!-- {x-version-update:google-cloud-nio:current} -->
+          <version>0.128.15-SNAPSHOT</version><!-- {x-version-update:google-cloud-nio:current} -->
         </dependency>
       </dependencies>
     </dependencyManagement>

java-storage-nio/google-cloud-nio/pom.xml

@@ -2,7 +2,7 @@
 <project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
   <modelVersion>4.0.0</modelVersion>
   <artifactId>google-cloud-nio</artifactId>
-  <version>0.128.14</version><!-- {x-version-update:google-cloud-nio:current} -->
+  <version>0.128.15-SNAPSHOT</version><!-- {x-version-update:google-cloud-nio:current} -->
   <packaging>jar</packaging>
   <name>Google Cloud NIO</name>
   <url>https://github.com/googleapis/google-cloud-java</url>

sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcClientCalls.java

@@ -90,9 +90,28 @@ public static <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
       channel = ((ChannelPool) channel).getChannel(grpcContext.getChannelAffinity());
     }
 
-    if (!grpcContext.getExtraHeaders().isEmpty()) {
-      ClientInterceptor interceptor =
-          MetadataUtils.newAttachHeadersInterceptor(grpcContext.getMetadata());
+    java.util.Map<String, String> traceContext = new java.util.HashMap<>();
+    grpcContext.getTracer().injectTraceContext(traceContext);
+
+    if (!grpcContext.getExtraHeaders().isEmpty() || !traceContext.isEmpty()) {
+      Metadata metadata = new Metadata();
+      metadata.merge(grpcContext.getMetadata());
+      for (java.util.Map.Entry<String, java.util.List<String>> entry :
+          grpcContext.getExtraHeaders().entrySet()) {
+        Metadata.Key<String> key =
+            Metadata.Key.of(entry.getKey(), Metadata.ASCII_STRING_MARSHALLER);
+        metadata.removeAll(key);
+        for (String value : entry.getValue()) {
+          metadata.put(key, value);
+        }
+      }
+      for (java.util.Map.Entry<String, String> entry : traceContext.entrySet()) {
+        Metadata.Key<String> key =
+            Metadata.Key.of(entry.getKey(), Metadata.ASCII_STRING_MARSHALLER);
+        metadata.removeAll(key);
+        metadata.put(key, entry.getValue());
+      }
+      ClientInterceptor interceptor = MetadataUtils.newAttachHeadersInterceptor(metadata);
       channel = ClientInterceptors.intercept(channel, interceptor);
     }
 

sdk-platform-java/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcClientCallsTest.java

@@ -340,4 +340,48 @@ void testUniverseDomainNotReady_shouldRetry() throws IOException {
     Truth.assertThat(exception.isRetryable()).isTrue();
     Mockito.verify(mockChannel, Mockito.never()).newCall(descriptor, callOptions);
   }
+
+  @Test
+  void testTraceContextHeaders() throws IOException {
+    Metadata emptyHeaders = new Metadata();
+
+    MethodDescriptor<Color, Money> descriptor = FakeServiceGrpc.METHOD_RECOGNIZE;
+
+    @SuppressWarnings("unchecked")
+    ClientCall<Color, Money> mockClientCall = Mockito.mock(ClientCall.class);
+
+    @SuppressWarnings("unchecked")
+    ClientCall.Listener<Money> mockListener = Mockito.mock(ClientCall.Listener.class);
+
+    Channel mockChannel = Mockito.mock(ManagedChannel.class);
+    com.google.api.gax.tracing.ApiTracer mockTracer =
+        Mockito.mock(com.google.api.gax.tracing.ApiTracer.class);
+
+    Mockito.doAnswer(
+            invocation -> {
+              java.util.Map<String, String> carrier = invocation.getArgument(0);
+              carrier.put("traceparent", "00-00000000000000000000000000000001-0000000000000002-01");
+              return null;
+            })
+        .when(mockTracer)
+        .injectTraceContext(Mockito.anyMap());
+
+    Mockito.doAnswer(
+            invocation -> {
+              Metadata clientCallHeaders = (Metadata) invocation.getArguments()[1];
+              Metadata.Key<String> traceparentKey =
+                  Metadata.Key.of("traceparent", Metadata.ASCII_STRING_MARSHALLER);
+              assertThat(clientCallHeaders.getAll(traceparentKey))
+                  .containsExactly("00-00000000000000000000000000000001-0000000000000002-01");
+              return null;
+            })
+        .when(mockClientCall)
+        .start(Mockito.<ClientCall.Listener<Money>>any(), Mockito.<Metadata>any());
+
+    Mockito.when(mockChannel.newCall(Mockito.eq(descriptor), Mockito.<CallOptions>any()))
+        .thenReturn(mockClientCall);
+
+    GrpcCallContext context = defaultCallContext.withChannel(mockChannel).withTracer(mockTracer);
+    GrpcClientCalls.newCall(descriptor, context).start(mockListener, emptyHeaders);
+  }
 }

sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonClientCalls.java

@@ -80,15 +80,25 @@ public static <RequestT, ResponseT> HttpJsonClientCall<RequestT, ResponseT> newC
     return httpJsonContext.getChannel().newCall(methodDescriptor, httpJsonContext.getCallOptions());
   }
 
+  static HttpJsonMetadata getMetadataWithTraceContext(HttpJsonCallContext context) {
+    java.util.Map<String, String> traceHeaders = new java.util.HashMap<>();
+    context.getTracer().injectTraceContext(traceHeaders);
+
+    java.util.Map<String, java.util.List<String>> finalHeaders =
+        new java.util.HashMap<>(context.getExtraHeaders());
+    for (java.util.Map.Entry<String, String> entry : traceHeaders.entrySet()) {
+      finalHeaders.put(entry.getKey(), java.util.Collections.singletonList(entry.getValue()));
+    }
+    return HttpJsonMetadata.newBuilder().build().withHeaders(finalHeaders);
+  }
+
   static <RequestT, ResponseT> ApiFuture<ResponseT> futureUnaryCall(
       HttpJsonClientCall<RequestT, ResponseT> clientCall,
       RequestT request,
       HttpJsonCallContext context) {
     // Start the call
     HttpJsonFuture<ResponseT> future = new HttpJsonFuture<>(clientCall);
-    clientCall.start(
-        new FutureListener<>(future),
-        HttpJsonMetadata.newBuilder().build().withHeaders(context.getExtraHeaders()));
+    clientCall.start(new FutureListener<>(future), getMetadataWithTraceContext(context));
 
     // Send the request
     try {

sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/HttpJsonClientCallsTest.java

@@ -139,4 +139,35 @@ void testUniverseDomainNotReady_shouldRetry() throws IOException {
         .isEqualTo(HttpJsonStatusCode.Code.UNAUTHENTICATED);
     Mockito.verify(mockChannel, Mockito.never()).newCall(descriptor, callOptions);
   }
+
+  @Test
+  void testGetMetadataWithTraceContext() {
+    com.google.api.gax.tracing.ApiTracer mockTracer =
+        Mockito.mock(com.google.api.gax.tracing.ApiTracer.class);
+    Mockito.doAnswer(
+            invocation -> {
+              java.util.Map<String, String> carrier = invocation.getArgument(0);
+              carrier.put("traceparent", "00-00000000000000000000000000000001-0000000000000002-01");
+              return null;
+            })
+        .when(mockTracer)
+        .injectTraceContext(Mockito.anyMap());
+
+    java.util.Map<String, java.util.List<String>> extraHeaders = new java.util.HashMap<>();
+    extraHeaders.put("existing-header", java.util.Collections.singletonList("existing-value"));
+
+    HttpJsonCallContext context =
+        (HttpJsonCallContext)
+            HttpJsonCallContext.createDefault()
+                .withTracer(mockTracer)
+                .withExtraHeaders(extraHeaders);
+
+    HttpJsonMetadata metadata = HttpJsonClientCalls.getMetadataWithTraceContext(context);
+
+    assertThat(metadata.getHeaders()).containsKey("existing-header");
+    assertThat(metadata.getHeaders().get("existing-header").toString()).contains("existing-value");
+    assertThat(metadata.getHeaders()).containsKey("traceparent");
+    assertThat(metadata.getHeaders().get("traceparent").toString())
+        .contains("00-00000000000000000000000000000001-0000000000000002-01");
+  }
 }

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/ApiTracer.java

@@ -196,6 +196,9 @@ default void requestSent() {}
   default void batchRequestSent(long elementCount, long requestSize) {}
   ;
 
+  /** Extract the trace context from the tracer and add it to the given headers map. */
+  default void injectTraceContext(java.util.Map<String, String> carrier) {}
+
   /**
    * Annotates the attempt with the full resolved HTTP URL. Only relevant for HTTP transport.
    *

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/CompositeTracer.java

@@ -213,4 +213,11 @@ public void batchRequestSent(long elementCount, long requestSize) {
       child.batchRequestSent(elementCount, requestSize);
     }
   }
+
+  @Override
+  public void injectTraceContext(java.util.Map<String, String> carrier) {
+    for (ApiTracer child : children) {
+      child.injectTraceContext(carrier);
+    }
+  }
 }

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/SpanTracer.java

@@ -57,6 +57,23 @@ public class SpanTracer implements ApiTracer {
   private final ApiTracerContext apiTracerContext;
   private Span attemptSpan;
 
+  @Override
+  public void injectTraceContext(java.util.Map<String, String> carrier) {
+    if (attemptSpan != null) {
+      io.opentelemetry.context.Context context =
+          io.opentelemetry.context.Context.current().with(attemptSpan);
+      io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator.getInstance()
+          .inject(
+              context,
+              carrier,
+              (c, k, v) -> {
+                if (c != null) {
+                  c.put(k, v);
+                }
+              });
+    }
+  }
+
   /**
    * Creates a new instance of {@code SpanTracer}.
    *

sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/tracing/CompositeTracerTest.java

@@ -270,4 +270,13 @@ void testBatchRequestSent() {
     verify(child1).batchRequestSent(10L, 100L);
     verify(child2).batchRequestSent(10L, 100L);
   }
+
+  @Test
+  void testInjectTraceContext() {
+    java.util.Map<String, String> carrier = new java.util.HashMap<>();
+    compositeTracer.injectTraceContext(carrier);
+    InOrder inOrder = inOrder(child1, child2);
+    inOrder.verify(child1).injectTraceContext(carrier);
+    inOrder.verify(child2).injectTraceContext(carrier);
+  }
 }