Add basic support for log4j ThreadContext for log injection.

dd-java-agent/instrumentation/log4j-thread-context/log4j-thread-context.gradle

@@ -0,0 +1,28 @@
+apply from: "${rootDir}/gradle/java.gradle"
+
+ext {
+  log4jVersion = '2.11.2'
+}
+
+//muzzle {
+//  pass {
+//    group = 'org.slf4j'
+//    module = 'slf4j-api'
+//    versions = '(,)'
+//  }
+//}
+
+dependencies {
+  compile project(':dd-trace-api')
+  compile project(':dd-java-agent:agent-tooling')
+
+  testCompile group: 'org.apache.logging.log4j', name: 'log4j-core', version: log4jVersion
+  testCompile group: 'org.apache.logging.log4j', name: 'log4j-api', version: log4jVersion
+
+  compile deps.bytebuddy
+  compile deps.opentracing
+  annotationProcessor deps.autoservice
+  implementation deps.autoservice
+
+  testCompile project(':dd-java-agent:testing')
+}

dd-java-agent/instrumentation/log4j-thread-context/src/main/java/datadog/trace/instrumentation/slf4j/mdc/ThreadContextInstrumentation.java

@@ -0,0 +1,111 @@
+package datadog.trace.instrumentation.slf4j.mdc;
+
+import static java.util.Collections.singletonMap;
+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.Map;
+import lombok.extern.slf4j.Slf4j;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.utility.JavaModule;
+
+@AutoService(Instrumenter.class)
+public class ThreadContextInstrumentation extends Instrumenter.Default {
+  public static final String MDC_INSTRUMENTATION_NAME = "log4j-thread-context";
+
+  private static final String mdcClassName = "org.apache.logging.log4j.ThreadContext";
+
+  public ThreadContextInstrumentation() {
+    super(MDC_INSTRUMENTATION_NAME);
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return Config.getBooleanSettingFromEnvironment(
+        Config.LOGS_INJECTION_ENABLED, Config.DEFAULT_LOGS_INJECTION_ENABLED);
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return named(mdcClassName);
+  }
+
+  @Override
+  public void postMatch(
+      final TypeDescription typeDescription,
+      final ClassLoader classLoader,
+      final JavaModule module,
+      final Class<?> classBeingRedefined,
+      final ProtectionDomain protectionDomain) {
+    if (classBeingRedefined != null) {
+      ThreadContextAdvice.mdcClassInitialized(classBeingRedefined);
+    }
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(isTypeInitializer(), ThreadContextAdvice.class.getName());
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {ThreadContextAdvice.class.getName() + "$ThreadContextScopeListener"};
+  }
+
+  public static class ThreadContextAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void mdcClassInitialized(@Advice.Origin final Class threadClass) {
+      try {
+        final Method putMethod = threadClass.getMethod("put", String.class, String.class);
+        final Method removeMethod = threadClass.getMethod("remove", String.class);
+        GlobalTracer.get().addScopeListener(new ThreadContextScopeListener(putMethod, removeMethod));
+      } catch (final NoSuchMethodException e) {
+        org.slf4j.LoggerFactory.getLogger(threadClass).debug("Failed to add log4j ThreadContext span listener", e);
+      }
+    }
+
+    @Slf4j
+    public static class ThreadContextScopeListener implements ScopeListener {
+      private final Method putMethod;
+      private final Method removeMethod;
+
+      public ThreadContextScopeListener(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 (final Exception e) {
+          log.debug("Exception setting thread context context", e);
+        }
+      }
+
+      @Override
+      public void afterScopeClosed() {
+        try {
+          removeMethod.invoke(null, CorrelationIdentifier.getTraceIdKey());
+          removeMethod.invoke(null, CorrelationIdentifier.getSpanIdKey());
+        } catch (final Exception e) {
+          log.debug("Exception removing thread context context", e);
+        }
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/log4j-thread-context/src/test/groovy/Log4jThreadContextTest.groovy

@@ -0,0 +1,97 @@
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.agent.test.utils.ConfigUtils
+import datadog.trace.api.CorrelationIdentifier
+import io.opentracing.Scope
+import io.opentracing.util.GlobalTracer
+import org.apache.logging.log4j.ThreadContext
+
+import java.util.concurrent.atomic.AtomicReference
+
+class Log4jThreadContextTest extends AgentTestRunner {
+
+  static {
+      System.setProperty("dd.logs.injection", "true")
+  }
+
+  def "ThreadContext shows trace and span ids for active scope"() {
+    when:
+    ThreadContext.put("foo", "bar")
+    Scope rootScope = GlobalTracer.get().buildSpan("root").startActive(true)
+
+    then:
+    ThreadContext.get(CorrelationIdentifier.getTraceIdKey()) == CorrelationIdentifier.getTraceId()
+    ThreadContext.get(CorrelationIdentifier.getSpanIdKey()) == CorrelationIdentifier.getSpanId()
+    ThreadContext.get("foo") == "bar"
+
+    when:
+    Scope childScope = GlobalTracer.get().buildSpan("child").startActive(true)
+
+    then:
+    ThreadContext.get(CorrelationIdentifier.getTraceIdKey()) == CorrelationIdentifier.getTraceId()
+    ThreadContext.get(CorrelationIdentifier.getSpanIdKey()) == CorrelationIdentifier.getSpanId()
+    ThreadContext.get("foo") == "bar"
+
+    when:
+    childScope.close()
+
+    then:
+    ThreadContext.get(CorrelationIdentifier.getTraceIdKey()) == CorrelationIdentifier.getTraceId()
+    ThreadContext.get(CorrelationIdentifier.getSpanIdKey()) == CorrelationIdentifier.getSpanId()
+    ThreadContext.get("foo") == "bar"
+
+    when:
+    rootScope.close()
+
+    then:
+    ThreadContext.get(CorrelationIdentifier.getTraceIdKey()) == null
+    ThreadContext.get(CorrelationIdentifier.getSpanIdKey()) == null
+    ThreadContext.get("foo") == "bar"
+  }
+
+  def "ThreadContext context scoped by thread"() {
+    setup:
+    ConfigUtils.updateConfig {
+      System.setProperty("dd.logs.injection", "true")
+    }
+    AtomicReference<String> thread1TraceId = new AtomicReference<>()
+    AtomicReference<String> thread2TraceId = new AtomicReference<>()
+
+    final Thread thread1 = new Thread() {
+      @Override
+      void run() {
+        // no trace in scope
+        thread1TraceId.set(ThreadContext.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(ThreadContext.get(CorrelationIdentifier.getTraceIdKey()))
+        } finally {
+          thread2Scope.close()
+        }
+      }
+    }
+    final Scope mainScope = GlobalTracer.get().buildSpan("root").startActive(true)
+    thread1.start()
+    thread2.start()
+    final String mainThreadTraceId = ThreadContext.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()
+  }
+}

settings.gradle

@@ -66,6 +66,7 @@ include ':dd-java-agent:instrumentation:jsp-2.3'
 include ':dd-java-agent:instrumentation:kafka-clients-0.11'
 include ':dd-java-agent:instrumentation:kafka-streams-0.11'
 include ':dd-java-agent:instrumentation:lettuce-5'
+include ':dd-java-agent:instrumentation:log4j-thread-context'
 include ':dd-java-agent:instrumentation:mongo'
 include ':dd-java-agent:instrumentation:mongo:driver-3.1'
 include ':dd-java-agent:instrumentation:mongo:driver-async-3.3'