slf4j MDC autoinjection integration

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java

@@ -92,7 +92,6 @@ public class AgentInstaller {
             .or(nameStartsWith("org.aspectj."))
             .or(nameStartsWith("org.groovy."))
             .or(nameStartsWith("com.p6spy."))
-            .or(nameStartsWith("org.slf4j."))
             .or(nameStartsWith("com.newrelic."))
             .or(nameContains("javassist"))
             .or(nameContains(".asm."))

dd-java-agent/instrumentation/slf4j-mdc/slf4j-mdc.gradle

@@ -0,0 +1,23 @@
+apply from: "${rootDir}/gradle/java.gradle"
+
+muzzle {
+  pass {
+    group = 'org.slf4j'
+    module = 'slf4j-api'
+    versions = '(,)'
+  }
+}
+
+dependencies {
+  compile project(':dd-trace-api')
+  compile project(':dd-java-agent:agent-tooling')
+
+  // no need to compileOnly against slf4j. Included with transitive dependency.
+
+  compile deps.bytebuddy
+  compile deps.opentracing
+  annotationProcessor deps.autoservice
+  implementation deps.autoservice
+
+  testCompile project(':dd-java-agent:testing')
+}

dd-java-agent/instrumentation/slf4j-mdc/src/main/java/datadog/trace/instrumentation/slf4j/mdc/MDCInjectionInstrumentation.java

@@ -0,0 +1,115 @@
+package datadog.trace.instrumentation.slf4j.mdc;
+
+import static net.bytebuddy.matcher.ElementMatchers.isTypeInitializer;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.Config;
+import datadog.trace.api.CorrelationIdentifier;
+import datadog.trace.api.GlobalTracer;
+import datadog.trace.context.ScopeListener;
+import java.lang.reflect.Method;
+import java.security.ProtectionDomain;
+import java.util.Collections;
+import java.util.Map;
+import lombok.extern.slf4j.Slf4j;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.utility.JavaModule;
+
+@AutoService(Instrumenter.class)
+public class MDCInjectionInstrumentation extends Instrumenter.Default {
+  public static final String MDC_INSTRUMENTATION_NAME = "mdc";
+
+  // Intentionally doing the string replace to bypass gradle shadow rename
+  // mdcClassName = org.slf4j.MDC
+  private static final String mdcClassName = "org.TMP.MDC".replaceFirst("TMP", "slf4j");
+
+  public MDCInjectionInstrumentation() {
+    super(MDC_INSTRUMENTATION_NAME);
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    final String enableInjection = getPropOrEnv("dd." + Config.LOGS_INJECTION_ENABLED);
+    return null == enableInjection
+        ? Config.DEFAULT_LOGS_INJECTION_ENABLED
+        : Boolean.parseBoolean(enableInjection);
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return named(mdcClassName);
+  }
+
+  @Override
+  public void postMatch(
+      TypeDescription typeDescription,
+      ClassLoader classLoader,
+      JavaModule module,
+      Class<?> classBeingRedefined,
+      ProtectionDomain protectionDomain) {
+    if (classBeingRedefined != null) {
+      MDCAdvice.mdcClassInitialized(classBeingRedefined);
+    }
+  }
+
+  @Override
+  public Map<? extends ElementMatcher, String> transformers() {
+    return Collections.<ElementMatcher, String>singletonMap(
+        isTypeInitializer(), MDCAdvice.class.getName());
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {MDCAdvice.class.getName() + "$MDCScopeListener"};
+  }
+
+  public static class MDCAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void mdcClassInitialized(@Advice.Origin final Class mdcClass) {
+      try {
+        final Method putMethod = mdcClass.getMethod("put", String.class, String.class);
+        final Method removeMethod = mdcClass.getMethod("remove", String.class);
+        GlobalTracer.get().addScopeListener(new MDCScopeListener(putMethod, removeMethod));
+      } catch (NoSuchMethodException e) {
+        org.slf4j.LoggerFactory.getLogger(mdcClass).debug("Failed to add MDC span listener", e);
+      }
+    }
+
+    @Slf4j
+    public static class MDCScopeListener implements ScopeListener {
+      private final Method putMethod;
+      private final Method removeMethod;
+
+      public MDCScopeListener(final Method putMethod, final Method removeMethod) {
+        this.putMethod = putMethod;
+        this.removeMethod = removeMethod;
+      }
+
+      @Override
+      public void afterScopeActivated() {
+        try {
+          putMethod.invoke(
+              null, CorrelationIdentifier.getTraceIdKey(), CorrelationIdentifier.getTraceId());
+          putMethod.invoke(
+              null, CorrelationIdentifier.getSpanIdKey(), CorrelationIdentifier.getSpanId());
+        } catch (Exception e) {
+          log.debug("Exception setting mdc context", e);
+        }
+      }
+
+      @Override
+      public void afterScopeClose() {
+        try {
+          removeMethod.invoke(null, CorrelationIdentifier.getTraceIdKey());
+          removeMethod.invoke(null, CorrelationIdentifier.getSpanIdKey());
+        } catch (Exception e) {
+          log.debug("Exception removing mdc context", e);
+        }
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/slf4j-mdc/src/test/groovy/Slf4jMDCTest.groovy

@@ -0,0 +1,92 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.CorrelationIdentifier
+import io.opentracing.Scope
+import io.opentracing.util.GlobalTracer
+import org.slf4j.MDC
+
+import java.util.concurrent.atomic.AtomicReference
+
+class Slf4jMDCTest extends AgentTestRunner {
+  static {
+    System.setProperty("dd.logs.injection", "true")
+  }
+
+  def "mdc shows trace and span ids for active scope"() {
+    when:
+    MDC.put("foo", "bar")
+    Scope rootScope = GlobalTracer.get().buildSpan("root").startActive(true)
+
+    then:
+    MDC.get(CorrelationIdentifier.getTraceIdKey()) == CorrelationIdentifier.getTraceId()
+    MDC.get(CorrelationIdentifier.getSpanIdKey()) == CorrelationIdentifier.getSpanId()
+    MDC.get("foo") == "bar"
+
+    when:
+    Scope childScope = GlobalTracer.get().buildSpan("child").startActive(true)
+
+    then:
+    MDC.get(CorrelationIdentifier.getTraceIdKey()) == CorrelationIdentifier.getTraceId()
+    MDC.get(CorrelationIdentifier.getSpanIdKey()) == CorrelationIdentifier.getSpanId()
+    MDC.get("foo") == "bar"
+
+    when:
+    childScope.close()
+
+    then:
+    MDC.get(CorrelationIdentifier.getTraceIdKey()) == CorrelationIdentifier.getTraceId()
+    MDC.get(CorrelationIdentifier.getSpanIdKey()) == CorrelationIdentifier.getSpanId()
+    MDC.get("foo") == "bar"
+
+    when:
+    rootScope.close()
+
+    then:
+    MDC.get(CorrelationIdentifier.getTraceIdKey()) == null
+    MDC.get(CorrelationIdentifier.getSpanIdKey()) == null
+    MDC.get("foo") == "bar"
+  }
+
+  def "mdc context scoped by thread"() {
+    setup:
+    AtomicReference<String> thread1TraceId = new AtomicReference<>()
+    AtomicReference<String> thread2TraceId = new AtomicReference<>()
+
+    final Thread thread1 = new Thread() {
+      @Override
+      void run() {
+        // no trace in scope
+        thread1TraceId.set(MDC.get(CorrelationIdentifier.getTraceIdKey()))
+      }
+    }
+
+    final Thread thread2 = new Thread() {
+      @Override
+      void run() {
+        // other trace in scope
+        final Scope thread2Scope = GlobalTracer.get().buildSpan("root2").startActive(true)
+        try {
+          thread2TraceId.set(MDC.get(CorrelationIdentifier.getTraceIdKey()))
+        } finally {
+          thread2Scope.close()
+        }
+      }
+    }
+    final Scope mainScope = GlobalTracer.get().buildSpan("root").startActive(true)
+    thread1.start()
+    thread2.start()
+    final String mainThreadTraceId = MDC.get(CorrelationIdentifier.getTraceIdKey())
+    final String expectedMainThreadTraceId = CorrelationIdentifier.getTraceId()
+
+    thread1.join()
+    thread2.join()
+
+    expect:
+    mainThreadTraceId == expectedMainThreadTraceId
+    thread1TraceId.get() == null
+    thread2TraceId.get() != null
+    thread2TraceId.get() != mainThreadTraceId
+
+    cleanup:
+    mainScope?.close()
+  }
+}

dd-trace-api/src/main/java/datadog/trace/api/Config.java

@@ -59,6 +59,9 @@ public class Config {
   public static final String DEFAULT_AGENT_HOST = "localhost";
   public static final int DEFAULT_TRACE_AGENT_PORT = 8126;
 
+  public static final String LOGS_INJECTION_ENABLED = "logs.injection";
+  public static final boolean DEFAULT_LOGS_INJECTION_ENABLED = false;
+
   private static final boolean DEFAULT_RUNTIME_CONTEXT_FIELD_INJECTION = true;
 
   private static final boolean DEFAULT_PRIORITY_SAMPLING_ENABLED = true;
@@ -91,6 +94,7 @@ public class Config {
   @Getter private final Integer jmxFetchRefreshBeansPeriod;
   @Getter private final String jmxFetchStatsdHost;
   @Getter private final Integer jmxFetchStatsdPort;
+  @Getter private final boolean logsInjectionEnabled;
 
   // Read order: System Properties -> Env Variables, [-> default value]
   // Visible for testing
@@ -128,6 +132,9 @@ public class Config {
     jmxFetchStatsdHost = getSettingFromEnvironment(JMX_FETCH_STATSD_HOST, null);
     jmxFetchStatsdPort =
         getIntegerSettingFromEnvironment(JMX_FETCH_STATSD_PORT, DEFAULT_JMX_FETCH_STATSD_PORT);
+
+    logsInjectionEnabled =
+        getBooleanSettingFromEnvironment(LOGS_INJECTION_ENABLED, DEFAULT_LOGS_INJECTION_ENABLED);
   }
 
   // Read order: Properties -> Parent
@@ -169,6 +176,9 @@ public class Config {
     jmxFetchStatsdHost = properties.getProperty(JMX_FETCH_STATSD_HOST, parent.jmxFetchStatsdHost);
     jmxFetchStatsdPort =
         getPropertyIntegerValue(properties, JMX_FETCH_STATSD_PORT, parent.jmxFetchStatsdPort);
+
+    logsInjectionEnabled =
+        getBooleanSettingFromEnvironment(LOGS_INJECTION_ENABLED, DEFAULT_LOGS_INJECTION_ENABLED);
   }
 
   public Map<String, String> getMergedSpanTags() {

dd-trace-api/src/main/java/datadog/trace/api/CorrelationIdentifier.java

@@ -6,6 +6,19 @@ package datadog.trace.api;
  * <p>Intended for use with MDC frameworks.
  */
 public class CorrelationIdentifier {
+  private static final String TRACE_ID_KEY = "dd.trace_id";
+  private static final String SPAN_ID_KEY = "dd.span_id";
+
+  /** @return The trace-id key to use with datadog logs integration */
+  public static String getTraceIdKey() {
+    return TRACE_ID_KEY;
+  }
+
+  /** @return The span-id key to use with datadog logs integration */
+  public static String getSpanIdKey() {
+    return SPAN_ID_KEY;
+  }
+
   public static String getTraceId() {
     return GlobalTracer.get().getTraceId();
   }

settings.gradle

@@ -62,6 +62,7 @@ include ':dd-java-agent:instrumentation:rabbitmq-amqp-2.7'
 include ':dd-java-agent:instrumentation:ratpack-1.4'
 include ':dd-java-agent:instrumentation:servlet-2'
 include ':dd-java-agent:instrumentation:servlet-3'
+include ':dd-java-agent:instrumentation:slf4j-mdc'
 include ':dd-java-agent:instrumentation:sparkjava-2.3'
 include ':dd-java-agent:instrumentation:spring-web'
 include ':dd-java-agent:instrumentation:spring-webflux'