Merge pull request #8778 from Ladicek/nested-mdc-scopes-improvement

Improve how nested MDC scopes are handled in the Jaeger extension
quarkusio · Apr 23, 2020 · def0378 · def0378
2 parents ed96897 + b2f595e
commit def0378
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 10 deletions.
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
@@ -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
@@ -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
@@ -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