Skip to content

Commit f5ae43d

Browse files
jinseopkim0jinseopkim0
committed
feat: add T3 client request operation spans to OpenTelemetryTracingTracer
1 parent cc4b980 commit f5ae43d

3 files changed

Lines changed: 250 additions & 18 deletions

File tree

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

Lines changed: 95 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -51,14 +51,17 @@ class OpenTelemetryTracingTracer implements ApiTracer {
5151
private final Tracer tracer;
5252
private final Map<String, Object> attemptAttributes;
5353
private final String attemptSpanName;
54+
private final String operationSpanName;
5455
private final ApiTracerContext apiTracerContext;
56+
private @Nullable Span operationSpan;
5557
private @Nullable Span attemptSpan;
5658

5759
@Override
5860
public void injectTraceContext(java.util.Map<String, String> carrier) {
59-
if (attemptSpan != null) {
61+
Span spanToInject = attemptSpan != null ? attemptSpan : operationSpan;
62+
if (spanToInject != null) {
6063
io.opentelemetry.context.Context context =
61-
io.opentelemetry.context.Context.current().with(attemptSpan);
64+
io.opentelemetry.context.Context.current().with(spanToInject);
6265
io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator.getInstance()
6366
.inject(
6467
context,
@@ -71,18 +74,24 @@ public void injectTraceContext(java.util.Map<String, String> carrier) {
7174
}
7275
}
7376

77+
@Override
78+
public Scope inScope() {
79+
Span currentSpan = attemptSpan != null ? attemptSpan : operationSpan;
80+
if (currentSpan == null) {
81+
return () -> {};
82+
}
83+
io.opentelemetry.context.Scope otelScope = currentSpan.makeCurrent();
84+
return otelScope::close;
85+
}
86+
7487
/**
7588
* Creates a new instance of {@code OpenTelemetryTracingTracer}.
7689
*
7790
* @param tracer the {@link Tracer} to use for recording spans
7891
* @param apiTracerContext the {@link ApiTracerContext} to use for recording spans
7992
*/
8093
OpenTelemetryTracingTracer(Tracer tracer, ApiTracerContext apiTracerContext) {
81-
this.tracer = tracer;
82-
this.apiTracerContext = apiTracerContext;
83-
this.attemptSpanName = resolveAttemptSpanName(apiTracerContext);
84-
this.attemptAttributes = new HashMap<>();
85-
buildAttributes();
94+
this(tracer, apiTracerContext, resolveAttemptSpanName(apiTracerContext));
8695
}
8796

8897
/**
@@ -97,20 +106,54 @@ public void injectTraceContext(java.util.Map<String, String> carrier) {
97106
OpenTelemetryTracingTracer(
98107
Tracer tracer, ApiTracerContext apiTracerContext, String attemptSpanName) {
99108
this.tracer = tracer;
100-
this.attemptSpanName = attemptSpanName;
101109
this.apiTracerContext = apiTracerContext;
110+
this.attemptSpanName = attemptSpanName;
111+
this.operationSpanName = resolveOperationSpanName(apiTracerContext, attemptSpanName);
102112
this.attemptAttributes = new HashMap<>();
103113
buildAttributes();
114+
startOperationSpan();
115+
}
116+
117+
private void startOperationSpan() {
118+
SpanBuilder operationSpanBuilder = tracer.spanBuilder(operationSpanName);
119+
operationSpanBuilder.setSpanKind(SpanKind.INTERNAL);
120+
operationSpanBuilder.setAllAttributes(
121+
ObservabilityUtils.toOtelAttributes(this.attemptAttributes));
122+
this.operationSpan = operationSpanBuilder.startSpan();
123+
}
124+
125+
private static String resolveOperationSpanName(
126+
ApiTracerContext apiTracerContext, String attemptSpanName) {
127+
if (apiTracerContext.transport() == ApiTracerContext.Transport.GRPC
128+
&& !Strings.isNullOrEmpty(apiTracerContext.fullMethodName())) {
129+
return apiTracerContext.fullMethodName();
130+
} else if (apiTracerContext.transport() == ApiTracerContext.Transport.HTTP
131+
&& !Strings.isNullOrEmpty(apiTracerContext.httpMethod())
132+
&& !Strings.isNullOrEmpty(apiTracerContext.httpPathTemplate())) {
133+
return String.format(
134+
"%s %s", apiTracerContext.httpMethod(), apiTracerContext.httpPathTemplate());
135+
} else if (!Strings.isNullOrEmpty(apiTracerContext.fullMethodName())) {
136+
return apiTracerContext.fullMethodName();
137+
} else if (!Strings.isNullOrEmpty(attemptSpanName)) {
138+
if (attemptSpanName.endsWith("/attempt")) {
139+
return attemptSpanName.substring(0, attemptSpanName.length() - "/attempt".length());
140+
}
141+
return attemptSpanName;
142+
}
143+
return "operation";
104144
}
105145

106146
private static String resolveAttemptSpanName(ApiTracerContext apiTracerContext) {
107-
if (apiTracerContext.transport() == ApiTracerContext.Transport.GRPC) {
147+
if (apiTracerContext.transport() == ApiTracerContext.Transport.GRPC
148+
&& !Strings.isNullOrEmpty(apiTracerContext.fullMethodName())) {
108149
// gRPC Uses the full method name as span name.
109150
return apiTracerContext.fullMethodName();
110151
} else if (apiTracerContext.httpMethod() == null
111152
|| apiTracerContext.httpPathTemplate() == null) {
112153
// HTTP method name without necessary components defaults to the full method name
113-
return apiTracerContext.fullMethodName();
154+
return apiTracerContext.fullMethodName() != null
155+
? apiTracerContext.fullMethodName()
156+
: "attempt";
114157
} else {
115158
// We construct the span name with HTTP method and path template.
116159
return String.format(
@@ -142,12 +185,54 @@ public void attemptStarted(Object request, int attemptNumber) {
142185
// Attempt spans are of the CLIENT kind
143186
spanBuilder.setSpanKind(SpanKind.CLIENT);
144187

188+
if (operationSpan != null) {
189+
spanBuilder.setParent(io.opentelemetry.context.Context.current().with(operationSpan));
190+
}
191+
145192
// Pass the combined attributes to the new SpanBuilder method
146193
spanBuilder.setAllAttributes(ObservabilityUtils.toOtelAttributes(currentAttemptAttributes));
147194

148195
this.attemptSpan = spanBuilder.startSpan();
149196
}
150197

198+
@Override
199+
public void operationSucceeded() {
200+
recordErrorAndEndOperation(null);
201+
}
202+
203+
@Override
204+
public void operationCancelled() {
205+
recordErrorAndEndOperation(new CancellationException());
206+
}
207+
208+
@Override
209+
public void operationFailed(Throwable error) {
210+
recordErrorAndEndOperation(error);
211+
}
212+
213+
private void recordErrorAndEndOperation(@Nullable Throwable error) {
214+
if (attemptSpan != null) {
215+
endAttempt();
216+
}
217+
if (operationSpan == null) {
218+
return;
219+
}
220+
221+
Map<String, Object> responseAttributes =
222+
ObservabilityUtils.getResponseAttributes(error, this.apiTracerContext.transport());
223+
if (!responseAttributes.isEmpty()) {
224+
operationSpan.setAllAttributes(ObservabilityUtils.toOtelAttributes(responseAttributes));
225+
}
226+
227+
if (error != null && !Strings.isNullOrEmpty(error.getMessage())) {
228+
operationSpan.setAttribute(
229+
ObservabilityAttributes.STATUS_MESSAGE_ATTRIBUTE, error.getMessage());
230+
}
231+
232+
operationSpan.end();
233+
operationSpan = null;
234+
}
235+
151236
@Override
152237
public void attemptSucceeded() {
153238
recordErrorAndEndAttempt(null);

‎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: 150 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -64,15 +64,37 @@ class OpenTelemetryTracingTracerTest {
6464
@Mock private Tracer tracer;
6565
@Mock private SpanBuilder spanBuilder;
6666
@Mock private Span span;
67+
@Mock private SpanBuilder operationSpanBuilder;
68+
@Mock private Span operationSpan;
6769
private OpenTelemetryTracingTracer openTelemetryTracingTracer;
6870
private static final String ATTEMPT_SPAN_NAME = "Service/Method/attempt";
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+
org.mockito.Mockito.lenient().when(tracer.spanBuilder(anyString())).thenReturn(spanBuilder);
75+
org.mockito.Mockito.lenient()
76+
.when(spanBuilder.setSpanKind(any(SpanKind.class)))
77+
.thenReturn(spanBuilder);
78+
org.mockito.Mockito.lenient().when(spanBuilder.setParent(any())).thenReturn(spanBuilder);
79+
org.mockito.Mockito.lenient()
80+
.when(spanBuilder.setAllAttributes(any(Attributes.class)))
81+
.thenReturn(spanBuilder);
82+
org.mockito.Mockito.lenient().when(spanBuilder.startSpan()).thenReturn(span);
83+
84+
org.mockito.Mockito.lenient()
85+
.when(operationSpanBuilder.setSpanKind(any(SpanKind.class)))
86+
.thenReturn(operationSpanBuilder);
87+
org.mockito.Mockito.lenient()
88+
.when(operationSpanBuilder.setAllAttributes(any(Attributes.class)))
89+
.thenReturn(operationSpanBuilder);
90+
org.mockito.Mockito.lenient().when(operationSpanBuilder.startSpan()).thenReturn(operationSpan);
91+
org.mockito.Mockito.lenient()
92+
.when(operationSpan.storeInContext(any(io.opentelemetry.context.Context.class)))
93+
.thenAnswer(invocation -> invocation.getArgument(0));
94+
org.mockito.Mockito.lenient()
95+
.when(tracer.spanBuilder("Service/Method"))
96+
.thenReturn(operationSpanBuilder);
97+
7698
openTelemetryTracingTracer =
7799
new OpenTelemetryTracingTracer(tracer, ApiTracerContext.empty(), ATTEMPT_SPAN_NAME);
78100
}
@@ -680,4 +702,128 @@ void testInjectTraceContext_addsHeaders() {
680702
assertThat(carrier.get("traceparent")).contains("00000000000000000000000000000001");
681703
assertThat(carrier.get("traceparent")).contains("0000000000000002");
682704
}
705+
706+
@Test
707+
void testOperationLifecycle_startsOperationSpanWithInternalKind() {
708+
verify(tracer).spanBuilder("Service/Method");
709+
verify(operationSpanBuilder).setSpanKind(SpanKind.INTERNAL);
710+
}
711+
712+
@Test
713+
void testAttemptStarted_setsParentToOperationSpan() {
714+
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
715+
verify(spanBuilder).setParent(any(io.opentelemetry.context.Context.class));
716+
}
717+
718+
@Test
719+
void testOperationSucceeded_endsOperationSpanWithStatusOk() {
720+
openTelemetryTracingTracer.operationSucceeded();
721+
722+
ArgumentCaptor<Attributes> attrsCaptor = ArgumentCaptor.forClass(Attributes.class);
723+
verify(operationSpan).setAllAttributes(attrsCaptor.capture());
724+
verify(operationSpan).end();
725+
726+
assertThat(attrsCaptor.getValue().asMap())
727+
.containsEntry(
728+
AttributeKey.stringKey(ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE), "OK");
729+
}
730+
731+
@Test
732+
void testOperationFailed_endsOperationSpanWithErrorAttributes() {
733+
ApiException exception =
734+
new ApiException(
735+
"custom failure",
736+
null,
737+
new StatusCode() {
738+
@Override
739+
public Code getCode() {
740+
return Code.UNAVAILABLE;
741+
}
742+
743+
@Override
744+
public Object getTransportCode() {
745+
return null;
746+
}
747+
},
748+
true);
749+
750+
openTelemetryTracingTracer.operationFailed(exception);
751+
752+
ArgumentCaptor<Attributes> attrsCaptor = ArgumentCaptor.forClass(Attributes.class);
753+
verify(operationSpan).setAllAttributes(attrsCaptor.capture());
754+
verify(operationSpan)
755+
.setAttribute(ObservabilityAttributes.STATUS_MESSAGE_ATTRIBUTE, "custom failure");
756+
verify(operationSpan).end();
757+
758+
Map<AttributeKey<?>, Object> captured = attrsCaptor.getValue().asMap();
759+
assertThat(captured)
760+
.containsEntry(
761+
AttributeKey.stringKey(ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE),
762+
"UNAVAILABLE");
763+
assertThat(captured)
764+
.containsEntry(
765+
AttributeKey.stringKey(ObservabilityAttributes.ERROR_TYPE_ATTRIBUTE), "ApiException");
766+
}
767+
768+
@Test
769+
void testOperationCancelled_endsOperationSpanWithCancelledStatus() {
770+
openTelemetryTracingTracer.operationCancelled();
771+
772+
ArgumentCaptor<Attributes> attrsCaptor = ArgumentCaptor.forClass(Attributes.class);
773+
verify(operationSpan).setAllAttributes(attrsCaptor.capture());
774+
verify(operationSpan).end();
775+
776+
assertThat(attrsCaptor.getValue().asMap())
777+
.containsEntry(
778+
AttributeKey.stringKey(ObservabilityAttributes.RPC_RESPONSE_STATUS_ATTRIBUTE),
779+
"CANCELLED");
780+
}
781+
782+
@Test
783+
void testInScope_withAttemptSpan() {
784+
io.opentelemetry.context.Scope mockScope =
785+
org.mockito.Mockito.mock(io.opentelemetry.context.Scope.class);
786+
org.mockito.Mockito.when(span.makeCurrent()).thenReturn(mockScope);
787+
788+
openTelemetryTracingTracer.attemptStarted(new Object(), 1);
789+
try (ApiTracer.Scope scope = openTelemetryTracingTracer.inScope()) {
790+
verify(span).makeCurrent();
791+
}
792+
verify(mockScope).close();
793+
}
794+
795+
@Test
796+
void testInScope_withOperationSpanFallback() {
797+
io.opentelemetry.context.Scope mockScope =
798+
org.mockito.Mockito.mock(io.opentelemetry.context.Scope.class);
799+
org.mockito.Mockito.when(operationSpan.makeCurrent()).thenReturn(mockScope);
800+
801+
try (ApiTracer.Scope scope = openTelemetryTracingTracer.inScope()) {
802+
verify(operationSpan).makeCurrent();
803+
}
804+
verify(mockScope).close();
805+
}
806+
807+
@Test
808+
void testInjectTraceContext_withOperationSpanFallback() {
809+
io.opentelemetry.api.trace.SpanContext mockSpanContext =
810+
io.opentelemetry.api.trace.SpanContext.create(
811+
"00000000000000000000000000000003",
812+
"0000000000000004",
813+
io.opentelemetry.api.trace.TraceFlags.getSampled(),
814+
io.opentelemetry.api.trace.TraceState.getDefault());
815+
io.opentelemetry.api.trace.Span realSpan =
816+
io.opentelemetry.api.trace.Span.wrap(mockSpanContext);
817+
org.mockito.Mockito.when(operationSpanBuilder.startSpan()).thenReturn(realSpan);
818+
819+
openTelemetryTracingTracer =
820+
new OpenTelemetryTracingTracer(tracer, ApiTracerContext.empty(), ATTEMPT_SPAN_NAME);
821+
822+
Map<String, String> carrier = new java.util.HashMap<>();
823+
openTelemetryTracingTracer.injectTraceContext(carrier);
824+
825+
assertThat(carrier).containsKey("traceparent");
826+
assertThat(carrier.get("traceparent")).contains("00000000000000000000000000000003");
827+
assertThat(carrier.get("traceparent")).contains("0000000000000004");
828+
}
683829
}

0 commit comments

Comments
 (0)