Skip to content

Commit 59e81c3

Browse files
westarlewestarle
committed
feat: Implement trace context extraction and injection with integration test
1 parent 4b5e8af commit 59e81c3

File tree

8 files changed

+332
-6
lines changed

8 files changed

+332
-6
lines changed

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,16 @@
100100
</dependency>
101101

102102
<!-- test dependencies -->
103+
<dependency>
104+
<groupId>io.opentelemetry</groupId>
105+
<artifactId>opentelemetry-sdk</artifactId>
106+
<scope>test</scope>
107+
</dependency>
108+
<dependency>
109+
<groupId>io.opentelemetry</groupId>
110+
<artifactId>opentelemetry-sdk-testing</artifactId>
111+
<scope>test</scope>
112+
</dependency>
103113
<dependency>
104114
<groupId>io.grpc</groupId>
105115
<artifactId>grpc-s2a</artifactId>

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -90,9 +90,16 @@ public static <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
9090
channel = ((ChannelPool) channel).getChannel(grpcContext.getChannelAffinity());
9191
}
9292

93-
if (!grpcContext.getExtraHeaders().isEmpty()) {
94-
ClientInterceptor interceptor =
95-
MetadataUtils.newAttachHeadersInterceptor(grpcContext.getMetadata());
93+
java.util.Map<String, String> traceContext = new java.util.HashMap<>();
94+
grpcContext.getTracer().injectTraceContext(traceContext);
95+
96+
if (!grpcContext.getExtraHeaders().isEmpty() || !traceContext.isEmpty()) {
97+
Metadata metadata = grpcContext.getMetadata();
98+
for (java.util.Map.Entry<String, String> entry : traceContext.entrySet()) {
99+
metadata.put(
100+
Metadata.Key.of(entry.getKey(), Metadata.ASCII_STRING_MARSHALLER), entry.getValue());
101+
}
102+
ClientInterceptor interceptor = MetadataUtils.newAttachHeadersInterceptor(metadata);
96103
channel = ClientInterceptors.intercept(channel, interceptor);
97104
}
98105

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,16 @@
8686
</dependency>
8787

8888
<!-- test dependencies -->
89+
<dependency>
90+
<groupId>io.opentelemetry</groupId>
91+
<artifactId>opentelemetry-sdk</artifactId>
92+
<scope>test</scope>
93+
</dependency>
94+
<dependency>
95+
<groupId>io.opentelemetry</groupId>
96+
<artifactId>opentelemetry-sdk-testing</artifactId>
97+
<scope>test</scope>
98+
</dependency>
8999
<dependency>
90100
<groupId>com.google.api</groupId>
91101
<artifactId>gax</artifactId>

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

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -80,15 +80,25 @@ public static <RequestT, ResponseT> HttpJsonClientCall<RequestT, ResponseT> newC
8080
return httpJsonContext.getChannel().newCall(methodDescriptor, httpJsonContext.getCallOptions());
8181
}
8282

83+
static HttpJsonMetadata getMetadataWithTraceContext(HttpJsonCallContext context) {
84+
java.util.Map<String, String> traceHeaders = new java.util.HashMap<>();
85+
context.getTracer().injectTraceContext(traceHeaders);
86+
87+
java.util.Map<String, java.util.List<String>> finalHeaders =
88+
new java.util.HashMap<>(context.getExtraHeaders());
89+
for (java.util.Map.Entry<String, String> entry : traceHeaders.entrySet()) {
90+
finalHeaders.put(entry.getKey(), java.util.Collections.singletonList(entry.getValue()));
91+
}
92+
return HttpJsonMetadata.newBuilder().build().withHeaders(finalHeaders);
93+
}
94+
8395
static <RequestT, ResponseT> ApiFuture<ResponseT> futureUnaryCall(
8496
HttpJsonClientCall<RequestT, ResponseT> clientCall,
8597
RequestT request,
8698
HttpJsonCallContext context) {
8799
// Start the call
88100
HttpJsonFuture<ResponseT> future = new HttpJsonFuture<>(clientCall);
89-
clientCall.start(
90-
new FutureListener<>(future),
91-
HttpJsonMetadata.newBuilder().build().withHeaders(context.getExtraHeaders()));
101+
clientCall.start(new FutureListener<>(future), getMetadataWithTraceContext(context));
92102

93103
// Send the request
94104
try {

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,9 @@ default void requestSent() {}
196196
default void batchRequestSent(long elementCount, long requestSize) {}
197197
;
198198

199+
/** Extract the trace context from the tracer and add it to the given headers map. */
200+
default void injectTraceContext(java.util.Map<String, String> carrier) {}
201+
199202
/**
200203
* Annotates the attempt with the full resolved HTTP URL. Only relevant for HTTP transport.
201204
*

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -213,4 +213,11 @@ public void batchRequestSent(long elementCount, long requestSize) {
213213
child.batchRequestSent(elementCount, requestSize);
214214
}
215215
}
216+
217+
@Override
218+
public void injectTraceContext(java.util.Map<String, String> carrier) {
219+
for (ApiTracer child : children) {
220+
child.injectTraceContext(carrier);
221+
}
222+
}
216223
}

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,23 @@ public class SpanTracer implements ApiTracer {
5757
private final ApiTracerContext apiTracerContext;
5858
private Span attemptSpan;
5959

60+
@Override
61+
public void injectTraceContext(java.util.Map<String, String> carrier) {
62+
if (attemptSpan != null) {
63+
io.opentelemetry.context.Context context =
64+
io.opentelemetry.context.Context.current().with(attemptSpan);
65+
io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator.getInstance()
66+
.inject(
67+
context,
68+
carrier,
69+
(c, k, v) -> {
70+
if (c != null) {
71+
c.put(k, v);
72+
}
73+
});
74+
}
75+
}
76+
6077
/**
6178
* Creates a new instance of {@code SpanTracer}.
6279
*

‎sdk-platform-java/java-showcase/gapic-showcase/src/test/java/com/google/showcase/v1beta1/it/ITOtelTracePropagation.java‎

Lines changed: 262 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,262 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.showcase.v1beta1.it;
18+
19+
import static com.google.common.truth.Truth.assertThat;
20+
21+
import com.google.api.gax.httpjson.ApiMethodDescriptor;
22+
import com.google.api.gax.httpjson.ForwardingHttpJsonClientCall;
23+
import com.google.api.gax.httpjson.ForwardingHttpJsonClientCallListener;
24+
import com.google.api.gax.httpjson.HttpJsonCallOptions;
25+
import com.google.api.gax.httpjson.HttpJsonChannel;
26+
import com.google.api.gax.httpjson.HttpJsonClientCall;
27+
import com.google.api.gax.httpjson.HttpJsonClientInterceptor;
28+
import com.google.api.gax.httpjson.HttpJsonMetadata;
29+
import com.google.api.gax.tracing.SpanTracerFactory;
30+
import com.google.common.collect.ImmutableList;
31+
import com.google.showcase.v1beta1.EchoClient;
32+
import com.google.showcase.v1beta1.EchoRequest;
33+
import com.google.showcase.v1beta1.EchoSettings;
34+
import com.google.showcase.v1beta1.it.util.TestClientInitializer;
35+
import io.grpc.CallOptions;
36+
import io.grpc.Channel;
37+
import io.grpc.ClientCall;
38+
import io.grpc.ClientInterceptor;
39+
import io.grpc.ForwardingClientCall;
40+
import io.grpc.ForwardingClientCallListener;
41+
import io.grpc.Metadata;
42+
import io.grpc.MethodDescriptor;
43+
import io.opentelemetry.api.GlobalOpenTelemetry;
44+
import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
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.data.SpanData;
50+
import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor;
51+
import java.util.List;
52+
import java.util.Map;
53+
import java.util.concurrent.TimeUnit;
54+
import org.junit.jupiter.api.AfterAll;
55+
import org.junit.jupiter.api.BeforeAll;
56+
import org.junit.jupiter.api.BeforeEach;
57+
import org.junit.jupiter.api.Test;
58+
59+
class ITOtelTracePropagation {
60+
private static final Metadata.Key<String> TRACEPARENT_GRPC_HEADER_KEY =
61+
Metadata.Key.of("traceparent", Metadata.ASCII_STRING_MARSHALLER);
62+
private static final String TRACEPARENT_HTTP_HEADER_KEY = "traceparent";
63+
64+
private static InMemorySpanExporter spanExporter;
65+
private static OpenTelemetrySdk openTelemetrySdk;
66+
67+
private static class GrpcResponseCapturingClientInterceptor implements ClientInterceptor {
68+
private Metadata responseHeaders;
69+
70+
@Override
71+
public <RequestT, ResponseT> ClientCall<RequestT, ResponseT> interceptCall(
72+
MethodDescriptor<RequestT, ResponseT> method, final CallOptions callOptions, Channel next) {
73+
ClientCall<RequestT, ResponseT> call = next.newCall(method, callOptions);
74+
return new ForwardingClientCall.SimpleForwardingClientCall<RequestT, ResponseT>(call) {
75+
@Override
76+
public void start(ClientCall.Listener<ResponseT> responseListener, Metadata headers) {
77+
ClientCall.Listener<ResponseT> forwardingResponseListener =
78+
new ForwardingClientCallListener.SimpleForwardingClientCallListener<ResponseT>(responseListener) {
79+
@Override
80+
public void onHeaders(Metadata headers) {
81+
responseHeaders = headers;
82+
super.onHeaders(headers);
83+
}
84+
};
85+
super.start(forwardingResponseListener, headers);
86+
}
87+
};
88+
}
89+
}
90+
91+
private static class HttpJsonResponseCapturingClientInterceptor implements HttpJsonClientInterceptor {
92+
private HttpJsonMetadata responseHeaders;
93+
94+
@Override
95+
public <RequestT, ResponseT> HttpJsonClientCall<RequestT, ResponseT> interceptCall(
96+
ApiMethodDescriptor<RequestT, ResponseT> method,
97+
HttpJsonCallOptions callOptions,
98+
HttpJsonChannel next) {
99+
HttpJsonClientCall<RequestT, ResponseT> call = next.newCall(method, callOptions);
100+
return new ForwardingHttpJsonClientCall.SimpleForwardingHttpJsonClientCall<
101+
RequestT, ResponseT>(call) {
102+
@Override
103+
public void start(Listener<ResponseT> responseListener, HttpJsonMetadata requestHeaders) {
104+
Listener<ResponseT> forwardingResponseListener =
105+
new ForwardingHttpJsonClientCallListener.SimpleForwardingHttpJsonClientCallListener<
106+
ResponseT>(responseListener) {
107+
@Override
108+
public void onHeaders(HttpJsonMetadata headers) {
109+
responseHeaders = headers;
110+
super.onHeaders(headers);
111+
}
112+
113+
@Override
114+
public void onMessage(ResponseT message) {
115+
super.onMessage(message);
116+
}
117+
118+
@Override
119+
public void onClose(int statusCode, HttpJsonMetadata trailers) {
120+
super.onClose(statusCode, trailers);
121+
}
122+
};
123+
super.start(forwardingResponseListener, requestHeaders);
124+
}
125+
};
126+
}
127+
}
128+
129+
private static GrpcResponseCapturingClientInterceptor grpcInterceptor;
130+
private static HttpJsonResponseCapturingClientInterceptor httpJsonInterceptor;
131+
132+
private static EchoClient grpcClient;
133+
private static EchoClient httpJsonClient;
134+
135+
@BeforeAll
136+
static void setup() throws Exception {
137+
spanExporter = InMemorySpanExporter.create();
138+
139+
SdkTracerProvider tracerProvider =
140+
SdkTracerProvider.builder()
141+
.addSpanProcessor(SimpleSpanProcessor.create(spanExporter))
142+
.build();
143+
144+
openTelemetrySdk =
145+
OpenTelemetrySdk.builder()
146+
.setTracerProvider(tracerProvider)
147+
.setPropagators(ContextPropagators.create(W3CTraceContextPropagator.getInstance()))
148+
.buildAndRegisterGlobal();
149+
150+
SpanTracerFactory tracingFactory = new SpanTracerFactory(openTelemetrySdk);
151+
152+
// Create gRPC Interceptor and Client
153+
grpcInterceptor = new GrpcResponseCapturingClientInterceptor();
154+
grpcClient =
155+
TestClientInitializer.createGrpcEchoClientOpentelemetry(
156+
tracingFactory,
157+
EchoSettings.defaultGrpcTransportProviderBuilder()
158+
.setChannelConfigurator(io.grpc.ManagedChannelBuilder::usePlaintext)
159+
.setInterceptorProvider(() -> ImmutableList.of(grpcInterceptor))
160+
.build());
161+
162+
// Create HttpJson Interceptor and Client
163+
httpJsonInterceptor = new HttpJsonResponseCapturingClientInterceptor();
164+
EchoSettings httpJsonEchoSettings =
165+
EchoSettings.newHttpJsonBuilder()
166+
.setCredentialsProvider(com.google.api.gax.core.NoCredentialsProvider.create())
167+
.setTransportChannelProvider(
168+
EchoSettings.defaultHttpJsonTransportProviderBuilder()
169+
.setEndpoint(TestClientInitializer.DEFAULT_HTTPJSON_ENDPOINT)
170+
.setInterceptorProvider(() -> ImmutableList.of(httpJsonInterceptor))
171+
.build())
172+
.build();
173+
174+
com.google.showcase.v1beta1.stub.EchoStubSettings echoStubSettings =
175+
(com.google.showcase.v1beta1.stub.EchoStubSettings)
176+
httpJsonEchoSettings.getStubSettings().toBuilder()
177+
.setTracerFactory(tracingFactory)
178+
.build();
179+
com.google.showcase.v1beta1.stub.EchoStub stub = echoStubSettings.createStub();
180+
httpJsonClient = EchoClient.create(stub);
181+
}
182+
183+
@BeforeEach
184+
void cleanUp() {
185+
spanExporter.reset();
186+
grpcInterceptor.responseHeaders = null;
187+
httpJsonInterceptor.responseHeaders = null;
188+
}
189+
190+
@AfterAll
191+
static void tearDown() throws InterruptedException {
192+
grpcClient.close();
193+
httpJsonClient.close();
194+
195+
grpcClient.awaitTermination(TestClientInitializer.AWAIT_TERMINATION_SECONDS, TimeUnit.SECONDS);
196+
httpJsonClient.awaitTermination(
197+
TestClientInitializer.AWAIT_TERMINATION_SECONDS, TimeUnit.SECONDS);
198+
199+
if (openTelemetrySdk != null) {
200+
openTelemetrySdk.close();
201+
}
202+
GlobalOpenTelemetry.resetForTest();
203+
}
204+
205+
@Test
206+
void testTracePropagation_grpc() {
207+
EchoRequest request = EchoRequest.newBuilder().setContent("test-grpc-propagation").build();
208+
grpcClient.echo(request);
209+
210+
List<SpanData> spans = spanExporter.getFinishedSpanItems();
211+
assertThat(spans).isNotEmpty();
212+
213+
SpanData attemptSpan =
214+
spans.stream()
215+
.filter(span -> span.getName().equals("google.showcase.v1beta1.Echo/Echo"))
216+
.findFirst()
217+
.orElseThrow(() -> new AssertionError("Attempt span not found"));
218+
219+
String expectedTraceId = attemptSpan.getSpanContext().getTraceId();
220+
String expectedSpanId = attemptSpan.getSpanContext().getSpanId();
221+
String expectedTraceparent = "00-" + expectedTraceId + "-" + expectedSpanId + "-01";
222+
223+
String headerValue = grpcInterceptor.responseHeaders.get(TRACEPARENT_GRPC_HEADER_KEY);
224+
assertThat(headerValue).isNotNull();
225+
assertThat(headerValue).isEqualTo(expectedTraceparent);
226+
}
227+
228+
@Test
229+
void testTracePropagation_httpjson() {
230+
EchoRequest request = EchoRequest.newBuilder().setContent("test-http-propagation").build();
231+
httpJsonClient.echo(request);
232+
233+
List<SpanData> spans = spanExporter.getFinishedSpanItems();
234+
assertThat(spans).isNotEmpty();
235+
236+
// The T4 CLIENT span generated by ApiTracer (SpanTracer)
237+
SpanData attemptSpan =
238+
spans.stream()
239+
.filter(span -> span.getName().equals("POST v1beta1/echo:echo"))
240+
.findFirst()
241+
.orElseThrow(() -> new AssertionError("Attempt span not found"));
242+
243+
String expectedTraceId = attemptSpan.getSpanContext().getTraceId();
244+
String expectedSpanId = attemptSpan.getSpanContext().getSpanId();
245+
String expectedTraceparent = "00-" + expectedTraceId + "-" + expectedSpanId + "-01";
246+
247+
assertThat(httpJsonInterceptor.responseHeaders).isNotNull();
248+
Map<String, Object> headers = httpJsonInterceptor.responseHeaders.getHeaders();
249+
String expectedHttpHeaderKey = "x-showcase-request-" + TRACEPARENT_HTTP_HEADER_KEY;
250+
assertThat(headers).containsKey(expectedHttpHeaderKey);
251+
252+
Object headerVal = headers.get(expectedHttpHeaderKey);
253+
if (headerVal instanceof List) {
254+
@SuppressWarnings("unchecked")
255+
List<String> traceparentHeaders = (List<String>) headerVal;
256+
assertThat(traceparentHeaders).hasSize(1);
257+
assertThat(traceparentHeaders.get(0)).isEqualTo(expectedTraceparent);
258+
} else {
259+
assertThat(String.valueOf(headerVal)).isEqualTo(expectedTraceparent);
260+
}
261+
}
262+
}

0 commit comments

Comments
 (0)