Skip to content

Commit 2a2d21a

Browse files
lpriimalpriima
committed
Add DD tags to log instrumentations
1 parent fdc5b57 commit 2a2d21a

File tree

15 files changed

+449
-28
lines changed

15 files changed

+449
-28
lines changed

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -204,6 +204,9 @@ private static void deleteTempDir(final File file) {
204204
final boolean deleted = file.delete();
205205
if (!deleted) {
206206
file.deleteOnExit();
207+
log.debug("file '{}' added to shutdown delete hook", file);
208+
} else {
209+
log.debug("file '{}' deleted", file);
207210
}
208211
}
209212
}

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/matcher/AdditionalLibraryIgnoresMatcher.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,10 @@ public boolean matches(final T target) {
203203
if (name.equals("ch.qos.logback.core.AsyncAppenderBase$Worker")) {
204204
return false;
205205
}
206+
// for inserting service, env, version in MDC of every thread
207+
if (name.equals("ch.qos.logback.classic.util.LogbackMDCAdapter")) {
208+
return false;
209+
}
206210

207211
return true;
208212
}

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,15 @@
22

33
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
44

5+
import datadog.trace.api.Config;
56
import datadog.trace.api.CorrelationIdentifier;
7+
import datadog.trace.bootstrap.instrumentation.api.Tags;
68
import datadog.trace.context.ScopeListener;
9+
import java.lang.reflect.InvocationTargetException;
710
import java.lang.reflect.Method;
11+
import java.util.Collections;
12+
import java.util.HashMap;
13+
import java.util.Map;
814
import lombok.extern.slf4j.Slf4j;
915

1016
/**
@@ -50,4 +56,34 @@ public void afterScopeClosed() {
5056
log.debug("Exception removing log context context", e);
5157
}
5258
}
59+
60+
static final Map<String, String> LOG_CONTEXT_DD_TAGS;
61+
62+
static {
63+
Map<String, String> logContextDDTags = new HashMap<>();
64+
logContextDDTags.put(Tags.DD_SERVICE, Config.get().getServiceName());
65+
final Map<String, String> mergedSpanTags = Config.get().getMergedSpanTags();
66+
String version = "";
67+
String env = "";
68+
if (mergedSpanTags != null) {
69+
version = mergedSpanTags.get("version");
70+
if (version == null) {
71+
version = "";
72+
}
73+
env = mergedSpanTags.get("env");
74+
if (env == null) {
75+
env = "";
76+
}
77+
}
78+
logContextDDTags.put(Tags.DD_VERSION, version);
79+
logContextDDTags.put(Tags.DD_ENV, env);
80+
LOG_CONTEXT_DD_TAGS = Collections.unmodifiableMap(logContextDDTags);
81+
}
82+
83+
public static void addDDTagsToMDC(final Method putMethod)
84+
throws InvocationTargetException, IllegalAccessException {
85+
for (final Map.Entry<String, String> e : LOG_CONTEXT_DD_TAGS.entrySet()) {
86+
putMethod.invoke(null, e.getKey(), e.getValue());
87+
}
88+
}
5389
}

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

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
package datadog.trace.agent.tooling.log;
2+
3+
import java.lang.reflect.InvocationTargetException;
4+
import java.lang.reflect.Method;
5+
import java.util.Map;
6+
import lombok.extern.slf4j.Slf4j;
7+
8+
/**
9+
* @param <T> something like String to String map. DD tags will be injected in T only if T is
10+
* instance of Map, or has T#put(String, Object) method, or putValue(String, Object) method.
11+
*/
12+
@Slf4j
13+
public class ThreadLocalWithDDTagsInitValue<T> extends ThreadLocal<T> {
14+
private static final String[] PUT_METHODS_POSSIBLE_NAMES_IN_ORDER =
15+
new String[] {"put", "putValue"};
16+
17+
private static Method findMethodWithName(final Method[] methods, final String methodName) {
18+
for (final Method method : methods) {
19+
if (method.getName().equals(methodName)) {
20+
return method;
21+
}
22+
}
23+
return null;
24+
}
25+
26+
/**
27+
* try to find method with name in {@param methods} in order of names listed in {@param
28+
* methodNames}
29+
*/
30+
private static Method findMethodWithNamesByOrder(
31+
final Method[] methods, final String[] methodNames) {
32+
for (final String methodName : methodNames) {
33+
final Method method = findMethodWithName(methods, methodName);
34+
if (method != null) {
35+
return method;
36+
}
37+
}
38+
return null;
39+
}
40+
41+
private static <T> T putToMap(final T stringToObjMap, final Map<String, ?> whatToPut)
42+
throws InvocationTargetException, IllegalAccessException {
43+
if (stringToObjMap instanceof Map) {
44+
((Map) stringToObjMap).putAll(whatToPut);
45+
} else {
46+
final Class<?> cl = stringToObjMap.getClass();
47+
final Method[] methods = cl.getDeclaredMethods();
48+
final Method method =
49+
findMethodWithNamesByOrder(methods, PUT_METHODS_POSSIBLE_NAMES_IN_ORDER);
50+
if (method != null) {
51+
final Class<?>[] parameterTypes = method.getParameterTypes();
52+
if (parameterTypes.length != 2 || !String.class.equals(parameterTypes[0])) {
53+
log.warn(
54+
"Can't find a way how to add DD tags to '{}'; was trying to use {}, "
55+
+ "but didn't like it's signature",
56+
stringToObjMap,
57+
method);
58+
return null;
59+
}
60+
method.setAccessible(true);
61+
for (final Map.Entry<String, ?> entry : whatToPut.entrySet()) {
62+
method.invoke(stringToObjMap, entry.getKey(), entry.getValue());
63+
}
64+
} else {
65+
log.warn("Can't find a method how to add DD tags to '{}'", stringToObjMap);
66+
return null;
67+
}
68+
}
69+
return stringToObjMap;
70+
}
71+
72+
public static <T> ThreadLocal<T> create(T origThreadLocalValue)
73+
throws InvocationTargetException, IllegalAccessException {
74+
if (origThreadLocalValue instanceof Map) {
75+
((Map) origThreadLocalValue).putAll(LogContextScopeListener.LOG_CONTEXT_DD_TAGS);
76+
} else {
77+
putToMap(origThreadLocalValue, LogContextScopeListener.LOG_CONTEXT_DD_TAGS);
78+
}
79+
return new ThreadLocalWithDDTagsInitValue<>(origThreadLocalValue);
80+
}
81+
82+
private final T initValue;
83+
84+
private ThreadLocalWithDDTagsInitValue() {
85+
log.warn("This constructor should never be called");
86+
initValue = null;
87+
}
88+
89+
private ThreadLocalWithDDTagsInitValue(final T initValue) {
90+
this.initValue = initValue;
91+
}
92+
93+
@Override
94+
protected T initialValue() {
95+
return initValue;
96+
}
97+
}

‎dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/ClassLoaderMatcherTest.groovy

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package datadog.trace.agent.test
22

33
import datadog.trace.agent.tooling.ClassLoaderMatcher
4+
import datadog.trace.agent.tooling.log.LogContextScopeListener
5+
import datadog.trace.agent.tooling.log.ThreadLocalWithDDTagsInitValue
46
import datadog.trace.bootstrap.DatadogClassLoader
57
import datadog.trace.util.test.DDSpecification
68

@@ -38,6 +40,12 @@ class ClassLoaderMatcherTest extends DDSpecification {
3840
DatadogClassLoader.name == "datadog.trace.bootstrap.DatadogClassLoader"
3941
}
4042

43+
def "helper class names are hardcoded in Log Instrumentations"() {
44+
expect:
45+
LogContextScopeListener.name == "datadog.trace.agent.tooling.log.LogContextScopeListener"
46+
ThreadLocalWithDDTagsInitValue.name == "datadog.trace.agent.tooling.log.ThreadLocalWithDDTagsInitValue"
47+
}
48+
4149
/*
4250
* A URLClassloader which only delegates java.* classes
4351
*/

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

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
11
package datadog.trace.instrumentation.log4j1;
22

33
import static java.util.Collections.singletonMap;
4-
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
4+
import static net.bytebuddy.matcher.ElementMatchers.isTypeInitializer;
55
import static net.bytebuddy.matcher.ElementMatchers.named;
66

77
import com.google.auto.service.AutoService;
88
import datadog.trace.agent.tooling.Instrumenter;
99
import datadog.trace.agent.tooling.log.LogContextScopeListener;
1010
import datadog.trace.api.Config;
1111
import datadog.trace.api.GlobalTracer;
12+
import java.lang.reflect.InvocationTargetException;
1213
import java.lang.reflect.Method;
1314
import java.util.Map;
1415
import net.bytebuddy.asm.Advice;
@@ -37,7 +38,7 @@ public ElementMatcher<? super TypeDescription> typeMatcher() {
3738
@Override
3839
public Map<? extends ElementMatcher<? super MethodDescription>, String> transformers() {
3940
return singletonMap(
40-
isConstructor(), Log4j1MDCInstrumentation.class.getName() + "$MDCContextAdvice");
41+
isTypeInitializer(), Log4j1MDCInstrumentation.class.getName() + "$MDCContextAdvice");
4142
}
4243

4344
@Override
@@ -47,18 +48,16 @@ public String[] helperClassNames() {
4748

4849
public static class MDCContextAdvice {
4950
@Advice.OnMethodExit(suppress = Throwable.class)
50-
public static void mdcClassInitialized(@Advice.This Object instance) {
51-
if (instance == null) {
52-
return;
53-
}
54-
51+
public static void mdcClassInitialized(@Advice.Origin final Class<?> mdcClass) {
5552
try {
56-
Class<?> mdcClass = instance.getClass();
5753
final Method putMethod = mdcClass.getMethod("put", String.class, Object.class);
5854
final Method removeMethod = mdcClass.getMethod("remove", String.class);
5955
GlobalTracer.get().addScopeListener(new LogContextScopeListener(putMethod, removeMethod));
60-
} catch (final NoSuchMethodException e) {
61-
org.slf4j.LoggerFactory.getLogger(instance.getClass())
56+
// log4j1 uses subclass of InheritableThreadLocal and we don't need to modify private thread
57+
// local field:
58+
LogContextScopeListener.addDDTagsToMDC(putMethod);
59+
} catch (final NoSuchMethodException | InvocationTargetException | IllegalAccessException e) {
60+
org.slf4j.LoggerFactory.getLogger(mdcClass)
6261
.debug("Failed to add log4j ThreadContext span listener", e);
6362
}
6463
}

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,4 +19,15 @@ class Log4j1MDCTest extends LogContextInjectionTestBase {
1919
def get(String key) {
2020
return MDC.get(key)
2121
}
22+
23+
@Override
24+
def remove(String key) {
25+
MDC.context
26+
return MDC.remove(key)
27+
}
28+
29+
@Override
30+
def clear() {
31+
return MDC.clear()
32+
}
2233
}

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

Lines changed: 77 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,13 @@
77
import com.google.auto.service.AutoService;
88
import datadog.trace.agent.tooling.Instrumenter;
99
import datadog.trace.agent.tooling.log.LogContextScopeListener;
10+
import datadog.trace.agent.tooling.log.ThreadLocalWithDDTagsInitValue;
1011
import datadog.trace.api.Config;
1112
import datadog.trace.api.GlobalTracer;
13+
import java.lang.reflect.Field;
14+
import java.lang.reflect.InvocationTargetException;
1215
import java.lang.reflect.Method;
16+
import java.util.HashMap;
1317
import java.util.Map;
1418
import net.bytebuddy.asm.Advice;
1519
import net.bytebuddy.description.method.MethodDescription;
@@ -18,7 +22,8 @@
1822

1923
@AutoService(Instrumenter.class)
2024
public class ThreadContextInstrumentation extends Instrumenter.Default {
21-
public static final String MDC_INSTRUMENTATION_NAME = "log4j";
25+
private static final String TYPE_NAME = "org.apache.logging.log4j.ThreadContext";
26+
private static final String MDC_INSTRUMENTATION_NAME = "log4j";
2227

2328
public ThreadContextInstrumentation() {
2429
super(MDC_INSTRUMENTATION_NAME);
@@ -31,7 +36,7 @@ protected boolean defaultEnabled() {
3136

3237
@Override
3338
public ElementMatcher<? super TypeDescription> typeMatcher() {
34-
return named("org.apache.logging.log4j.ThreadContext");
39+
return named(TYPE_NAME);
3540
}
3641

3742
@Override
@@ -42,19 +47,83 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
4247

4348
@Override
4449
public String[] helperClassNames() {
45-
return new String[] {"datadog.trace.agent.tooling.log.LogContextScopeListener"};
50+
return new String[] {
51+
"datadog.trace.agent.tooling.log.LogContextScopeListener",
52+
"datadog.trace.agent.tooling.log.ThreadLocalWithDDTagsInitValue",
53+
};
4654
}
4755

4856
public static class ThreadContextAdvice {
4957
@Advice.OnMethodExit(suppress = Throwable.class)
50-
public static void mdcClassInitialized(@Advice.Origin final Class threadClass) {
58+
public static void mdcClassInitialized(@Advice.Origin final Class<?> threadContextClass) {
5159
try {
52-
final Method putMethod = threadClass.getMethod("put", String.class, String.class);
53-
final Method removeMethod = threadClass.getMethod("remove", String.class);
60+
final Method putMethod = threadContextClass.getMethod("put", String.class, String.class);
61+
final Method removeMethod = threadContextClass.getMethod("remove", String.class);
5462
GlobalTracer.get().addScopeListener(new LogContextScopeListener(putMethod, removeMethod));
55-
} catch (final NoSuchMethodException e) {
56-
org.slf4j.LoggerFactory.getLogger(threadClass)
63+
64+
final Field contextMapField = threadContextClass.getDeclaredField("contextMap");
65+
contextMapField.setAccessible(true);
66+
final Object contextMap = contextMapField.get(null);
67+
if (contextMap
68+
.getClass()
69+
.getCanonicalName()
70+
.equals("org.apache.logging.slf4j.MDCContextMap")) {
71+
org.slf4j.LoggerFactory.getLogger(threadContextClass)
72+
.debug(
73+
"Log4j to SLF4J Adapter detected. "
74+
+ TYPE_NAME
75+
+ "'s ThreadLocal"
76+
+ " field will not be instrumented because it delegates to slf4-MDC");
77+
return;
78+
}
79+
final Field localMapField = contextMap.getClass().getDeclaredField("localMap");
80+
if (!ThreadLocal.class.isAssignableFrom(localMapField.getType())) {
81+
org.slf4j.LoggerFactory.getLogger(threadContextClass)
82+
.debug("Can't find thread local field: {}", localMapField);
83+
return;
84+
}
85+
86+
localMapField.setAccessible(true);
87+
final Object threadLocalInitValue = ((ThreadLocal) localMapField.get(contextMap)).get();
88+
final String fullTypeWithGeneric = localMapField.getGenericType().toString();
89+
if ("java.lang.ThreadLocal<java.util.Map<java.lang.String, java.lang.String>>"
90+
.equals(fullTypeWithGeneric)) {
91+
final Map<String, String> threadLocalInitValueAsMap =
92+
threadLocalInitValue != null
93+
? (Map<String, String>) threadLocalInitValue
94+
: new HashMap<String, String>();
95+
org.slf4j.LoggerFactory.getLogger(threadContextClass)
96+
.debug("Setting {} for ThreadLocalWithDDTagsInitValue ", threadLocalInitValueAsMap);
97+
localMapField.set(
98+
contextMap, ThreadLocalWithDDTagsInitValue.create(threadLocalInitValueAsMap));
99+
} else if ("java.lang.ThreadLocal<org.apache.logging.log4j.util.StringMap>"
100+
.equals(fullTypeWithGeneric)) {
101+
final Object threadLocalInitValueAsStringMap =
102+
threadLocalInitValue != null
103+
? threadLocalInitValue
104+
: Class.forName("org.apache.logging.log4j.util.SortedArrayStringMap")
105+
.newInstance();
106+
org.slf4j.LoggerFactory.getLogger(threadContextClass)
107+
.debug(
108+
"Setting {} for ThreadLocalWithDDTagsInitValue ",
109+
threadLocalInitValueAsStringMap);
110+
localMapField.set(
111+
contextMap, ThreadLocalWithDDTagsInitValue.create(threadLocalInitValueAsStringMap));
112+
} else {
113+
org.slf4j.LoggerFactory.getLogger(threadContextClass)
114+
.warn(
115+
"can't find thread local for {}; skipping adding extra tags",
116+
fullTypeWithGeneric);
117+
}
118+
} catch (final NoSuchMethodException
119+
| NoSuchFieldException
120+
| IllegalAccessException
121+
| InvocationTargetException e) {
122+
org.slf4j.LoggerFactory.getLogger(threadContextClass)
57123
.debug("Failed to add log4j ThreadContext span listener", e);
124+
} catch (final Throwable t) {
125+
org.slf4j.LoggerFactory.getLogger(threadContextClass)
126+
.debug("Unexpected exception while adding log4j ThreadContext span listener", t);
58127
}
59128
}
60129
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,4 +12,14 @@ class Log4jThreadContextTest extends LogContextInjectionTestBase {
1212
def get(String key) {
1313
return ThreadContext.get(key)
1414
}
15+
16+
@Override
17+
def remove(String key) {
18+
return ThreadContext.remove(key)
19+
}
20+
21+
@Override
22+
def clear() {
23+
return ThreadContext.clearAll()
24+
}
1525
}

0 commit comments

Comments
 (0)