Skip to content

Commit 7068ce7

Browse files
jinseopkim0jinseopkim0
committed
fix: capture operationContext at creation and use explicit lock for thread safety
1 parent 03dd4b3 commit 7068ce7

2 files changed

Lines changed: 80 additions & 75 deletions

File tree

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

Lines changed: 62 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,9 @@ class OpenTelemetryTracingTracer implements ApiTracer {
5353
private final String attemptSpanName;
5454
private final String operationSpanName;
5555
private final ApiTracerContext apiTracerContext;
56+
private final io.opentelemetry.context.Context operationContext;
57+
private final java.util.concurrent.locks.ReentrantLock lock =
58+
new java.util.concurrent.locks.ReentrantLock();
5659
private volatile @Nullable Span operationSpan;
5760
private volatile @Nullable Span attemptSpan;
5861

@@ -114,6 +117,7 @@ public Scope inScope() {
114117
this.attemptAttributes = new HashMap<>();
115118
buildAttributes();
116119
startOperationSpan();
120+
this.operationContext = io.opentelemetry.context.Context.current().with(this.operationSpan);
117121
}
118122

119123
private void startOperationSpan() {
@@ -163,33 +167,37 @@ private void buildAttributes() {
163167

164168
@Override
165169
public void attemptStarted(Object request, int attemptNumber) {
166-
Map<String, Object> currentAttemptAttributes = new HashMap<>(this.attemptAttributes);
167-
168-
if (attemptNumber > 0) {
169-
ApiTracerContext.Transport transport = apiTracerContext.transport();
170-
if (transport == ApiTracerContext.Transport.GRPC) {
171-
currentAttemptAttributes.put(
172-
ObservabilityAttributes.GRPC_RESEND_COUNT_ATTRIBUTE, (long) attemptNumber);
173-
} else if (transport == ApiTracerContext.Transport.HTTP) {
174-
currentAttemptAttributes.put(
175-
ObservabilityAttributes.HTTP_RESEND_COUNT_ATTRIBUTE, (long) attemptNumber);
170+
lock.lock();
171+
try {
172+
Map<String, Object> currentAttemptAttributes = new HashMap<>(this.attemptAttributes);
173+
174+
if (attemptNumber > 0) {
175+
ApiTracerContext.Transport transport = apiTracerContext.transport();
176+
if (transport == ApiTracerContext.Transport.GRPC) {
177+
currentAttemptAttributes.put(
178+
ObservabilityAttributes.GRPC_RESEND_COUNT_ATTRIBUTE, (long) attemptNumber);
179+
} else if (transport == ApiTracerContext.Transport.HTTP) {
180+
currentAttemptAttributes.put(
181+
ObservabilityAttributes.HTTP_RESEND_COUNT_ATTRIBUTE, (long) attemptNumber);
182+
}
176183
}
177-
}
178184

179-
SpanBuilder spanBuilder = tracer.spanBuilder(attemptSpanName);
185+
SpanBuilder spanBuilder = tracer.spanBuilder(attemptSpanName);
180186

181-
// Attempt spans are of the CLIENT kind
182-
spanBuilder.setSpanKind(SpanKind.CLIENT);
187+
// Attempt spans are of the CLIENT kind
188+
spanBuilder.setSpanKind(SpanKind.CLIENT);
183189

184-
Span localOperationSpan = operationSpan;
185-
if (localOperationSpan != null) {
186-
spanBuilder.setParent(io.opentelemetry.context.Context.current().with(localOperationSpan));
187-
}
190+
if (operationSpan != null) {
191+
spanBuilder.setParent(operationContext);
192+
}
188193

189-
// Pass the combined attributes to the new SpanBuilder method
190-
spanBuilder.setAllAttributes(ObservabilityUtils.toOtelAttributes(currentAttemptAttributes));
194+
// Pass the combined attributes to the new SpanBuilder method
195+
spanBuilder.setAllAttributes(ObservabilityUtils.toOtelAttributes(currentAttemptAttributes));
191196

192-
this.attemptSpan = spanBuilder.startSpan();
197+
this.attemptSpan = spanBuilder.startSpan();
198+
} finally {
199+
lock.unlock();
200+
}
193201
}
194202

195203
@Override
@@ -208,31 +216,34 @@ public void operationFailed(Throwable error) {
208216
}
209217

210218
private void recordErrorAndEndOperation(@Nullable Throwable error) {
211-
Span localOperationSpan = operationSpan;
212-
if (localOperationSpan == null) {
213-
return;
214-
}
215-
operationSpan = null;
219+
Span localOperationSpan;
220+
lock.lock();
216221
try {
222+
localOperationSpan = operationSpan;
223+
if (localOperationSpan == null) {
224+
return;
225+
}
226+
operationSpan = null;
217227
Span localAttemptSpan = attemptSpan;
218228
if (localAttemptSpan != null) {
219229
recordErrorAndEndAttempt(error);
220230
}
221231
} finally {
222-
Map<String, Object> responseAttributes =
223-
ObservabilityUtils.getResponseAttributes(error, this.apiTracerContext.transport());
224-
if (!responseAttributes.isEmpty()) {
225-
localOperationSpan.setAllAttributes(
226-
ObservabilityUtils.toOtelAttributes(responseAttributes));
227-
}
232+
lock.unlock();
233+
}
228234

229-
if (error != null && !Strings.isNullOrEmpty(error.getMessage())) {
230-
localOperationSpan.setAttribute(
231-
ObservabilityAttributes.STATUS_MESSAGE_ATTRIBUTE, error.getMessage());
232-
}
235+
Map<String, Object> responseAttributes =
236+
ObservabilityUtils.getResponseAttributes(error, this.apiTracerContext.transport());
237+
if (!responseAttributes.isEmpty()) {
238+
localOperationSpan.setAllAttributes(ObservabilityUtils.toOtelAttributes(responseAttributes));
239+
}
233240

234-
localOperationSpan.end();
241+
if (error != null && !Strings.isNullOrEmpty(error.getMessage())) {
242+
localOperationSpan.setAttribute(
243+
ObservabilityAttributes.STATUS_MESSAGE_ATTRIBUTE, error.getMessage());
235244
}
245+
246+
localOperationSpan.end();
236247
}
237248

238249
@Override
@@ -303,30 +314,30 @@ public void attemptPermanentFailure(Throwable error) {
303314
}
304315

305316
private void recordErrorAndEndAttempt(@Nullable Throwable error) {
306-
if (attemptSpan == null) {
307-
return;
317+
Span localAttemptSpan;
318+
lock.lock();
319+
try {
320+
localAttemptSpan = attemptSpan;
321+
if (localAttemptSpan == null) {
322+
return;
323+
}
324+
attemptSpan = null;
325+
} finally {
326+
lock.unlock();
308327
}
328+
309329
Map<String, Object> responseAttributes =
310330
ObservabilityUtils.getResponseAttributes(error, this.apiTracerContext.transport());
311331
if (!responseAttributes.isEmpty()) {
312-
attemptSpan.setAllAttributes(ObservabilityUtils.toOtelAttributes(responseAttributes));
332+
localAttemptSpan.setAllAttributes(ObservabilityUtils.toOtelAttributes(responseAttributes));
313333
}
314334

315335
if (error != null && !Strings.isNullOrEmpty(error.getMessage())) {
316-
attemptSpan.setAttribute(
336+
localAttemptSpan.setAttribute(
317337
ObservabilityAttributes.STATUS_MESSAGE_ATTRIBUTE, error.getMessage());
318338
}
319339

320-
endAttempt();
321-
}
322-
323-
private void endAttempt() {
324-
if (attemptSpan == null) {
325-
return;
326-
}
327-
328-
attemptSpan.end();
329-
attemptSpan = null;
340+
localAttemptSpan.end();
330341
}
331342

332343
@Override

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

Lines changed: 18 additions & 24 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;
@@ -74,29 +76,23 @@ class OpenTelemetryTracingTracerTest {
7476

7577
@BeforeEach
7678
void setUp() {
77-
org.mockito.Mockito.lenient().when(tracer.spanBuilder(anyString())).thenReturn(spanBuilder);
78-
org.mockito.Mockito.lenient()
79-
.when(spanBuilder.setSpanKind(any(SpanKind.class)))
80-
.thenReturn(spanBuilder);
81-
org.mockito.Mockito.lenient().when(spanBuilder.setParent(any())).thenReturn(spanBuilder);
82-
org.mockito.Mockito.lenient()
83-
.when(spanBuilder.setAllAttributes(any(Attributes.class)))
84-
.thenReturn(spanBuilder);
85-
org.mockito.Mockito.lenient().when(spanBuilder.startSpan()).thenReturn(span);
86-
87-
org.mockito.Mockito.lenient()
79+
lenient().when(tracer.spanBuilder(anyString())).thenReturn(spanBuilder);
80+
lenient().when(spanBuilder.setSpanKind(any(SpanKind.class))).thenReturn(spanBuilder);
81+
lenient().when(spanBuilder.setParent(any())).thenReturn(spanBuilder);
82+
lenient().when(spanBuilder.setAllAttributes(any(Attributes.class))).thenReturn(spanBuilder);
83+
lenient().when(spanBuilder.startSpan()).thenReturn(span);
84+
85+
lenient()
8886
.when(operationSpanBuilder.setSpanKind(any(SpanKind.class)))
8987
.thenReturn(operationSpanBuilder);
90-
org.mockito.Mockito.lenient()
88+
lenient()
9189
.when(operationSpanBuilder.setAllAttributes(any(Attributes.class)))
9290
.thenReturn(operationSpanBuilder);
93-
org.mockito.Mockito.lenient().when(operationSpanBuilder.startSpan()).thenReturn(operationSpan);
94-
org.mockito.Mockito.lenient()
91+
lenient().when(operationSpanBuilder.startSpan()).thenReturn(operationSpan);
92+
lenient()
9593
.when(operationSpan.storeInContext(any(io.opentelemetry.context.Context.class)))
9694
.thenAnswer(invocation -> invocation.getArgument(0));
97-
org.mockito.Mockito.lenient()
98-
.when(tracer.spanBuilder("Service/Method"))
99-
.thenReturn(operationSpanBuilder);
95+
lenient().when(tracer.spanBuilder("Service/Method")).thenReturn(operationSpanBuilder);
10096

10197
openTelemetryTracingTracer =
10298
new OpenTelemetryTracingTracer(tracer, ApiTracerContext.empty(), ATTEMPT_SPAN_NAME);
@@ -783,9 +779,8 @@ void testOperationCancelled_endsOperationSpanWithCancelledStatus() {
783779

784780
@Test
785781
void testInScope_withAttemptSpan() {
786-
io.opentelemetry.context.Scope mockScope =
787-
org.mockito.Mockito.mock(io.opentelemetry.context.Scope.class);
788-
org.mockito.Mockito.when(span.makeCurrent()).thenReturn(mockScope);
782+
io.opentelemetry.context.Scope mockScope = mock(io.opentelemetry.context.Scope.class);
783+
when(span.makeCurrent()).thenReturn(mockScope);
789784

790785
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
791786
try (ApiTracer.Scope scope = openTelemetryTracingTracer.inScope()) {
@@ -796,9 +791,8 @@ void testInScope_withAttemptSpan() {
796791

797792
@Test
798793
void testInScope_withOperationSpanFallback() {
799-
io.opentelemetry.context.Scope mockScope =
800-
org.mockito.Mockito.mock(io.opentelemetry.context.Scope.class);
801-
org.mockito.Mockito.when(operationSpan.makeCurrent()).thenReturn(mockScope);
794+
io.opentelemetry.context.Scope mockScope = mock(io.opentelemetry.context.Scope.class);
795+
when(operationSpan.makeCurrent()).thenReturn(mockScope);
802796

803797
try (ApiTracer.Scope scope = openTelemetryTracingTracer.inScope()) {
804798
verify(operationSpan).makeCurrent();
@@ -815,7 +809,7 @@ void testInjectTraceContext_withOperationSpanFallback() {
815809
TraceFlags.getSampled(),
816810
TraceState.getDefault());
817811
Span realSpan = Span.wrap(mockSpanContext);
818-
org.mockito.Mockito.when(operationSpanBuilder.startSpan()).thenReturn(realSpan);
812+
when(operationSpanBuilder.startSpan()).thenReturn(realSpan);
819813

820814
openTelemetryTracingTracer =
821815
new OpenTelemetryTracingTracer(tracer, ApiTracerContext.empty(), ATTEMPT_SPAN_NAME);

0 commit comments

Comments
 (0)