Skip to content

Commit 65a795f

Browse files
westarlewestarle
committed
feat: Implement trace context extraction and injection interceptors
1 parent 39133b5 commit 65a795f

File tree

13 files changed

+569
-11
lines changed

13 files changed

+569
-11
lines changed

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,14 @@
1515
</parent>
1616

1717
<dependencies>
18+
<dependency>
19+
<groupId>io.opentelemetry</groupId>
20+
<artifactId>opentelemetry-api</artifactId>
21+
</dependency>
22+
<dependency>
23+
<groupId>io.opentelemetry</groupId>
24+
<artifactId>opentelemetry-context</artifactId>
25+
</dependency>
1826
<dependency>
1927
<groupId>com.google.api</groupId>
2028
<artifactId>gax</artifactId>
@@ -100,6 +108,16 @@
100108
</dependency>
101109

102110
<!-- test dependencies -->
111+
<dependency>
112+
<groupId>io.opentelemetry</groupId>
113+
<artifactId>opentelemetry-sdk</artifactId>
114+
<scope>test</scope>
115+
</dependency>
116+
<dependency>
117+
<groupId>io.opentelemetry</groupId>
118+
<artifactId>opentelemetry-sdk-testing</artifactId>
119+
<scope>test</scope>
120+
</dependency>
103121
<dependency>
104122
<groupId>io.grpc</groupId>
105123
<artifactId>grpc-s2a</artifactId>

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import com.google.api.gax.rpc.ClientStream;
3535
import com.google.api.gax.rpc.ClientStreamReadyObserver;
3636
import com.google.api.gax.rpc.ResponseObserver;
37+
import com.google.api.gax.tracing.ApiTracer.Scope;
3738
import com.google.common.base.Preconditions;
3839
import io.grpc.ClientCall;
3940
import io.grpc.MethodDescriptor;
@@ -93,7 +94,9 @@ public void run() {
9394
onReady.onReady(clientStream);
9495
}
9596
});
96-
controller.startBidi();
97+
try (Scope ignored = context.getTracer().inScope()) {
98+
controller.startBidi();
99+
}
97100

98101
return clientStream;
99102
}

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

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import com.google.api.core.ListenableFutureToApiFuture;
3535
import com.google.api.gax.rpc.ApiCallContext;
3636
import com.google.api.gax.rpc.UnaryCallable;
37+
import com.google.api.gax.tracing.ApiTracer.Scope;
3738
import com.google.common.base.Preconditions;
3839
import io.grpc.ClientCall;
3940
import io.grpc.MethodDescriptor;
@@ -61,10 +62,12 @@ public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext inputCon
6162

6263
ClientCall<RequestT, ResponseT> clientCall = GrpcClientCalls.newCall(descriptor, inputContext);
6364

64-
if (awaitTrailers) {
65-
return new ListenableFutureToApiFuture<>(ClientCalls.futureUnaryCall(clientCall, request));
66-
} else {
67-
return GrpcClientCalls.eagerFutureUnaryCall(clientCall, request);
65+
try (Scope ignored = inputContext.getTracer().inScope()) {
66+
if (awaitTrailers) {
67+
return new ListenableFutureToApiFuture<>(ClientCalls.futureUnaryCall(clientCall, request));
68+
} else {
69+
return GrpcClientCalls.eagerFutureUnaryCall(clientCall, request);
70+
}
6871
}
6972
}
7073

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import com.google.api.gax.rpc.ApiCallContext;
3333
import com.google.api.gax.rpc.ApiStreamObserver;
3434
import com.google.api.gax.rpc.ClientStreamingCallable;
35+
import com.google.api.gax.tracing.ApiTracer.Scope;
3536
import com.google.common.base.Preconditions;
3637
import io.grpc.ClientCall;
3738
import io.grpc.MethodDescriptor;
@@ -57,8 +58,10 @@ public ApiStreamObserver<RequestT> clientStreamingCall(
5758
ApiStreamObserver<ResponseT> responseObserver, ApiCallContext context) {
5859
Preconditions.checkNotNull(responseObserver);
5960
ClientCall<RequestT, ResponseT> call = GrpcClientCalls.newCall(descriptor, context);
60-
return new StreamObserverDelegate<>(
61-
ClientCalls.asyncClientStreamingCall(
62-
call, new ApiStreamObserverDelegate<>(responseObserver)));
61+
try (Scope ignored = context.getTracer().inScope()) {
62+
return new StreamObserverDelegate<>(
63+
ClientCalls.asyncClientStreamingCall(
64+
call, new ApiStreamObserverDelegate<>(responseObserver)));
65+
}
6366
}
6467
}

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import com.google.api.gax.rpc.ResponseObserver;
3434
import com.google.api.gax.rpc.ServerStreamingCallable;
3535
import com.google.api.gax.rpc.StreamController;
36+
import com.google.api.gax.tracing.ApiTracer.Scope;
3637
import com.google.common.base.Preconditions;
3738
import io.grpc.ClientCall;
3839
import io.grpc.MethodDescriptor;
@@ -65,6 +66,8 @@ public void call(
6566
ClientCall<RequestT, ResponseT> call = GrpcClientCalls.newCall(descriptor, context);
6667
GrpcDirectStreamController<RequestT, ResponseT> controller =
6768
new GrpcDirectStreamController<>(call, responseObserver);
68-
controller.start(request);
69+
try (Scope ignored = context.getTracer().inScope()) {
70+
controller.start(request);
71+
}
6972
}
7073
}

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

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Redistribution and use in source and binary forms, with or without
5+
* modification, are permitted provided that the following conditions are
6+
* met:
7+
*
8+
* * Redistributions of source code must retain the above copyright
9+
* notice, this list of conditions and the following disclaimer.
10+
* * Redistributions in binary form must reproduce the above
11+
* copyright notice, this list of conditions and the following disclaimer
12+
* in the documentation and/or other materials provided with the
13+
* distribution.
14+
* * Neither the name of Google LLC nor the names of its
15+
* contributors may be used to endorse or promote products derived from
16+
* this software without specific prior written permission.
17+
*
18+
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
19+
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
20+
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
21+
* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
22+
* OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
23+
* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
24+
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
25+
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
26+
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
27+
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
28+
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29+
*/
30+
package com.google.api.gax.grpc;
31+
32+
import com.google.api.core.InternalApi;
33+
import io.grpc.CallOptions;
34+
import io.grpc.Channel;
35+
import io.grpc.ClientCall;
36+
import io.grpc.ClientInterceptor;
37+
import io.grpc.ForwardingClientCall.SimpleForwardingClientCall;
38+
import io.grpc.Metadata;
39+
import io.grpc.MethodDescriptor;
40+
import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
41+
import io.opentelemetry.context.Context;
42+
import io.opentelemetry.context.propagation.TextMapPropagator;
43+
import io.opentelemetry.context.propagation.TextMapSetter;
44+
45+
/**
46+
* An interceptor to handle dynamic trace context propagation.
47+
*
48+
* <p>Package-private for internal usage.
49+
*/
50+
@InternalApi
51+
public class GrpcTracePropagationInterceptor implements ClientInterceptor {
52+
53+
private static volatile Boolean isOpentelemetryAvailable;
54+
55+
private static boolean isOpenTelemetryAvailable() {
56+
if (isOpentelemetryAvailable == null) {
57+
synchronized (GrpcTracePropagationInterceptor.class) {
58+
if (isOpentelemetryAvailable == null) {
59+
try {
60+
Class.forName("io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator");
61+
isOpentelemetryAvailable = true;
62+
} catch (ClassNotFoundException e) {
63+
// OpenTelemetry API is not available
64+
isOpentelemetryAvailable = false;
65+
}
66+
}
67+
}
68+
}
69+
return isOpentelemetryAvailable;
70+
}
71+
72+
public GrpcTracePropagationInterceptor() {}
73+
74+
@Override
75+
public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
76+
MethodDescriptor<ReqT, RespT> method, final CallOptions callOptions, Channel next) {
77+
if (!isOpenTelemetryAvailable()) {
78+
return next.newCall(method, callOptions);
79+
}
80+
return OpenTelemetryContextInjector.interceptCall(method, callOptions, next);
81+
}
82+
83+
private static class OpenTelemetryContextInjector {
84+
private static final TextMapSetter<Metadata> setter =
85+
new TextMapSetter<Metadata>() {
86+
@Override
87+
public void set(Metadata carrier, String key, String value) {
88+
if (carrier != null) {
89+
carrier.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value);
90+
}
91+
}
92+
};
93+
94+
static <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(
95+
MethodDescriptor<ReqT, RespT> method, final CallOptions callOptions, Channel next) {
96+
ClientCall<ReqT, RespT> call = next.newCall(method, callOptions);
97+
return new SimpleForwardingClientCall<ReqT, RespT>(call) {
98+
@Override
99+
public void start(ClientCall.Listener<RespT> responseListener, Metadata headers) {
100+
try {
101+
TextMapPropagator propagator = W3CTraceContextPropagator.getInstance();
102+
propagator.inject(Context.current(), headers, setter);
103+
} catch (NoSuchMethodError e) {
104+
// Silently ignore if incompatible OpenTelemetry version
105+
}
106+
super.start(responseListener, headers);
107+
}
108+
};
109+
}
110+
}
111+
}

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

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Redistribution and use in source and binary forms, with or without
5+
* modification, are permitted provided that the following conditions are
6+
* met:
7+
*
8+
* * Redistributions of source code must retain the above copyright
9+
* notice, this list of conditions and the following disclaimer.
10+
* * Redistributions in binary form must reproduce the above
11+
* copyright notice, this list of conditions and the following disclaimer
12+
* in the documentation and/or other materials provided with the
13+
* distribution.
14+
* * Neither the name of Google LLC nor the names of its
15+
* contributors may be used to endorse or promote products derived from
16+
* this software without specific prior written permission.
17+
*
18+
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
19+
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
20+
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
21+
* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
22+
* OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
23+
* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
24+
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
25+
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
26+
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
27+
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
28+
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29+
*/
30+
package com.google.api.gax.grpc;
31+
32+
import static org.junit.jupiter.api.Assertions.assertEquals;
33+
import static org.mockito.Mockito.mock;
34+
import static org.mockito.Mockito.when;
35+
36+
import io.grpc.CallOptions;
37+
import io.grpc.Channel;
38+
import io.grpc.ClientCall;
39+
import io.grpc.Metadata;
40+
import io.grpc.MethodDescriptor;
41+
import io.opentelemetry.api.GlobalOpenTelemetry;
42+
import io.opentelemetry.api.trace.Span;
43+
import io.opentelemetry.api.trace.Tracer;
44+
import io.opentelemetry.context.Scope;
45+
import io.opentelemetry.context.propagation.ContextPropagators;
46+
import io.opentelemetry.sdk.OpenTelemetrySdk;
47+
import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter;
48+
import io.opentelemetry.sdk.trace.SdkTracerProvider;
49+
import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
50+
import org.junit.jupiter.api.AfterAll;
51+
import org.junit.jupiter.api.BeforeAll;
52+
import org.junit.jupiter.api.Test;
53+
54+
public class GrpcTracePropagationInterceptorTest {
55+
56+
private static OpenTelemetrySdk openTelemetry;
57+
private static Tracer tracer;
58+
59+
@BeforeAll
60+
public static void setUp() {
61+
SdkTracerProvider tracerProvider =
62+
SdkTracerProvider.builder()
63+
.addSpanProcessor(SimpleSpanProcessor.create(InMemorySpanExporter.create()))
64+
.build();
65+
openTelemetry =
66+
OpenTelemetrySdk.builder()
67+
.setTracerProvider(tracerProvider)
68+
.setPropagators(
69+
ContextPropagators.create(
70+
io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator.getInstance()))
71+
.buildAndRegisterGlobal();
72+
tracer = openTelemetry.getTracer("test");
73+
}
74+
75+
@AfterAll
76+
public static void tearDown() {
77+
GlobalOpenTelemetry.resetForTest();
78+
}
79+
80+
@Test
81+
public void testPropagation() {
82+
GrpcTracePropagationInterceptor interceptor = new GrpcTracePropagationInterceptor();
83+
Channel channel = mock(Channel.class);
84+
MethodDescriptor<String, String> methodDescriptor =
85+
MethodDescriptor.<String, String>newBuilder()
86+
.setType(MethodDescriptor.MethodType.UNARY)
87+
.setFullMethodName("test/test")
88+
.setRequestMarshaller(mock(MethodDescriptor.Marshaller.class))
89+
.setResponseMarshaller(mock(MethodDescriptor.Marshaller.class))
90+
.build();
91+
92+
@SuppressWarnings("unchecked")
93+
ClientCall<String, String> clientCall = mock(ClientCall.class);
94+
when(channel.newCall(methodDescriptor, CallOptions.DEFAULT)).thenReturn(clientCall);
95+
96+
ClientCall<String, String> interceptedCall =
97+
interceptor.interceptCall(methodDescriptor, CallOptions.DEFAULT, channel);
98+
99+
Metadata metadata = new Metadata();
100+
@SuppressWarnings("unchecked")
101+
ClientCall.Listener<String> listener = mock(ClientCall.Listener.class);
102+
103+
Span span = tracer.spanBuilder("test-span").startSpan();
104+
try (Scope ignored = span.makeCurrent()) {
105+
interceptedCall.start(listener, metadata);
106+
} finally {
107+
span.end();
108+
}
109+
110+
String traceparent =
111+
metadata.get(Metadata.Key.of("traceparent", Metadata.ASCII_STRING_MARSHALLER));
112+
113+
// Assert the traceparent header was added and matches the span
114+
String expectedTraceId = span.getSpanContext().getTraceId();
115+
String expectedSpanId = span.getSpanContext().getSpanId();
116+
assertEquals("00-" + expectedTraceId + "-" + expectedSpanId + "-01", traceparent);
117+
}
118+
}

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,14 @@
2020
</properties>
2121

2222
<dependencies>
23+
<dependency>
24+
<groupId>io.opentelemetry</groupId>
25+
<artifactId>opentelemetry-api</artifactId>
26+
</dependency>
27+
<dependency>
28+
<groupId>io.opentelemetry</groupId>
29+
<artifactId>opentelemetry-context</artifactId>
30+
</dependency>
2331
<dependency>
2432
<groupId>com.google.api</groupId>
2533
<artifactId>gax</artifactId>
@@ -86,6 +94,16 @@
8694
</dependency>
8795

8896
<!-- test dependencies -->
97+
<dependency>
98+
<groupId>io.opentelemetry</groupId>
99+
<artifactId>opentelemetry-sdk</artifactId>
100+
<scope>test</scope>
101+
</dependency>
102+
<dependency>
103+
<groupId>io.opentelemetry</groupId>
104+
<artifactId>opentelemetry-sdk-testing</artifactId>
105+
<scope>test</scope>
106+
</dependency>
89107
<dependency>
90108
<groupId>com.google.api</groupId>
91109
<artifactId>gax</artifactId>

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import com.google.api.core.ApiFuture;
3333
import com.google.api.gax.rpc.ApiCallContext;
3434
import com.google.api.gax.rpc.UnaryCallable;
35+
import com.google.api.gax.tracing.ApiTracer.Scope;
3536
import com.google.common.base.Preconditions;
3637
import com.google.protobuf.TypeRegistry;
3738

@@ -65,7 +66,9 @@ public ApiFuture<ResponseT> futureCall(RequestT request, ApiCallContext inputCon
6566

6667
HttpJsonClientCall<RequestT, ResponseT> clientCall =
6768
HttpJsonClientCalls.newCall(descriptor, context);
68-
return HttpJsonClientCalls.futureUnaryCall(clientCall, request, context);
69+
try (Scope ignored = context.getTracer().inScope()) {
70+
return HttpJsonClientCalls.futureUnaryCall(clientCall, request, context);
71+
}
6972
}
7073

7174
@Override

0 commit comments

Comments
 (0)