slf4j MDC autoinjection integration
MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java
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."))

ADD
dd-java-agent/instrumentation/slf4j-mdc/slf4j-mdc.gradle
+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')
+}

ADD
dd-java-agent/instrumentation/slf4j-mdc/src/main/java/datadog/trace/instrumentation/slf4j/mdc/MDCInjectionInstrumentation.java
+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);
+ }
+ }
+ }
+ }
+}

ADD
dd-java-agent/instrumentation/slf4j-mdc/src/test/groovy/Slf4jMDCTest.groovy
+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()
+ }
+}

MODIFY
dd-trace-api/src/main/java/datadog/trace/api/Config.java
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() {

MODIFY
dd-trace-api/src/main/java/datadog/trace/api/CorrelationIdentifier.java
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();
}

MODIFY
settings.gradle
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'

slf4j MDC autoinjection integration