Support log trace injection for log4j 1.x and log4j 2.x when used without Slf4j

buildSrc/src/main/groovy/MuzzlePlugin.groovy

@@ -263,9 +263,15 @@ class MuzzlePlugin implements Plugin<Project> {
   private static Task addMuzzleTask(MuzzleDirective muzzleDirective, Artifact versionArtifact, Project instrumentationProject, Task runAfter, Project bootstrapProject, Project toolingProject) {
     def taskName = "muzzle-Assert${muzzleDirective.assertPass ? "Pass" : "Fail"}-$versionArtifact.groupId-$versionArtifact.artifactId-$versionArtifact.version${muzzleDirective.name ? "-${muzzleDirective.getNameSlug()}" : ""}"
     def config = instrumentationProject.configurations.create(taskName)
-    config.dependencies.add(instrumentationProject.dependencies.create("$versionArtifact.groupId:$versionArtifact.artifactId:$versionArtifact.version") {
+    def dep =  instrumentationProject.dependencies.create("$versionArtifact.groupId:$versionArtifact.artifactId:$versionArtifact.version") {
       transitive = true
-    })
+    }
+    // The following optional transitive dependencies are brought in by some legacy module such as log4j 1.x but are no
+    // longer bundled with the JVM and have to be excluded for the muzzle tests to be able to run.
+    dep.exclude group: 'com.sun.jdmk', module: 'jmxtools'
+    dep.exclude group: 'com.sun.jmx', module: 'jmxri'
+
+    config.dependencies.add(dep)
     for (String additionalDependency : muzzleDirective.additionalDependencies) {
       config.dependencies.add(instrumentationProject.dependencies.create(additionalDependency) {
         transitive = true

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

@@ -127,7 +127,7 @@ public static ResettableClassFileTransformer installBytebuddyAgent(
             .or(nameStartsWith("com.singularity."))
             .or(nameStartsWith("com.jinspired."))
             .or(nameStartsWith("org.jinspired."))
-            .or(nameStartsWith("org.apache.log4j."))
+            .or(nameStartsWith("org.apache.log4j.").and(not(named("org.apache.log4j.MDC"))))
             .or(nameStartsWith("org.slf4j.").and(not(named("org.slf4j.MDC"))))
             .or(nameContains("$JaxbAccessor"))
             .or(nameContains("CGLIB$$"))

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/log/LogContextScopeListener.java

@@ -0,0 +1,47 @@
+package datadog.trace.agent.tooling.log;
+
+import datadog.trace.api.CorrelationIdentifier;
+import datadog.trace.context.ScopeListener;
+import java.lang.reflect.Method;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * A scope listener that receives the MDC/ThreadContext put and receive methods and update the trace
+ * and span reference anytime a new scope is activated or closed.
+ */
+@Slf4j
+public class LogContextScopeListener implements ScopeListener {
+
+  /** A reference to the log context method that sets a new attribute in the log context */
+  private final Method putMethod;
+
+  /** A reference to the log context method that removes an attribute from the log context */
+  private final Method removeMethod;
+
+  public LogContextScopeListener(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 log 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 log context context", e);
+    }
+  }
+}

dd-java-agent/instrumentation/log4j1/log4j1.gradle

@@ -0,0 +1,37 @@
+apply from: "${rootDir}/gradle/java.gradle"
+
+ext {
+  log4jVersion = '1.2.17'
+}
+
+muzzle {
+  pass {
+    group = 'log4j'
+    module = 'log4j'
+    versions = '(,)'
+  }
+}
+
+configurations {
+  // In order to test the real log4j library we need to remove the log4j transitive
+  // dependency 'log4j-over-slf4j' brought in by :dd-java-agent:testing which would shadow
+  // the log4j module under test using a proxy to slf4j instead.
+  testCompile.exclude group: 'org.slf4j', module: 'log4j-over-slf4j'
+
+  // See: https://stackoverflow.com/a/9047963/2749853
+  testCompile.exclude group: 'javax.jms', module: 'jms'
+}
+
+dependencies {
+  compile project(':dd-trace-api')
+  compile project(':dd-java-agent:agent-tooling')
+
+  testCompile group: 'log4j', name: 'log4j', version: log4jVersion
+
+  compile deps.bytebuddy
+  compile deps.opentracing
+  annotationProcessor deps.autoservice
+  implementation deps.autoservice
+
+  testCompile project(':dd-java-agent:testing')
+}

dd-java-agent/instrumentation/log4j1/src/main/java/datadog/trace/instrumentation/log4j1/Log4j1MDCInstrumentation.java

@@ -0,0 +1,65 @@
+package datadog.trace.instrumentation.log4j1;
+
+import static java.util.Collections.singletonMap;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.log.LogContextScopeListener;
+import datadog.trace.api.Config;
+import datadog.trace.api.GlobalTracer;
+import java.lang.reflect.Method;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public class Log4j1MDCInstrumentation extends Instrumenter.Default {
+  public static final String MDC_INSTRUMENTATION_NAME = "log4j1";
+
+  public Log4j1MDCInstrumentation() {
+    super(MDC_INSTRUMENTATION_NAME);
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return Config.get().isLogsInjectionEnabled();
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return named("org.apache.log4j.MDC");
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(isConstructor(), MDCContextAdvice.class.getName());
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {LogContextScopeListener.class.getName()};
+  }
+
+  public static class MDCContextAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void mdcClassInitialized(@Advice.This Object instance) {
+      if (instance == null) {
+        return;
+      }
+
+      try {
+        Class<?> mdcClass = instance.getClass();
+        final Method putMethod = mdcClass.getMethod("put", String.class, Object.class);
+        final Method removeMethod = mdcClass.getMethod("remove", String.class);
+        GlobalTracer.get().addScopeListener(new LogContextScopeListener(putMethod, removeMethod));
+      } catch (final NoSuchMethodException e) {
+        org.slf4j.LoggerFactory.getLogger(instance.getClass())
+            .debug("Failed to add log4j ThreadContext span listener", e);
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/log4j1/src/test/groovy/Log4j1MDCTest.groovy

@@ -0,0 +1,15 @@
+import datadog.trace.agent.test.log.injection.LogContextInjectionTestBase
+import org.apache.log4j.MDC
+
+class Log4j1MDCTest extends LogContextInjectionTestBase {
+
+  @Override
+  def put(String key, Object value) {
+    return MDC.put(key, value)
+  }
+
+  @Override
+  def get(String key) {
+    return MDC.get(key)
+  }
+}

dd-java-agent/instrumentation/log4j2/log4j2.gradle

@@ -0,0 +1,40 @@
+apply from: "${rootDir}/gradle/java.gradle"
+
+ext {
+  log4jVersion = '2.11.2'
+}
+
+configurations {
+  // In order to test the real log4j library we need to remove the log4j transitive
+  // dependency 'log4j-over-slf4j' brought in by :dd-java-agent:testing which would shadow
+  // the log4j module under test using a proxy to slf4j instead.
+  testCompile.exclude group: 'org.slf4j', module: 'log4j-over-slf4j'
+}
+
+muzzle {
+  pass {
+    group = 'org.apache.logging.log4j'
+    module = 'log4j-core'
+    versions = '(,)'
+  }
+
+  pass {
+    group = 'org.apache.logging.log4j'
+    module = 'log4j-api'
+    versions = '(,)'
+  }
+}
+
+dependencies {
+  compile project(':dd-trace-api')
+  compile project(':dd-java-agent:agent-tooling')
+
+  compile deps.bytebuddy
+  compile deps.opentracing
+  annotationProcessor deps.autoservice
+  implementation deps.autoservice
+
+  testCompile project(':dd-java-agent:testing')
+  testCompile group: 'org.apache.logging.log4j', name: 'log4j-core', version: log4jVersion
+  testCompile group: 'org.apache.logging.log4j', name: 'log4j-api', version: log4jVersion
+}

dd-java-agent/instrumentation/log4j2/src/main/java/datadog/trace/instrumentation/log4j2/ThreadContextInstrumentation.java

@@ -0,0 +1,60 @@
+package datadog.trace.instrumentation.log4j2;
+
+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.agent.tooling.log.LogContextScopeListener;
+import datadog.trace.api.Config;
+import datadog.trace.api.GlobalTracer;
+import java.lang.reflect.Method;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.method.MethodDescription;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(Instrumenter.class)
+public class ThreadContextInstrumentation extends Instrumenter.Default {
+  public static final String MDC_INSTRUMENTATION_NAME = "log4j";
+
+  public ThreadContextInstrumentation() {
+    super(MDC_INSTRUMENTATION_NAME);
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return Config.get().isLogsInjectionEnabled();
+  }
+
+  @Override
+  public ElementMatcher<? super TypeDescription> typeMatcher() {
+    return named("org.apache.logging.log4j.ThreadContext");
+  }
+
+  @Override
+  public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
+    return singletonMap(isTypeInitializer(), ThreadContextAdvice.class.getName());
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {LogContextScopeListener.class.getName()};
+  }
+
+  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 LogContextScopeListener(putMethod, removeMethod));
+      } catch (final NoSuchMethodException e) {
+        org.slf4j.LoggerFactory.getLogger(threadClass)
+            .debug("Failed to add log4j ThreadContext span listener", e);
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/log4j2/src/test/groovy/Log4jThreadContextTest.groovy

@@ -0,0 +1,15 @@
+import datadog.trace.agent.test.log.injection.LogContextInjectionTestBase
+import org.apache.logging.log4j.ThreadContext
+
+class Log4jThreadContextTest extends LogContextInjectionTestBase {
+
+  @Override
+  def put(String key, Object value) {
+    return ThreadContext.put(key, value as String)
+  }
+
+  @Override
+  def get(String key) {
+    return ThreadContext.get(key)
+  }
+}

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

@@ -6,14 +6,12 @@
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.log.LogContextScopeListener;
 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;
@@ -61,7 +59,7 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
 
   @Override
   public String[] helperClassNames() {
-    return new String[] {MDCAdvice.class.getName() + "$MDCScopeListener"};
+    return new String[] {LogContextScopeListener.class.getName()};
   }
 
   public static class MDCAdvice {
@@ -70,43 +68,10 @@ 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));
+        GlobalTracer.get().addScopeListener(new LogContextScopeListener(putMethod, removeMethod));
       } catch (final 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 (final Exception e) {
-          log.debug("Exception setting mdc 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 mdc context", e);
-        }
-      }
-    }
   }
 }