Create a reusable log context listener to be used for slf4j, log4j1 and log4j2
ADD
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/log/LogContextScopeListener.java
+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.
+ */
+@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);
+ }
+ }
+}

MODIFY
dd-java-agent/instrumentation/log4j1/src/main/java/datadog/trace/instrumentation/log4j1/Log4j1MDCInstrumentation.java
dd-java-agent/instrumentation/log4j1/src/main/java/datadog/trace/instrumentation/log4j1/Log4j1MDCInstrumentation.java
@@ -6,13 +6,11 @@ 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.CorrelationIdentifier;
import datadog.trace.api.GlobalTracer;
-import datadog.trace.context.ScopeListener;
import java.lang.reflect.Method;
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;
@@ -46,7 +44,7 @@ public class Log4j1MDCInstrumentation extends Instrumenter.Default {
@Override
public String[] helperClassNames() {
- return new String[] {MDCContextAdvice.class.getName() + "$MDCScopeListener"};
+ return new String[] {LogContextScopeListener.class.getName()};
}
public static class MDCContextAdvice {
@@ -60,43 +58,9 @@ public class Log4j1MDCInstrumentation extends Instrumenter.Default {
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 MDCScopeListener(putMethod, removeMethod));
+ GlobalTracer.get().addScopeListener(new LogContextScopeListener(putMethod, removeMethod));
} catch (final NoSuchMethodException e) {
}
}
-
- @Slf4j
- public static class MDCScopeListener implements ScopeListener {
- private final Method putMethod;
- private final Method removeMethod;
-
- public MDCScopeListener(final Method putMethod, final Method removeMethod) {
- System.out.println("Initializing scope listener");
- 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);
- }
- }
- }
}
}

MODIFY
dd-java-agent/instrumentation/log4j2/src/main/java/datadog/trace/instrumentation/log4j2/ThreadContextInstrumentation.java
dd-java-agent/instrumentation/log4j2/src/main/java/datadog/trace/instrumentation/log4j2/ThreadContextInstrumentation.java
@@ -6,14 +6,12 @@ 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.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;
@@ -60,7 +58,7 @@ public class ThreadContextInstrumentation extends Instrumenter.Default {
@Override
public String[] helperClassNames() {
- return new String[] {ThreadContextAdvice.class.getName() + "$ThreadContextScopeListener"};
+ return new String[] {LogContextScopeListener.class.getName()};
}
public static class ThreadContextAdvice {
@@ -69,45 +67,11 @@ public class ThreadContextInstrumentation extends Instrumenter.Default {
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));
+ 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);
}
}
-
- @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);
- }
- }
- }
}
}

MODIFY
dd-java-agent/instrumentation/slf4j-mdc/src/main/java/datadog/trace/instrumentation/slf4j/mdc/MDCInjectionInstrumentation.java
dd-java-agent/instrumentation/slf4j-mdc/src/main/java/datadog/trace/instrumentation/slf4j/mdc/MDCInjectionInstrumentation.java
@@ -6,14 +6,12 @@ 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.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 class MDCInjectionInstrumentation extends Instrumenter.Default {
@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 class MDCInjectionInstrumentation extends Instrumenter.Default {
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);
- }
- }
- }
}
}

Create a reusable log context listener to be used for slf4j, log4j1 and log4j2