Add MDCScopeManager for correlating logs with trace context (#718)

* Support MDC initial Signed-off-by: avim <[email protected]> * add license Signed-off-by: avim <[email protected]> * add license Signed-off-by: avim <[email protected]> * PR feedback add tests Signed-off-by: avim <[email protected]> * fix style Signed-off-by: avim <[email protected]> * PR test feedback Signed-off-by: avim <[email protected]> * test feedback Signed-off-by: avim <[email protected]> * remove empty line Signed-off-by: avim <[email protected]> * add license expand imports Signed-off-by: avim <[email protected]> * allow 4 consecutive capital letters Signed-off-by: avim <[email protected]> * fix style errors Signed-off-by: avim <[email protected]> * revert global stylecheck config escape only MDC classes Signed-off-by: avim <[email protected]> Co-authored-by: avim <[email protected]>
jaegertracing · Jun 27, 2020 · 74f4aa2 · 74f4aa2
1 parent 760fdda
commit 74f4aa2
Show file tree

Hide file tree

Showing 3 changed files with 274 additions and 0 deletions.
diff --git a/config/checkstyle/suppressions.xml b/config/checkstyle/suppressions.xml
@@ -6,4 +6,8 @@
 <suppressions>
   <suppress checks="AbbreviationAsWordInName"
     files="SpanContext.java"/>
+  <suppress checks="AbbreviationAsWordInName"
+            files="MDCScopeManagerTest.java"/>
+  <suppress checks="AbbreviationAsWordInName"
+            files="MDCScopeManager.java"/>
 </suppressions>
diff --git a/jaeger-core/src/main/java/io/jaegertracing/internal/MDCScopeManager.java b/jaeger-core/src/main/java/io/jaegertracing/internal/MDCScopeManager.java
@@ -0,0 +1,128 @@
+/*
+ * Copyright (c) 2020, The Jaeger Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License
+ * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+ * or implied. See the License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.jaegertracing.internal;
+
+import io.opentracing.Scope;
+import io.opentracing.ScopeManager;
+import io.opentracing.Span;
+import io.opentracing.util.ThreadLocalScopeManager;
+import org.slf4j.MDC;
+
+public class MDCScopeManager implements ScopeManager {
+  private final ScopeManager wrappedScopeManager;
+  private final String mdcTraceIdKey;
+  private final String mdcSpanIdKey;
+  private final String mdcSampledKey;
+
+  private MDCScopeManager(Builder builder) {
+    this.wrappedScopeManager = builder.scopeManager;
+    this.mdcTraceIdKey = builder.mdcTraceIdKey;
+    this.mdcSpanIdKey = builder.mdcSpanIdKey;
+    this.mdcSampledKey = builder.mdcSampledKey;
+  }
+
+  @Override
+  public Scope activate(Span span) {
+    return new MDCScope(wrappedScopeManager.activate(span), span);
+  }
+
+  @Override
+  public Span activeSpan() {
+    return wrappedScopeManager.activeSpan();
+  }
+
+  /**
+   * Builds an {@link MDCScopeManager} with options.
+   * Calling {@code new MDCScopeManager.Builder().build()}
+   * Builds an {@link MDCScopeManager} with configuration as follows:
+   * mdcTraceIDKey set to "traceId"
+   * mdcSpanIdKey set to "spanId"
+   * mdcSampledKey set to "sampled"
+   */
+  public static class Builder {
+    private ScopeManager scopeManager = new ThreadLocalScopeManager();
+    private String mdcTraceIdKey = "traceId";
+    private String mdcSpanIdKey = "spanId";
+    private String mdcSampledKey = "sampled";
+
+    public Builder withScopeManager(ScopeManager scopeManager) {
+      this.scopeManager = scopeManager;
+      return this;
+    }
+
+    public Builder withMDCTraceIdKey(String mdcTraceIdKey) {
+      this.mdcTraceIdKey = mdcTraceIdKey;
+      return this;
+    }
+
+    public Builder withMDCSpanIdKey(String mdcSpanIdKey) {
+      this.mdcSpanIdKey = mdcSpanIdKey;
+      return this;
+    }
+
+    public Builder withMDCSampledKey(String mdcSampledKey) {
+      this.mdcSampledKey = mdcSampledKey;
+      return this;
+    }
+
+    public MDCScopeManager build() {
+      return new MDCScopeManager(this);
+    }
+
+  }
+
+  private class MDCScope implements Scope {
+    private final Scope wrappedScope;
+    private final String previousTraceId;
+    private final String previousSpanId;
+    private final String previousSampled;
+
+    /**
+     * mdcScope.
+     */
+    MDCScope(Scope scope, Span span) {
+      this.wrappedScope = scope;
+      this.previousTraceId = MDC.get(mdcTraceIdKey);
+      this.previousSpanId = MDC.get(mdcSpanIdKey);
+      this.previousSampled = MDC.get(mdcSampledKey);
+
+      if (span.context() instanceof JaegerSpanContext) {
+        putContext((JaegerSpanContext) span.context());
+      }
+    }
+
+    protected void putContext(JaegerSpanContext spanContext) {
+      replace(mdcTraceIdKey, spanContext.toTraceId());
+      replace(mdcSpanIdKey, spanContext.toSpanId());
+      replace(mdcSampledKey, String.valueOf(spanContext.isSampled()));
+    }
+
+    private void replace(String key, String value) {
+      if (value == null) {
+        MDC.remove(key);
+      } else {
+        MDC.put(key, value);
+      }
+    }
+
+    @Override
+    public void close() {
+      wrappedScope.close();
+      replace(mdcTraceIdKey, previousTraceId);
+      replace(mdcSpanIdKey, previousSpanId);
+      replace(mdcSampledKey, previousSampled);
+    }
+  }
+}
diff --git a/jaeger-core/src/test/java/io/jaegertracing/internal/MDCScopeManagerTest.java b/jaeger-core/src/test/java/io/jaegertracing/internal/MDCScopeManagerTest.java
@@ -0,0 +1,142 @@
+/*
+ * Copyright (c) 2020, The Jaeger Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License
+ * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+ * or implied. See the License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.jaegertracing.internal;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
+
+import io.jaegertracing.internal.reporters.InMemoryReporter;
+import io.jaegertracing.internal.samplers.ConstSampler;
+import io.opentracing.Scope;
+import io.opentracing.ScopeManager;
+import io.opentracing.Span;
+import io.opentracing.Tracer;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+import org.slf4j.MDC;
+
+public class MDCScopeManagerTest {
+
+  private InMemoryReporter reporter;
+  private JaegerTracer defaultTracer;
+  private static final String TRACE_ID = "traceId";
+  private static final String SPAN_ID = "spanId";
+  private static final String SAMPLED = "sampled";
+
+
+  @Before
+  public void setUp() {
+    reporter = new InMemoryReporter();
+    defaultTracer = createTracer(new MDCScopeManager.Builder().build());
+  }
+
+  @Test
+  public void testActiveSpan() {
+    Span mockSpan = Mockito.mock(JaegerSpan.class);
+    try (Scope scope = defaultTracer.activateSpan(mockSpan)) {
+      assertEquals(mockSpan, defaultTracer.activeSpan());
+    }
+  }
+
+  @Test
+  public void testNestedSpans() {
+    Span parentSpan = defaultTracer.buildSpan("parent").start();
+    try (Scope scope = defaultTracer.activateSpan(parentSpan)) {
+      assertSpanContextEqualsToMDC((JaegerSpanContext) parentSpan.context(), TRACE_ID, SPAN_ID, SAMPLED);
+      Span childSpan = defaultTracer.buildSpan("child").start();
+      try (Scope childScope = defaultTracer.activateSpan(childSpan)) {
+        assertSpanContextEqualsToMDC((JaegerSpanContext) childSpan.context(), TRACE_ID, SPAN_ID, SAMPLED);
+      }
+      assertSpanContextEqualsToMDC((JaegerSpanContext) parentSpan.context(), TRACE_ID, SPAN_ID, SAMPLED);
+    }
+    assertNullMDCKeys(TRACE_ID, SPAN_ID, SAMPLED);
+  }
+
+  @Test
+  public void testDefaultCreation() {
+    Span span = defaultTracer.buildSpan("test Default").start();
+    Scope scope = defaultTracer.activateSpan(span);
+
+    assertSpanContextEqualsToMDC((JaegerSpanContext) span.context(), TRACE_ID, SPAN_ID, SAMPLED);
+
+    scope.close();
+    assertNullMDCKeys(TRACE_ID, SPAN_ID, SAMPLED);
+  }
+
+  @Test
+  public void testCustomKeysCreation() {
+    ScopeManager mdcScopeManager = new MDCScopeManager
+            .Builder()
+            .withMDCTraceIdKey("CustomTraceId")
+            .withMDCSampledKey("customSampled")
+            .withMDCSpanIdKey("customSpanId")
+            .build();
+
+    Tracer tracer = createTracer(mdcScopeManager);
+    Span span = tracer.buildSpan("testCustomKeysCreation").start();
+    Scope scope = tracer.activateSpan(span);
+
+    assertSpanContextEqualsToMDC((JaegerSpanContext) span.context(), "CustomTraceId", "customSpanId", "customSampled");
+
+    scope.close();
+
+    assertNullMDCKeys("CustomTraceId", "customSampled", "customSpanId");
+  }
+
+  @Test
+  public void testCustomAndDefaultKeysCreation() {
+    ScopeManager mdcScopeManager = new MDCScopeManager
+            .Builder()
+            .withMDCSampledKey("customSampled")
+            .withMDCSpanIdKey("customSpanId")
+            .build();
+
+    Tracer tracer = createTracer(mdcScopeManager);
+    Span span = tracer.buildSpan("testCustomAndDefaultKeysCreation").start();
+    Scope scope = tracer.activateSpan(span);
+
+    assertSpanContextEqualsToMDC((JaegerSpanContext) span.context(), TRACE_ID, "customSpanId", "customSampled");
+
+    scope.close();
+
+    assertNullMDCKeys(TRACE_ID, "customSpanId", "customSampled");
+  }
+
+  private JaegerTracer createTracer(ScopeManager scopeManager) {
+    return new JaegerTracer.Builder("MDCScopeManagerTest")
+            .withReporter(reporter)
+            .withSampler(new ConstSampler(true))
+            .withScopeManager(scopeManager)
+            .build();
+  }
+
+  private void assertSpanContextEqualsToMDC(JaegerSpanContext context,
+                                            String traceIDKey,
+                                            String spanIdKey,
+                                            String sampledKey) {
+
+    assertEquals(context.toTraceId(), MDC.get(traceIDKey));
+    assertEquals(context.toSpanId(), MDC.get(spanIdKey));
+    assertEquals(String.valueOf(context.isSampled()), MDC.get(sampledKey));
+  }
+
+  private void assertNullMDCKeys(String traceIDKey, String spanIdKey, String sampleKey) {
+    assertNull(MDC.get(traceIDKey));
+    assertNull(MDC.get(spanIdKey));
+    assertNull(MDC.get(sampleKey));
+  }
+
+}