feat: Implement trace context extraction and injection interceptors

westarle · westarle · commit 2f77e2d34b37 · 2026-03-31T12:43:43.000-04:00
diff --git a/sdk-platform-java/gax-java/gax-grpc/pom.xml b/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>
diff --git a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcClientCalls.java b/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);
     }
 
diff --git a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcDirectBidiStreamingCallable.java b/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcDirectBidiStreamingCallable.java
@@ -34,6 +34,7 @@
 import com.google.api.gax.rpc.ClientStream;
 import com.google.api.gax.rpc.ClientStreamReadyObserver;
 import com.google.api.gax.rpc.ResponseObserver;
+import com.google.api.gax.tracing.ApiTracer.Scope;
 import com.google.common.base.Preconditions;
 import io.grpc.ClientCall;
 import io.grpc.MethodDescriptor;
@@ -93,7 +94,9 @@ public void run() {
                 onReady.onReady(clientStream);
               }
             });
-    controller.startBidi();
+    try (Scope ignored = context.getTracer().inScope()) {
+      controller.startBidi();
+    }
 
     return clientStream;
   }
diff --git a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcDirectCallable.java b/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcDirectCallable.java
@@ -34,6 +34,7 @@
 import com.google.api.core.ListenableFutureToApiFuture;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.UnaryCallable;
+import com.google.api.gax.tracing.ApiTracer.Scope;
 import com.google.common.base.Preconditions;
 import io.grpc.ClientCall;
 import io.grpc.MethodDescriptor;
@@ -61,10 +62,12 @@ public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext inputCon
 
     ClientCall<RequestT, ResponseT> clientCall = GrpcClientCalls.newCall(descriptor, inputContext);
 
-    if (awaitTrailers) {
-      return new ListenableFutureToApiFuture<>(ClientCalls.futureUnaryCall(clientCall, request));
-    } else {
-      return GrpcClientCalls.eagerFutureUnaryCall(clientCall, request);
+    try (Scope ignored = inputContext.getTracer().inScope()) {
+      if (awaitTrailers) {
+        return new ListenableFutureToApiFuture<>(ClientCalls.futureUnaryCall(clientCall, request));
+      } else {
+        return GrpcClientCalls.eagerFutureUnaryCall(clientCall, request);
+      }
     }
   }
 
diff --git a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcDirectClientStreamingCallable.java b/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcDirectClientStreamingCallable.java
@@ -32,6 +32,7 @@
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.ApiStreamObserver;
 import com.google.api.gax.rpc.ClientStreamingCallable;
+import com.google.api.gax.tracing.ApiTracer.Scope;
 import com.google.common.base.Preconditions;
 import io.grpc.ClientCall;
 import io.grpc.MethodDescriptor;
@@ -57,8 +58,10 @@ public ApiStreamObserver<RequestT> clientStreamingCall(
       ApiStreamObserver<ResponseT> responseObserver, ApiCallContext context) {
     Preconditions.checkNotNull(responseObserver);
     ClientCall<RequestT, ResponseT> call = GrpcClientCalls.newCall(descriptor, context);
-    return new StreamObserverDelegate<>(
-        ClientCalls.asyncClientStreamingCall(
-            call, new ApiStreamObserverDelegate<>(responseObserver)));
+    try (Scope ignored = context.getTracer().inScope()) {
+      return new StreamObserverDelegate<>(
+          ClientCalls.asyncClientStreamingCall(
+              call, new ApiStreamObserverDelegate<>(responseObserver)));
+    }
   }
 }
diff --git a/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcDirectServerStreamingCallable.java b/sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcDirectServerStreamingCallable.java
@@ -33,6 +33,7 @@
 import com.google.api.gax.rpc.ResponseObserver;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.StreamController;
+import com.google.api.gax.tracing.ApiTracer.Scope;
 import com.google.common.base.Preconditions;
 import io.grpc.ClientCall;
 import io.grpc.MethodDescriptor;
@@ -65,6 +66,8 @@ public void call(
     ClientCall<RequestT, ResponseT> call = GrpcClientCalls.newCall(descriptor, context);
     GrpcDirectStreamController<RequestT, ResponseT> controller =
         new GrpcDirectStreamController<>(call, responseObserver);
-    controller.start(request);
+    try (Scope ignored = context.getTracer().inScope()) {
+      controller.start(request);
+    }
   }
 }
diff --git a/sdk-platform-java/gax-java/gax-httpjson/pom.xml b/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>
diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonClientCalls.java b/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 {
diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonDirectCallable.java b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonDirectCallable.java
@@ -32,6 +32,7 @@
 import com.google.api.core.ApiFuture;
 import com.google.api.gax.rpc.ApiCallContext;
 import com.google.api.gax.rpc.UnaryCallable;
+import com.google.api.gax.tracing.ApiTracer.Scope;
 import com.google.common.base.Preconditions;
 import com.google.protobuf.TypeRegistry;
 
@@ -65,7 +66,9 @@ public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext inputCon
 
     HttpJsonClientCall<RequestT, ResponseT> clientCall =
         HttpJsonClientCalls.newCall(descriptor, context);
-    return HttpJsonClientCalls.futureUnaryCall(clientCall, request, context);
+    try (Scope ignored = context.getTracer().inScope()) {
+      return HttpJsonClientCalls.futureUnaryCall(clientCall, request, context);
+    }
   }
 
   @Override
diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonDirectServerStreamingCallable.java b/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonDirectServerStreamingCallable.java
@@ -33,6 +33,7 @@
 import com.google.api.gax.rpc.ResponseObserver;
 import com.google.api.gax.rpc.ServerStreamingCallable;
 import com.google.api.gax.rpc.StreamController;
+import com.google.api.gax.tracing.ApiTracer.Scope;
 import com.google.common.base.Preconditions;
 
 /**
@@ -64,6 +65,8 @@ public void call(
     HttpJsonClientCall<RequestT, ResponseT> call = HttpJsonClientCalls.newCall(descriptor, context);
     HttpJsonDirectStreamController<RequestT, ResponseT> controller =
         new HttpJsonDirectStreamController<>(call, responseObserver);
-    controller.start(request, context);
+    try (Scope ignored = context.getTracer().inScope()) {
+      controller.start(request, context);
+    }
   }
 }
diff --git a/sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/HttpJsonDirectStreamController.java b/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);
diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/ApiTracer.java b/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
diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/tracing/SpanTracer.java b/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}.
@@ -103,6 +144,15 @@ private void buildAttributes() {
     this.attemptAttributes.putAll(this.apiTracerContext.getAttemptAttributes());
   }
 
+  @Override
+  public Scope inScope() {
+    if (attemptSpan != null) {
+      io.opentelemetry.context.Scope otelScope = attemptSpan.makeCurrent();
+      return otelScope::close;
+    }
+    return () -> {};
+  }
+
   @Override
   public void attemptStarted(Object request, int attemptNumber) {
     Map<String, Object> currentAttemptAttributes = new HashMap<>(this.attemptAttributes);
