feat: Implement trace context extraction and injection interceptors

Author: westarle

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

@@ -15,6 +15,14 @@
   </parent>
 
   <dependencies>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-api</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-context</artifactId>
+    </dependency>
     <dependency>
       <groupId>com.google.api</groupId>
       <artifactId>gax</artifactId>
@@ -100,6 +108,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

@@ -20,6 +20,14 @@
   </properties>
 
   <dependencies>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-api</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-context</artifactId>
+    </dependency>
     <dependency>
       <groupId>com.google.api</groupId>
       <artifactId>gax</artifactId>
@@ -86,6 +94,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-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonDirectStreamController.java

@@ -82,11 +82,12 @@ public void request(int count) {
   }
 
   void start(RequestT request, ApiCallContext context) {
+    HttpJsonCallContext httpJsonContext = HttpJsonCallContext.createDefault().nullToSelf(context);
     responseObserver.onStart(this);
     this.hasStarted = true;
     clientCall.start(
         new ResponseObserverAdapter(),
-        HttpJsonMetadata.newBuilder().build().withHeaders(context.getExtraHeaders()));
+        HttpJsonClientCalls.getMetadataWithTraceContext(httpJsonContext));
 
     if (autoflowControl) {
       clientCall.request(1);

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

@@ -192,6 +192,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) {}
+
   /**
    * A context class to be used with {@link #inScope()} and a try-with-resources block. Closing a
    * {@link Scope} removes any context that the underlying implementation might've set in {@link

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

@@ -52,6 +52,47 @@ public class SpanTracer implements ApiTracer {
   private final String attemptSpanName;
   private final ApiTracerContext apiTracerContext;
   private Span attemptSpan;
+  private static volatile Boolean isOpentelemetryAvailable;
+
+  private static boolean isOpenTelemetryAvailable() {
+    if (isOpentelemetryAvailable == null) {
+      synchronized (SpanTracer.class) {
+        if (isOpentelemetryAvailable == null) {
+          try {
+            Class.forName("io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator");
+            isOpentelemetryAvailable = true;
+          } catch (ClassNotFoundException e) {
+            isOpentelemetryAvailable = false;
+          }
+        }
+      }
+    }
+    return isOpentelemetryAvailable;
+  }
+
+  @Override
+  public void injectTraceContext(java.util.Map<String, String> carrier) {
+    if (!isOpenTelemetryAvailable()) {
+      return;
+    }
+    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}.