diff --git a/extensions/jaeger/deployment/src/test/java/io/quarkus/jaeger/test/NestedMdcScopesTest.java b/extensions/jaeger/deployment/src/test/java/io/quarkus/jaeger/test/NestedMdcScopesTest.java index 8513a558fb4ac..861c44da57484 100644 --- a/extensions/jaeger/deployment/src/test/java/io/quarkus/jaeger/test/NestedMdcScopesTest.java +++ b/extensions/jaeger/deployment/src/test/java/io/quarkus/jaeger/test/NestedMdcScopesTest.java @@ -2,6 +2,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; import java.util.Map; @@ -23,26 +24,28 @@ public void mdcIsRestoredCorrectly() { assertNull(mdcScopeManager.active()); assertNull(threadLocalScopeManager.active()); + assertNull(MDC.get("traceId")); JaegerSpanContext span = new JaegerSpanContext(1, 1, 1, 1, Byte.valueOf("0")); Scope scope = mdcScopeManager.activate(new TestSpan(span), true); - assertEquals(span, threadLocalScopeManager.active().span().context()); + assertSame(span, threadLocalScopeManager.active().span().context()); assertEquals("10000000000000001", MDC.get("traceId")); JaegerSpanContext subSpan = new JaegerSpanContext(2, 2, 2, 1, Byte.valueOf("0")); Scope subScope = mdcScopeManager.activate(new TestSpan(subSpan), true); - assertEquals(subSpan, threadLocalScopeManager.active().span().context()); + assertSame(subSpan, threadLocalScopeManager.active().span().context()); assertEquals("20000000000000002", MDC.get("traceId")); subScope.close(); - assertEquals(span, threadLocalScopeManager.active().span().context()); + assertSame(span, threadLocalScopeManager.active().span().context()); assertEquals("10000000000000001", MDC.get("traceId")); scope.close(); assertNull(mdcScopeManager.active()); assertNull(threadLocalScopeManager.active()); + assertNull(MDC.get("traceId")); } static class TestSpan implements Span { diff --git a/extensions/jaeger/runtime/src/main/java/io/quarkus/jaeger/runtime/MDCScope.java b/extensions/jaeger/runtime/src/main/java/io/quarkus/jaeger/runtime/MDCScope.java index 5fe89f6017303..8e8ae138f46e7 100644 --- a/extensions/jaeger/runtime/src/main/java/io/quarkus/jaeger/runtime/MDCScope.java +++ b/extensions/jaeger/runtime/src/main/java/io/quarkus/jaeger/runtime/MDCScope.java @@ -19,11 +19,15 @@ public class MDCScope implements Scope { private static final String SAMPLED = "sampled"; private final Scope wrapped; - private final Scope toRestore; + private final Object originalTraceId; + private final Object originalSpanId; + private final Object originalSampled; - public MDCScope(Scope toRestore, Scope scope) { - this.toRestore = toRestore; + public MDCScope(Scope scope) { this.wrapped = scope; + this.originalTraceId = MDC.get(TRACE_ID); + this.originalSpanId = MDC.get(SPAN_ID); + this.originalSampled = MDC.get(SAMPLED); if (scope.span().context() instanceof JaegerSpanContext) { putContext((JaegerSpanContext) scope.span().context()); } @@ -36,8 +40,14 @@ public void close() { MDC.remove(SPAN_ID); MDC.remove(SAMPLED); - if (toRestore != null && toRestore.span().context() instanceof JaegerSpanContext) { - putContext(((JaegerSpanContext) toRestore.span().context())); + if (originalTraceId != null) { + MDC.put(TRACE_ID, originalTraceId); + } + if (originalSpanId != null) { + MDC.put(SPAN_ID, originalSpanId); + } + if (originalSampled != null) { + MDC.put(SAMPLED, originalSampled); } } diff --git a/extensions/jaeger/runtime/src/main/java/io/quarkus/jaeger/runtime/MDCScopeManager.java b/extensions/jaeger/runtime/src/main/java/io/quarkus/jaeger/runtime/MDCScopeManager.java index 805c8ee135431..3c1cedef03266 100644 --- a/extensions/jaeger/runtime/src/main/java/io/quarkus/jaeger/runtime/MDCScopeManager.java +++ b/extensions/jaeger/runtime/src/main/java/io/quarkus/jaeger/runtime/MDCScopeManager.java @@ -14,8 +14,7 @@ public MDCScopeManager(ScopeManager scopeManager) { @Override public Scope activate(Span span, boolean finishSpanOnClose) { - Scope current = wrapped.active(); - return new MDCScope(current, wrapped.activate(span, finishSpanOnClose)); + return new MDCScope(wrapped.activate(span, finishSpanOnClose)); } @Override