Skip to content

Commit

Permalink
OTel Scope.close() warning improvement
Browse files Browse the repository at this point in the history
  • Loading branch information
brunobat committed Sep 12, 2023
1 parent 718a8d5 commit 095785b
Show file tree
Hide file tree
Showing 4 changed files with 107 additions and 11 deletions.
5 changes: 5 additions & 0 deletions extensions/opentelemetry/runtime/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,11 @@
<artifactId>assertj-core</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.mockito</groupId>
<artifactId>mockito-core</artifactId>
<scope>test</scope>
</dependency>
</dependencies>

<build>
Expand Down
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
package io.quarkus.opentelemetry.runtime;

import java.util.Collections;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.SpanContext;
Expand All @@ -16,6 +19,7 @@ public final class OpenTelemetryUtil {
public static final String SPAN_ID = "spanId";
public static final String SAMPLED = "sampled";
public static final String PARENT_ID = "parentId";
private static final Set<String> SPAN_DATA_KEYS = Set.of(TRACE_ID, SPAN_ID, SAMPLED, PARENT_ID);

private OpenTelemetryUtil() {
}
Expand Down Expand Up @@ -54,22 +58,43 @@ public static Map<String, String> convertKeyValueListToMap(List<String> headers)
* @param vertxContext vertx context
*/
public static void setMDCData(Context context, io.vertx.core.Context vertxContext) {
setMDCData(getSpanData(context), vertxContext);
}

public static void setMDCData(Map<String, String> spanData, io.vertx.core.Context vertxContext) {
if (spanData != null && !spanData.isEmpty()) {
for (Entry<String, String> entry : spanData.entrySet()) {
if (SPAN_DATA_KEYS.contains(entry.getKey())) {
VertxMDC.INSTANCE.put(entry.getKey(), entry.getValue(), vertxContext);
}
}
}
}

/**
* Gets current span data from the MDC context.
*
* @param context opentelemetry context
*/
public static Map<String, String> getSpanData(Context context) {
if (context == null) {
return Collections.emptyMap();
}
Span span = Span.fromContextOrNull(context);
Map<String, String> spanData = new HashMap<>();
if (span != null) {
SpanContext spanContext = span.getSpanContext();
VertxMDC vertxMDC = VertxMDC.INSTANCE;
vertxMDC.put(SPAN_ID, spanContext.getSpanId(), vertxContext);
vertxMDC.put(TRACE_ID, spanContext.getTraceId(), vertxContext);
vertxMDC.put(SAMPLED, Boolean.toString(spanContext.isSampled()), vertxContext);
spanData.put(SPAN_ID, spanContext.getSpanId());
spanData.put(TRACE_ID, spanContext.getTraceId());
spanData.put(SAMPLED, Boolean.toString(spanContext.isSampled()));
if (span instanceof ReadableSpan) {
SpanContext parentSpanContext = ((ReadableSpan) span).getParentSpanContext();
if (parentSpanContext.isValid()) {
vertxMDC.put(PARENT_ID, parentSpanContext.getSpanId(), vertxContext);
} else {
vertxMDC.remove(PARENT_ID, vertxContext);
if (parentSpanContext != null && parentSpanContext.isValid()) {
spanData.put(PARENT_ID, parentSpanContext.getSpanId());
}
}
}
return spanData;
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
import static io.quarkus.vertx.core.runtime.context.VertxContextSafetyToggle.setContextSafe;
import static io.smallrye.common.vertx.VertxContext.isDuplicatedContext;

import java.util.Map;

import org.jboss.logging.Logger;

import io.opentelemetry.context.Context;
Expand Down Expand Up @@ -64,14 +66,18 @@ public Scope attach(io.vertx.core.Context vertxContext, Context toAttach) {
return Scope.noop();
}
vertxContext.putLocal(OTEL_CONTEXT, toAttach);
OpenTelemetryUtil.setMDCData(toAttach, vertxContext);
final Map<String, String> spanDataToAttach = OpenTelemetryUtil.getSpanData(toAttach);
OpenTelemetryUtil.setMDCData(spanDataToAttach, vertxContext);

return new Scope() {

@Override
public void close() {
if (getContext(vertxContext) != toAttach) {
log.warn("Context in storage not the expected context, Scope.close was not called correctly");
final Context before = getContext(vertxContext);
if (before != toAttach) {
log.warn("Context in storage not the expected context, Scope.close was not called correctly. Details:" +
" OTel context before: " + OpenTelemetryUtil.getSpanData(before) +
". OTel context toAttach: " + spanDataToAttach);
}

if (beforeAttach == null) {
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
package io.quarkus.opentelemetry.runtime;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.mockito.Mockito.mock;

import java.util.AbstractMap;
import java.util.Arrays;
import java.util.Collections;
Expand All @@ -8,6 +11,13 @@
import org.assertj.core.api.Assertions;
import org.junit.jupiter.api.Test;

import io.opentelemetry.api.trace.Span;
import io.opentelemetry.api.trace.SpanBuilder;
import io.opentelemetry.api.trace.Tracer;
import io.opentelemetry.context.Context;
import io.opentelemetry.sdk.trace.SdkTracerProvider;
import io.opentelemetry.sdk.trace.SpanProcessor;

public class OpenTelemetryUtilTest {

@Test
Expand Down Expand Up @@ -58,4 +68,54 @@ public void testConvertKeyValueListToMap_empty_value() {
.convertKeyValueListToMap(Collections.emptyList());
Assertions.assertThat(actual).containsExactly();
}

@Test
public void testGetSpanData() {
SpanProcessor mockedSpanProcessor = mock(SpanProcessor.class);

SdkTracerProvider tracerSdkFactory = SdkTracerProvider.builder()
.addSpanProcessor(mockedSpanProcessor)
.build();
Tracer spanBuilderSdkTest = tracerSdkFactory.get("SpanBuilderSdkTest");
SpanBuilder spanBuilder = spanBuilderSdkTest.spanBuilder("SpanName");

Span parent = spanBuilder.startSpan();
Context contextParent = Context.current().with(parent);

Span child = spanBuilder.setParent(contextParent).startSpan();
Context contextChild = Context.current().with(child);

Map<String, String> actual = OpenTelemetryUtil.getSpanData(contextChild);
assertEquals(4, actual.size());
assertEquals(child.getSpanContext().getSpanId(), actual.get("spanId"));
assertEquals(child.getSpanContext().getTraceId(), actual.get("traceId"));
assertEquals("true", actual.get("sampled"));
assertEquals(parent.getSpanContext().getSpanId(), actual.get("parentId"));
}

@Test
public void testGetSpanData_noParent() {
SpanProcessor mockedSpanProcessor = mock(SpanProcessor.class);
SdkTracerProvider tracerSdkFactory = SdkTracerProvider.builder()
.addSpanProcessor(mockedSpanProcessor)
.build();
Tracer spanBuilderSdkTest = tracerSdkFactory.get("SpanBuilderSdkTest");

SpanBuilder spanBuilder = spanBuilderSdkTest.spanBuilder("SpanName");

Span child = spanBuilder.startSpan();
Context contextChild = Context.current().with(child);

Map<String, String> actual = OpenTelemetryUtil.getSpanData(contextChild);
assertEquals(3, actual.size());
assertEquals(child.getSpanContext().getSpanId(), actual.get("spanId"));
assertEquals(child.getSpanContext().getTraceId(), actual.get("traceId"));
assertEquals("true", actual.get("sampled"));
}

@Test
public void testGetSpanData_nullValue() {
Map<String, String> actual = OpenTelemetryUtil.getSpanData(null);
assertEquals(0, actual.size());
}
}

0 comments on commit 095785b

Please sign in to comment.