Skip to content

Commit f829aaa

Browse files
jinseopkim0jinseopkim0
committed
feat(gax): link T3 attempt spans to parent context and clean up on operation completion
1 parent fe1f1b1 commit f829aaa

3 files changed

Lines changed: 161 additions & 45 deletions

File tree

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

Lines changed: 88 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,9 @@ class OpenTelemetryTracingTracer implements ApiTracer {
5353
private final Map<String, Object> attemptAttributes;
5454
private final String attemptSpanName;
5555
private final ApiTracerContext apiTracerContext;
56+
private final io.opentelemetry.context.Context parentContext;
57+
private final java.util.concurrent.locks.ReentrantLock lock =
58+
new java.util.concurrent.locks.ReentrantLock();
5659
private volatile @Nullable Span attemptSpan;
5760
private final AtomicReference<io.opentelemetry.context.@Nullable Scope> scope =
5861
new AtomicReference<>();
@@ -85,6 +88,7 @@ public void injectTraceContext(java.util.Map<String, String> carrier) {
8588
this.apiTracerContext = apiTracerContext;
8689
this.attemptSpanName = resolveAttemptSpanName(apiTracerContext);
8790
this.attemptAttributes = new HashMap<>();
91+
this.parentContext = io.opentelemetry.context.Context.current();
8892
buildAttributes();
8993
}
9094

@@ -103,6 +107,7 @@ public void injectTraceContext(java.util.Map<String, String> carrier) {
103107
this.attemptSpanName = attemptSpanName;
104108
this.apiTracerContext = apiTracerContext;
105109
this.attemptAttributes = new HashMap<>();
110+
this.parentContext = io.opentelemetry.context.Context.current();
106111
buildAttributes();
107112
}
108113

@@ -128,30 +133,63 @@ private void buildAttributes() {
128133
@Override
129134
@SuppressWarnings("MustBeClosedChecker")
130135
public void attemptStarted(Object request, int attemptNumber) {
131-
Map<String, Object> currentAttemptAttributes = new HashMap<>(this.attemptAttributes);
132-
133-
if (attemptNumber > 0) {
134-
ApiTracerContext.Transport transport = apiTracerContext.transport();
135-
if (transport == ApiTracerContext.Transport.GRPC) {
136-
currentAttemptAttributes.put(
137-
ObservabilityAttributes.GRPC_RESEND_COUNT_ATTRIBUTE, (long) attemptNumber);
138-
} else if (transport == ApiTracerContext.Transport.HTTP) {
139-
currentAttemptAttributes.put(
140-
ObservabilityAttributes.HTTP_RESEND_COUNT_ATTRIBUTE, (long) attemptNumber);
136+
Span oldSpan = null;
137+
io.opentelemetry.context.Scope oldScope = null;
138+
lock.lock();
139+
try {
140+
if (attemptSpan != null) {
141+
oldSpan = attemptSpan;
142+
attemptSpan = null;
143+
oldScope = this.scope.getAndSet(null);
144+
}
145+
Map<String, Object> currentAttemptAttributes = new HashMap<>(this.attemptAttributes);
146+
147+
if (attemptNumber > 0) {
148+
ApiTracerContext.Transport transport = apiTracerContext.transport();
149+
if (transport == ApiTracerContext.Transport.GRPC) {
150+
currentAttemptAttributes.put(
151+
ObservabilityAttributes.GRPC_RESEND_COUNT_ATTRIBUTE, (long) attemptNumber);
152+
} else if (transport == ApiTracerContext.Transport.HTTP) {
153+
currentAttemptAttributes.put(
154+
ObservabilityAttributes.HTTP_RESEND_COUNT_ATTRIBUTE, (long) attemptNumber);
155+
}
141156
}
142-
}
143157

144-
SpanBuilder spanBuilder = tracer.spanBuilder(attemptSpanName);
158+
SpanBuilder spanBuilder = tracer.spanBuilder(attemptSpanName);
145159

146-
// Attempt spans are of the CLIENT kind
147-
spanBuilder.setSpanKind(SpanKind.CLIENT);
160+
// Attempt spans are of the CLIENT kind
161+
spanBuilder.setSpanKind(SpanKind.CLIENT);
162+
163+
// Link attempt span to parent context
164+
spanBuilder.setParent(parentContext);
165+
166+
// Pass the combined attributes to the new SpanBuilder method
167+
spanBuilder.setAllAttributes(ObservabilityUtils.toOtelAttributes(currentAttemptAttributes));
168+
169+
this.attemptSpan = spanBuilder.startSpan();
170+
// Make the span active on the current thread so logs can capture the trace ID.
171+
this.scope.set(attemptSpan.makeCurrent());
172+
} finally {
173+
lock.unlock();
174+
}
175+
if (oldSpan != null) {
176+
endAttemptSpan(oldSpan, oldScope, null);
177+
}
178+
}
148179

149-
// Pass the combined attributes to the new SpanBuilder method
150-
spanBuilder.setAllAttributes(ObservabilityUtils.toOtelAttributes(currentAttemptAttributes));
180+
@Override
181+
public void operationSucceeded() {
182+
recordErrorAndEndAttempt(null);
183+
}
151184

152-
this.attemptSpan = spanBuilder.startSpan();
153-
// Make the span active on the current thread so logs can capture the trace ID.
154-
this.scope.set(attemptSpan.makeCurrent());
185+
@Override
186+
public void operationCancelled() {
187+
recordErrorAndEndAttempt(new CancellationException());
188+
}
189+
190+
@Override
191+
public void operationFailed(Throwable error) {
192+
recordErrorAndEndAttempt(error);
155193
}
156194

157195
@Override
@@ -161,12 +199,13 @@ public void attemptSucceeded() {
161199

162200
@Override
163201
public void responseHeadersReceived(java.util.Map<String, Object> headers) {
164-
if (attemptSpan == null) {
202+
Span currentSpan = attemptSpan;
203+
if (currentSpan == null) {
165204
return;
166205
}
167206
long contentLength = extractContentLength(headers);
168207
if (contentLength >= 0) {
169-
attemptSpan.setAttribute(ObservabilityAttributes.HTTP_RESPONSE_BODY_SIZE, contentLength);
208+
currentSpan.setAttribute(ObservabilityAttributes.HTTP_RESPONSE_BODY_SIZE, contentLength);
170209
}
171210
}
172211

@@ -222,49 +261,57 @@ public void attemptPermanentFailure(Throwable error) {
222261
}
223262

224263
private void recordErrorAndEndAttempt(@Nullable Throwable error) {
225-
if (attemptSpan == null) {
226-
return;
264+
Span localAttemptSpan;
265+
io.opentelemetry.context.Scope localScope;
266+
lock.lock();
267+
try {
268+
localAttemptSpan = attemptSpan;
269+
if (localAttemptSpan == null) {
270+
return;
271+
}
272+
attemptSpan = null;
273+
localScope = this.scope.getAndSet(null);
274+
} finally {
275+
lock.unlock();
227276
}
277+
278+
endAttemptSpan(localAttemptSpan, localScope, error);
279+
}
280+
281+
private void endAttemptSpan(
282+
Span localAttemptSpan,
283+
io.opentelemetry.context.@Nullable Scope localScope,
284+
@Nullable Throwable error) {
228285
Map<String, Object> responseAttributes =
229286
ObservabilityUtils.getResponseAttributes(error, this.apiTracerContext.transport());
230287
if (!responseAttributes.isEmpty()) {
231-
attemptSpan.setAllAttributes(ObservabilityUtils.toOtelAttributes(responseAttributes));
288+
localAttemptSpan.setAllAttributes(ObservabilityUtils.toOtelAttributes(responseAttributes));
232289
}
233290

234291
if (error != null && !Strings.isNullOrEmpty(error.getMessage())) {
235-
attemptSpan.setAttribute(
292+
localAttemptSpan.setAttribute(
236293
ObservabilityAttributes.STATUS_MESSAGE_ATTRIBUTE, error.getMessage());
237294
}
238295

239-
endAttempt();
240-
}
241-
242-
private void endAttempt() {
243-
Span currentSpan = this.attemptSpan;
244-
this.attemptSpan = null;
245-
246-
io.opentelemetry.context.Scope currentScope = this.scope.getAndSet(null);
247-
// Remove the span from the current thread before closing the span.
248296
try {
249-
if (currentScope != null) {
250-
currentScope.close();
297+
if (localScope != null) {
298+
localScope.close();
251299
}
252300
} finally {
253-
if (currentSpan != null) {
254-
currentSpan.end();
255-
}
301+
localAttemptSpan.end();
256302
}
257303
}
258304

259305
@Override
260306
public void requestUrlResolved(String url) {
261-
if (attemptSpan == null) {
307+
Span currentSpan = attemptSpan;
308+
if (currentSpan == null) {
262309
return;
263310
}
264311
String sanitizedUrlString = ObservabilityUtils.sanitizeUrlFull(url);
265312
if (sanitizedUrlString.isEmpty()) {
266313
return;
267314
}
268-
attemptSpan.setAttribute(ObservabilityAttributes.HTTP_URL_FULL_ATTRIBUTE, sanitizedUrlString);
315+
currentSpan.setAttribute(ObservabilityAttributes.HTTP_URL_FULL_ATTRIBUTE, sanitizedUrlString);
269316
}
270317
}

‎sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/tracing/OpenTelemetryTracingTracerFactoryTest.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ void setUp() {
7575
when(openTelemetry.getTracer(anyString())).thenReturn(tracer);
7676
when(tracer.spanBuilder(anyString())).thenReturn(spanBuilder);
7777
when(spanBuilder.setSpanKind(any())).thenReturn(spanBuilder);
78+
when(spanBuilder.setParent(any())).thenReturn(spanBuilder);
7879
when(spanBuilder.setAllAttributes(any(Attributes.class))).thenReturn(spanBuilder);
7980
when(spanBuilder.startSpan()).thenReturn(span);
8081

@@ -228,7 +229,7 @@ void testNewTracer_withContext_grpc_usesFullMethodName() {
228229

229230
tracerInstance.attemptStarted(null, 1);
230231

231-
verify(tracer).spanBuilder("google.cloud.v1.Service/Method");
232+
verify(tracer, Mockito.atLeastOnce()).spanBuilder("google.cloud.v1.Service/Method");
232233
}
233234

234235
@ParameterizedTest
@@ -255,7 +256,7 @@ void testNewTracer_withContext_http_usesHttpMethodAndPathTemplate(
255256

256257
tracerInstance.attemptStarted(null, 1);
257258

258-
verify(tracer).spanBuilder(expectedSpanName);
259+
verify(tracer, Mockito.atLeastOnce()).spanBuilder(expectedSpanName);
259260
}
260261

261262
@Test
@@ -273,7 +274,7 @@ void testNewTracer_withContext_http_noHttpMethodOrPathTemplate_usesFullMethodNam
273274

274275
tracerInstance.attemptStarted(null, 1);
275276

276-
verify(tracer).spanBuilder("google.cloud.v1.Service.Method");
277+
verify(tracer, Mockito.atLeastOnce()).spanBuilder("google.cloud.v1.Service.Method");
277278
}
278279

279280
@Test
@@ -309,7 +310,7 @@ void testNewTracer_mergesFactoryContext() {
309310
tracerInstance.attemptStarted(null, 1);
310311

311312
ArgumentCaptor<Attributes> attributesCaptor = ArgumentCaptor.forClass(Attributes.class);
312-
verify(spanBuilder).setAllAttributes(attributesCaptor.capture());
313+
verify(spanBuilder, Mockito.atLeastOnce()).setAllAttributes(attributesCaptor.capture());
313314

314315
Attributes attributes = attributesCaptor.getValue();
315316
assertThat(attributes.asMap())

‎sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/tracing/OpenTelemetryTracingTracerTest.java‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import static org.mockito.ArgumentMatchers.anyString;
3535
import static org.mockito.ArgumentMatchers.eq;
3636
import static org.mockito.Mockito.lenient;
37+
import static org.mockito.Mockito.mock;
3738
import static org.mockito.Mockito.never;
3839
import static org.mockito.Mockito.times;
3940
import static org.mockito.Mockito.verify;
@@ -76,6 +77,7 @@ class OpenTelemetryTracingTracerTest {
7677
void setUp() {
7778
lenient().when(tracer.spanBuilder(anyString())).thenReturn(spanBuilder);
7879
lenient().when(spanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(spanBuilder);
80+
lenient().when(spanBuilder.setParent(any())).thenReturn(spanBuilder);
7981
lenient().when(spanBuilder.setAllAttributes(any(Attributes.class))).thenReturn(spanBuilder);
8082
lenient().when(spanBuilder.startSpan()).thenReturn(span);
8183
lenient().when(span.makeCurrent()).thenReturn(scope);
@@ -754,4 +756,70 @@ void testAttemptEnded_closesScope_permanentFailure() {
754756
openTelemetryTracingTracer.attemptPermanentFailure(new RuntimeException());
755757
verify(scope).close();
756758
}
759+
760+
@Test
761+
void testAttemptStarted_setsParentToParentContext() {
762+
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
763+
verify(spanBuilder).setParent(any(io.opentelemetry.context.Context.class));
764+
}
765+
766+
@Test
767+
void testOperationSucceeded_endsActiveAttemptSpan() {
768+
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
769+
openTelemetryTracingTracer.operationSucceeded();
770+
771+
verify(span).end();
772+
}
773+
774+
@Test
775+
void testOperationFailed_endsActiveAttemptSpanWithErrorAttributes() {
776+
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
777+
openTelemetryTracingTracer.operationFailed(new RuntimeException("operation failed"));
778+
779+
verify(span).setAttribute(ObservabilityAttributes.STATUS_MESSAGE_ATTRIBUTE, "operation failed");
780+
verify(span).end();
781+
}
782+
783+
@Test
784+
void testOperationCancelled_endsActiveAttemptSpanWithCancellation() {
785+
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
786+
openTelemetryTracingTracer.operationCancelled();
787+
788+
ArgumentCaptor<Attributes> attrsCaptor = ArgumentCaptor.forClass(Attributes.class);
789+
verify(span).setAllAttributes(attrsCaptor.capture());
790+
verify(span).end();
791+
792+
assertThat(attrsCaptor.getValue().asMap())
793+
.containsEntry(
794+
AttributeKey.stringKey(ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE),
795+
"CANCELLED");
796+
}
797+
798+
@Test
799+
void testAttemptStarted_whenPreviousAttemptActive_closesOldScopeAndSpan() {
800+
Span span1 = mock(Span.class);
801+
Scope scope1 = mock(Scope.class);
802+
when(span1.makeCurrent()).thenReturn(scope1);
803+
804+
Span span2 = mock(Span.class);
805+
Scope scope2 = mock(Scope.class);
806+
when(span2.makeCurrent()).thenReturn(scope2);
807+
808+
when(spanBuilder.startSpan()).thenReturn(span1, span2);
809+
810+
openTelemetryTracingTracer.attemptStarted(new Object(), 0);
811+
verify(span1).makeCurrent();
812+
813+
// Start a second attempt before the first attempt was ended
814+
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
815+
verify(scope1).close();
816+
verify(span1).end();
817+
verify(scope2, never()).close();
818+
verify(span2, never()).end();
819+
820+
// Now complete the second attempt
821+
openTelemetryTracingTracer.attemptSucceeded();
822+
verify(scope2).close();
823+
verify(span2).end();
824+
}
757825
}

0 commit comments

Comments
 (0)