feat: Implement trace context extraction and injection interceptors

Author: westarle

sdk-platform-java/gax-java/gax-grpc/pom.xml

@@ -100,6 +100,16 @@
     </dependency>
 
     <!-- test dependencies -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-testing</artifactId>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>io.grpc</groupId>
       <artifactId>grpc-s2a</artifactId>

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

@@ -90,9 +90,16 @@ 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 = grpcContext.getMetadata();
+      for (java.util.Map.Entry<String, String> entry : traceContext.entrySet()) {
+        metadata.put(
+            Metadata.Key.of(entry.getKey(), Metadata.ASCII_STRING_MARSHALLER), entry.getValue());
+      }
+      ClientInterceptor interceptor = MetadataUtils.newAttachHeadersInterceptor(metadata);
       channel = ClientInterceptors.intercept(channel, interceptor);
     }
 

sdk-platform-java/gax-java/gax-httpjson/pom.xml

@@ -86,6 +86,16 @@
     </dependency>
 
     <!-- test dependencies -->
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-sdk-testing</artifactId>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>com.google.api</groupId>
       <artifactId>gax</artifactId>

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/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/SpanTracer.java

@@ -55,6 +55,27 @@ 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) {
+      try {
+        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);
+                  }
+                });
+      } catch (NoSuchMethodError e) {
+        // Silently ignore if incompatible OpenTelemetry version
+      }
+    }
+  }
+
   /**
    * Creates a new instance of {@code SpanTracer}.
    *