Add MDCScopeManager for correlating logs with trace context

jaeger-core/src/main/java/io/jaegertracing/internal/MDCScopeManager.java

@@ -0,0 +1,128 @@
+/*
+ * Copyright (c) 2018, 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);
+    }
+  }
+}

jaeger-core/src/test/java/io/jaegertracing/internal/MDCScopeManagerTest.java

@@ -0,0 +1,103 @@
+package io.jaegertracing.internal;
+
+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;
+
+import static org.junit.Assert.*;
+
+public class MDCScopeManagerTest {
+
+  InMemoryReporter reporter;
+  JaegerTracer defaultTracer;
+
+  @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 testDefaultCreation() {
+    Span span = defaultTracer.buildSpan("test Default").start();
+    Scope scope = defaultTracer.activateSpan(span);
+    assertNotNull(MDC.get("traceId"));
+    assertNotNull(MDC.get("spanId"));
+    assertNotNull(MDC.get("sampled"));
+
+    scope.close();
+
+    assertNull(MDC.get("traceId"));
+    assertNull(MDC.get("spanId"));
+    assertNull(MDC.get("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);
+
+    assertNotNull(MDC.get("CustomTraceId"));
+    assertNotNull(MDC.get("customSampled"));
+    assertNotNull(MDC.get("customSpanId"));
+
+    scope.close();
+
+  }
+
+  @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);
+
+    assertNotNull(MDC.get("traceId"));
+    assertNotNull(MDC.get("customSampled"));
+    assertNotNull(MDC.get("customSpanId"));
+
+    scope.close();
+
+  }
+
+  private JaegerTracer createTracer(ScopeManager scopeManager) {
+    return new JaegerTracer.Builder("MDCScopeManagerTest")
+            .withReporter(reporter)
+            .withSampler(new ConstSampler(true))
+            .withScopeManager(scopeManager)
+            .build();
+  }
+
+}