Skip to content

Commit f85641e

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

3 files changed

Lines changed: 143 additions & 40 deletions

File tree

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

Lines changed: 78 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,10 @@ class OpenTelemetryTracingTracer implements ApiTracer {
5252
private final Map<String, Object> attemptAttributes;
5353
private final String attemptSpanName;
5454
private final ApiTracerContext apiTracerContext;
55-
private @Nullable Span attemptSpan;
55+
private final io.opentelemetry.context.Context parentContext;
56+
private final java.util.concurrent.locks.ReentrantLock lock =
57+
new java.util.concurrent.locks.ReentrantLock();
58+
private volatile @Nullable Span attemptSpan;
5659

5760
@Override
5861
public void injectTraceContext(java.util.Map<String, String> carrier) {
@@ -82,6 +85,7 @@ public void injectTraceContext(java.util.Map<String, String> carrier) {
8285
this.apiTracerContext = apiTracerContext;
8386
this.attemptSpanName = resolveAttemptSpanName(apiTracerContext);
8487
this.attemptAttributes = new HashMap<>();
88+
this.parentContext = io.opentelemetry.context.Context.current();
8589
buildAttributes();
8690
}
8791

@@ -100,6 +104,7 @@ public void injectTraceContext(java.util.Map<String, String> carrier) {
100104
this.attemptSpanName = attemptSpanName;
101105
this.apiTracerContext = apiTracerContext;
102106
this.attemptAttributes = new HashMap<>();
107+
this.parentContext = io.opentelemetry.context.Context.current();
103108
buildAttributes();
104109
}
105110

@@ -124,28 +129,59 @@ private void buildAttributes() {
124129

125130
@Override
126131
public void attemptStarted(Object request, int attemptNumber) {
127-
Map<String, Object> currentAttemptAttributes = new HashMap<>(this.attemptAttributes);
128-
129-
if (attemptNumber > 0) {
130-
ApiTracerContext.Transport transport = apiTracerContext.transport();
131-
if (transport == ApiTracerContext.Transport.GRPC) {
132-
currentAttemptAttributes.put(
133-
ObservabilityAttributes.GRPC_RESEND_COUNT_ATTRIBUTE, (long) attemptNumber);
134-
} else if (transport == ApiTracerContext.Transport.HTTP) {
135-
currentAttemptAttributes.put(
136-
ObservabilityAttributes.HTTP_RESEND_COUNT_ATTRIBUTE, (long) attemptNumber);
132+
Span oldSpan = null;
133+
lock.lock();
134+
try {
135+
if (attemptSpan != null) {
136+
oldSpan = attemptSpan;
137+
attemptSpan = null;
138+
}
139+
Map<String, Object> currentAttemptAttributes = new HashMap<>(this.attemptAttributes);
140+
141+
if (attemptNumber > 0) {
142+
ApiTracerContext.Transport transport = apiTracerContext.transport();
143+
if (transport == ApiTracerContext.Transport.GRPC) {
144+
currentAttemptAttributes.put(
145+
ObservabilityAttributes.GRPC_RESEND_COUNT_ATTRIBUTE, (long) attemptNumber);
146+
} else if (transport == ApiTracerContext.Transport.HTTP) {
147+
currentAttemptAttributes.put(
148+
ObservabilityAttributes.HTTP_RESEND_COUNT_ATTRIBUTE, (long) attemptNumber);
149+
}
137150
}
138-
}
139151

140-
SpanBuilder spanBuilder = tracer.spanBuilder(attemptSpanName);
152+
SpanBuilder spanBuilder = tracer.spanBuilder(attemptSpanName);
141153

142-
// Attempt spans are of the CLIENT kind
143-
spanBuilder.setSpanKind(SpanKind.CLIENT);
154+
// Attempt spans are of the CLIENT kind
155+
spanBuilder.setSpanKind(SpanKind.CLIENT);
144156

145-
// Pass the combined attributes to the new SpanBuilder method
146-
spanBuilder.setAllAttributes(ObservabilityUtils.toOtelAttributes(currentAttemptAttributes));
157+
// Link attempt span to parent context
158+
spanBuilder.setParent(parentContext);
159+
160+
// Pass the combined attributes to the new SpanBuilder method
161+
spanBuilder.setAllAttributes(ObservabilityUtils.toOtelAttributes(currentAttemptAttributes));
162+
163+
this.attemptSpan = spanBuilder.startSpan();
164+
} finally {
165+
lock.unlock();
166+
}
167+
if (oldSpan != null) {
168+
endAttemptSpan(oldSpan, null);
169+
}
170+
}
171+
172+
@Override
173+
public void operationSucceeded() {
174+
recordErrorAndEndAttempt(null);
175+
}
176+
177+
@Override
178+
public void operationCancelled() {
179+
recordErrorAndEndAttempt(new CancellationException());
180+
}
147181

148-
this.attemptSpan = spanBuilder.startSpan();
182+
@Override
183+
public void operationFailed(Throwable error) {
184+
recordErrorAndEndAttempt(error);
149185
}
150186

151187
@Override
@@ -155,12 +191,13 @@ public void attemptSucceeded() {
155191

156192
@Override
157193
public void responseHeadersReceived(java.util.Map<String, Object> headers) {
158-
if (attemptSpan == null) {
194+
Span currentSpan = attemptSpan;
195+
if (currentSpan == null) {
159196
return;
160197
}
161198
long contentLength = extractContentLength(headers);
162199
if (contentLength >= 0) {
163-
attemptSpan.setAttribute(ObservabilityAttributes.HTTP_RESPONSE_BODY_SIZE, contentLength);
200+
currentSpan.setAttribute(ObservabilityAttributes.HTTP_RESPONSE_BODY_SIZE, contentLength);
164201
}
165202
}
166203

@@ -216,41 +253,46 @@ public void attemptPermanentFailure(Throwable error) {
216253
}
217254

218255
private void recordErrorAndEndAttempt(@Nullable Throwable error) {
219-
if (attemptSpan == null) {
220-
return;
256+
Span localAttemptSpan;
257+
lock.lock();
258+
try {
259+
localAttemptSpan = attemptSpan;
260+
if (localAttemptSpan == null) {
261+
return;
262+
}
263+
attemptSpan = null;
264+
} finally {
265+
lock.unlock();
221266
}
267+
268+
endAttemptSpan(localAttemptSpan, error);
269+
}
270+
271+
private void endAttemptSpan(Span localAttemptSpan, @Nullable Throwable error) {
222272
Map<String, Object> responseAttributes =
223273
ObservabilityUtils.getResponseAttributes(error, this.apiTracerContext.transport());
224274
if (!responseAttributes.isEmpty()) {
225-
attemptSpan.setAllAttributes(ObservabilityUtils.toOtelAttributes(responseAttributes));
275+
localAttemptSpan.setAllAttributes(ObservabilityUtils.toOtelAttributes(responseAttributes));
226276
}
227277

228278
if (error != null && !Strings.isNullOrEmpty(error.getMessage())) {
229-
attemptSpan.setAttribute(
279+
localAttemptSpan.setAttribute(
230280
ObservabilityAttributes.STATUS_MESSAGE_ATTRIBUTE, error.getMessage());
231281
}
232282

233-
endAttempt();
234-
}
235-
236-
private void endAttempt() {
237-
if (attemptSpan == null) {
238-
return;
239-
}
240-
241-
attemptSpan.end();
242-
attemptSpan = null;
283+
localAttemptSpan.end();
243284
}
244285

245286
@Override
246287
public void requestUrlResolved(String url) {
247-
if (attemptSpan == null) {
288+
Span currentSpan = attemptSpan;
289+
if (currentSpan == null) {
248290
return;
249291
}
250292
String sanitizedUrlString = ObservabilityUtils.sanitizeUrlFull(url);
251293
if (sanitizedUrlString.isEmpty()) {
252294
return;
253295
}
254-
attemptSpan.setAttribute(ObservabilityAttributes.HTTP_URL_FULL_ATTRIBUTE, sanitizedUrlString);
296+
currentSpan.setAttribute(ObservabilityAttributes.HTTP_URL_FULL_ATTRIBUTE, sanitizedUrlString);
255297
}
256298
}

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

Lines changed: 1 addition & 0 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

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

Lines changed: 64 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@
3333
import static org.mockito.ArgumentMatchers.any;
3434
import static org.mockito.ArgumentMatchers.anyString;
3535
import static org.mockito.ArgumentMatchers.eq;
36+
import static org.mockito.Mockito.lenient;
37+
import static org.mockito.Mockito.mock;
3638
import static org.mockito.Mockito.never;
3739
import static org.mockito.Mockito.verify;
3840
import static org.mockito.Mockito.when;
@@ -69,10 +71,11 @@ class OpenTelemetryTracingTracerTest {
6971

7072
@BeforeEach
7173
void setUp() {
72-
when(tracer.spanBuilder(anyString())).thenReturn(spanBuilder);
73-
when(spanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(spanBuilder);
74-
when(spanBuilder.setAllAttributes(any(Attributes.class))).thenReturn(spanBuilder);
75-
when(spanBuilder.startSpan()).thenReturn(span);
74+
lenient().when(tracer.spanBuilder(anyString())).thenReturn(spanBuilder);
75+
lenient().when(spanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(spanBuilder);
76+
lenient().when(spanBuilder.setParent(any())).thenReturn(spanBuilder);
77+
lenient().when(spanBuilder.setAllAttributes(any(Attributes.class))).thenReturn(spanBuilder);
78+
lenient().when(spanBuilder.startSpan()).thenReturn(span);
7679
openTelemetryTracingTracer =
7780
new OpenTelemetryTracingTracer(tracer, ApiTracerContext.empty(), ATTEMPT_SPAN_NAME);
7881
}
@@ -680,4 +683,61 @@ void testInjectTraceContext_addsHeaders() {
680683
assertThat(carrier.get("traceparent")).contains("00000000000000000000000000000001");
681684
assertThat(carrier.get("traceparent")).contains("0000000000000002");
682685
}
686+
687+
@Test
688+
void testAttemptStarted_setsParentToParentContext() {
689+
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
690+
verify(spanBuilder).setParent(any(io.opentelemetry.context.Context.class));
691+
}
692+
693+
@Test
694+
void testOperationSucceeded_endsActiveAttemptSpan() {
695+
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
696+
openTelemetryTracingTracer.operationSucceeded();
697+
698+
verify(span).end();
699+
}
700+
701+
@Test
702+
void testOperationFailed_endsActiveAttemptSpanWithErrorAttributes() {
703+
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
704+
openTelemetryTracingTracer.operationFailed(new RuntimeException("operation failed"));
705+
706+
verify(span).setAttribute(ObservabilityAttributes.STATUS_MESSAGE_ATTRIBUTE, "operation failed");
707+
verify(span).end();
708+
}
709+
710+
@Test
711+
void testOperationCancelled_endsActiveAttemptSpanWithCancellation() {
712+
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
713+
openTelemetryTracingTracer.operationCancelled();
714+
715+
ArgumentCaptor<Attributes> attrsCaptor = ArgumentCaptor.forClass(Attributes.class);
716+
verify(span).setAllAttributes(attrsCaptor.capture());
717+
verify(span).end();
718+
719+
assertThat(attrsCaptor.getValue().asMap())
720+
.containsEntry(
721+
AttributeKey.stringKey(ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE),
722+
"CANCELLED");
723+
}
724+
725+
@Test
726+
void testAttemptStarted_whenPreviousAttemptActive_closesOldSpan() {
727+
Span span1 = mock(Span.class);
728+
Span span2 = mock(Span.class);
729+
730+
when(spanBuilder.startSpan()).thenReturn(span1, span2);
731+
732+
openTelemetryTracingTracer.attemptStarted(new Object(), 0);
733+
734+
// Start a second attempt before the first attempt was ended
735+
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
736+
verify(span1).end();
737+
verify(span2, never()).end();
738+
739+
// Now complete the second attempt
740+
openTelemetryTracingTracer.attemptSucceeded();
741+
verify(span2).end();
742+
}
683743
}

0 commit comments

Comments
 (0)