Skip to content

Commit b8bee9e

Browse files
lpriimalpriima
committed
replace ThreadLocal field in logges which not inherrit context
1 parent 439d308 commit b8bee9e

File tree

7 files changed

+145
-44
lines changed

7 files changed

+145
-44
lines changed

‎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
@@ -195,6 +195,10 @@ public boolean matches(final T target) {
195195
if (name.equals("ch.qos.logback.core.AsyncAppenderBase$Worker")) {
196196
return false;
197197
}
198+
// for inserting service, env, version in MDC of every thread
199+
if (name.equals("ch.qos.logback.classic.util.LogbackMDCAdapter")) {
200+
return false;
201+
}
198202

199203
return true;
200204
}

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

Lines changed: 27 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@
88
import datadog.trace.context.ScopeListener;
99
import java.lang.reflect.InvocationTargetException;
1010
import java.lang.reflect.Method;
11+
import java.util.Collections;
12+
import java.util.HashMap;
1113
import java.util.Map;
1214
import lombok.extern.slf4j.Slf4j;
1315

@@ -55,25 +57,33 @@ public void afterScopeClosed() {
5557
}
5658
}
5759

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+
5883
public static void addDDTagsToMDC(final Method putMethod)
5984
throws InvocationTargetException, IllegalAccessException {
60-
putMethod.invoke(null, Tags.DD_SERVICE, Config.get().getServiceName());
61-
{
62-
final Map<String, String> mergedSpanTags = Config.get().getMergedSpanTags();
63-
String version = "";
64-
String env = "";
65-
if (mergedSpanTags != null) {
66-
version = mergedSpanTags.get("version");
67-
if (version == null) {
68-
version = "";
69-
}
70-
env = mergedSpanTags.get("env");
71-
if (env == null) {
72-
env = "";
73-
}
74-
}
75-
putMethod.invoke(null, Tags.DD_VERSION, version);
76-
putMethod.invoke(null, Tags.DD_ENV, env);
85+
for (final Map.Entry<String, String> e : LOG_CONTEXT_DD_TAGS.entrySet()) {
86+
putMethod.invoke(null, e.getKey(), e.getValue());
7787
}
7888
}
7989
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
package datadog.trace.agent.tooling.log;
2+
3+
import java.util.Map;
4+
5+
public class ThreadLocalWithDDTagsInitValue extends ThreadLocal<Map<String, String>> {
6+
@Override
7+
protected Map<String, String> initialValue() {
8+
return LogContextScopeListener.LOG_CONTEXT_DD_TAGS;
9+
}
10+
}

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,8 +53,10 @@ public static void mdcClassInitialized(@Advice.Origin final Class<?> mdcClass) {
5353
final Method putMethod = mdcClass.getMethod("put", String.class, Object.class);
5454
final Method removeMethod = mdcClass.getMethod("remove", String.class);
5555
GlobalTracer.get().addScopeListener(new LogContextScopeListener(putMethod, removeMethod));
56+
// log4j1 uses subclass of InheritableThreadLocal and we don't need to modify private thread
57+
// local field:
5658
LogContextScopeListener.addDDTagsToMDC(putMethod);
57-
} catch (final NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
59+
} catch (final NoSuchMethodException | InvocationTargetException | IllegalAccessException e) {
5860
org.slf4j.LoggerFactory.getLogger(mdcClass)
5961
.debug("Failed to add log4j ThreadContext span listener", e);
6062
}

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

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,10 @@
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;
12-
import java.lang.reflect.InvocationTargetException;
13+
import java.lang.reflect.Field;
1314
import java.lang.reflect.Method;
1415
import java.util.Map;
1516
import net.bytebuddy.asm.Advice;
@@ -43,19 +44,30 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
4344

4445
@Override
4546
public String[] helperClassNames() {
46-
return new String[] {LogContextScopeListener.class.getName()};
47+
return new String[] {
48+
LogContextScopeListener.class.getName(), ThreadLocalWithDDTagsInitValue.class.getName(),
49+
};
4750
}
4851

4952
public static class ThreadContextAdvice {
5053
@Advice.OnMethodExit(suppress = Throwable.class)
51-
public static void mdcClassInitialized(@Advice.Origin final Class<?> threadClass) {
54+
public static void mdcClassInitialized(@Advice.Origin final Class<?> threadContextClass) {
55+
// @Advice.FieldValue("contextMap") Object contextMap
5256
try {
53-
final Method putMethod = threadClass.getMethod("put", String.class, String.class);
54-
final Method removeMethod = threadClass.getMethod("remove", String.class);
57+
final Method putMethod = threadContextClass.getMethod("put", String.class, String.class);
58+
final Method removeMethod = threadContextClass.getMethod("remove", String.class);
5559
GlobalTracer.get().addScopeListener(new LogContextScopeListener(putMethod, removeMethod));
56-
LogContextScopeListener.addDDTagsToMDC(putMethod);
57-
} catch (final NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
58-
org.slf4j.LoggerFactory.getLogger(threadClass)
60+
61+
final Field contextMapField = threadContextClass.getDeclaredField("contextMap");
62+
contextMapField.setAccessible(true);
63+
Object contextMap = contextMapField.get(null);
64+
65+
final Field localMapField = contextMap.getClass().getDeclaredField("localMap");
66+
localMapField.setAccessible(true);
67+
localMapField.set(contextMap, new ThreadLocalWithDDTagsInitValue());
68+
69+
} catch (final NoSuchMethodException | NoSuchFieldException | IllegalAccessException e) {
70+
org.slf4j.LoggerFactory.getLogger(threadContextClass)
5971
.debug("Failed to add log4j ThreadContext span listener", e);
6072
}
6173
}

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

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,10 @@
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;
12-
import java.lang.reflect.InvocationTargetException;
13+
import java.lang.reflect.Field;
1314
import java.lang.reflect.Method;
1415
import java.security.ProtectionDomain;
1516
import java.util.Map;
@@ -19,6 +20,7 @@
1920
import net.bytebuddy.matcher.BooleanMatcher;
2021
import net.bytebuddy.matcher.ElementMatcher;
2122
import net.bytebuddy.utility.JavaModule;
23+
import org.slf4j.spi.MDCAdapter;
2224

2325
@AutoService(Instrumenter.class)
2426
public class MDCInjectionInstrumentation extends Instrumenter.Default {
@@ -71,7 +73,9 @@ public boolean matches(final MethodDescription target) {
7173

7274
@Override
7375
public String[] helperClassNames() {
74-
return new String[] {LogContextScopeListener.class.getName()};
76+
return new String[] {
77+
LogContextScopeListener.class.getName(), ThreadLocalWithDDTagsInitValue.class.getName(),
78+
};
7579
}
7680

7781
public static class MDCAdvice {
@@ -81,8 +85,16 @@ public static void mdcClassInitialized(@Advice.Origin final Class<?> mdcClass) {
8185
final Method putMethod = mdcClass.getMethod("put", String.class, String.class);
8286
final Method removeMethod = mdcClass.getMethod("remove", String.class);
8387
GlobalTracer.get().addScopeListener(new LogContextScopeListener(putMethod, removeMethod));
84-
LogContextScopeListener.addDDTagsToMDC(putMethod);
85-
} catch (final NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
88+
89+
final Field mdcAdapterField = mdcClass.getDeclaredField("mdcAdapter");
90+
mdcAdapterField.setAccessible(true);
91+
final MDCAdapter mdcAdapterInstance = (MDCAdapter) mdcAdapterField.get(null);
92+
final Field copyOnThreadLocalField =
93+
mdcAdapterInstance.getClass().getDeclaredField("copyOnThreadLocal");
94+
copyOnThreadLocalField.setAccessible(true);
95+
copyOnThreadLocalField.set(mdcAdapterInstance, new ThreadLocalWithDDTagsInitValue());
96+
97+
} catch (final NoSuchMethodException | IllegalAccessException | NoSuchFieldException e) {
8698
org.slf4j.LoggerFactory.getLogger(mdcClass).debug("Failed to add MDC span listener", e);
8799
}
88100
}

‎dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/log/injection/LogContextInjectionTestBase.groovy

Lines changed: 65 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@ import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan
2020
* satisfy in order to support log injection.
2121
*/
2222
abstract class LogContextInjectionTestBase extends AgentTestRunner {
23+
private static final TEST_ENV = "test"
24+
private static final TEST_VERSION = "0.1"
2325

2426
/**
2527
* Set in the framework-specific context the given value at the given key
@@ -34,7 +36,7 @@ abstract class LogContextInjectionTestBase extends AgentTestRunner {
3436
static {
3537
ConfigUtils.updateConfig {
3638
System.setProperty("dd.logs.injection", "true")
37-
System.setProperty(PREFIX + GLOBAL_TAGS, "env:test,version:0.1")
39+
System.setProperty(PREFIX + GLOBAL_TAGS, "env:${TEST_ENV},version:${TEST_VERSION}")
3840
}
3941
}
4042

@@ -49,8 +51,8 @@ abstract class LogContextInjectionTestBase extends AgentTestRunner {
4951
get(CorrelationIdentifier.getSpanIdKey()) == CorrelationIdentifier.getSpanId()
5052
get("foo") == "bar"
5153
get(Tags.DD_SERVICE) == Config.DEFAULT_SERVICE_NAME
52-
get(Tags.DD_VERSION) == "0.1"
53-
get(Tags.DD_ENV) == "test"
54+
get(Tags.DD_VERSION) == TEST_VERSION
55+
get(Tags.DD_ENV) == TEST_ENV
5456

5557
when:
5658
AgentSpan childSpan = startSpan("child")
@@ -61,8 +63,8 @@ abstract class LogContextInjectionTestBase extends AgentTestRunner {
6163
get(CorrelationIdentifier.getSpanIdKey()) == CorrelationIdentifier.getSpanId()
6264
get("foo") == "bar"
6365
get(Tags.DD_SERVICE) == Config.DEFAULT_SERVICE_NAME
64-
get(Tags.DD_VERSION) == "0.1"
65-
get(Tags.DD_ENV) == "test"
66+
get(Tags.DD_VERSION) == TEST_VERSION
67+
get(Tags.DD_ENV) == TEST_ENV
6668

6769
when:
6870
childScope.close()
@@ -73,8 +75,8 @@ abstract class LogContextInjectionTestBase extends AgentTestRunner {
7375
get(CorrelationIdentifier.getSpanIdKey()) == CorrelationIdentifier.getSpanId()
7476
get("foo") == "bar"
7577
get(Tags.DD_SERVICE) == Config.DEFAULT_SERVICE_NAME
76-
get(Tags.DD_VERSION) == "0.1"
77-
get(Tags.DD_ENV) == "test"
78+
get(Tags.DD_VERSION) == TEST_VERSION
79+
get(Tags.DD_ENV) == TEST_ENV
7880

7981
when:
8082
rootScope.close()
@@ -85,25 +87,39 @@ abstract class LogContextInjectionTestBase extends AgentTestRunner {
8587
get(CorrelationIdentifier.getSpanIdKey()) == null
8688
get("foo") == "bar"
8789
get(Tags.DD_SERVICE) == Config.DEFAULT_SERVICE_NAME
88-
get(Tags.DD_VERSION) == "0.1"
89-
get(Tags.DD_ENV) == "test"
90+
get(Tags.DD_VERSION) == TEST_VERSION
91+
get(Tags.DD_ENV) == TEST_ENV
9092
}
9193

9294
def "Log context is scoped by thread"() {
9395
AtomicReference<String> thread1TraceId = new AtomicReference<>()
9496
AtomicReference<String> thread2TraceId = new AtomicReference<>()
97+
def t1VersionBeg
98+
def t1VersionEnd
99+
def t2VersionBeg
100+
def t2VersionEnd
101+
def t1EnvBeg
102+
def t1EnvEnd
103+
def t2EnvBeg
104+
def t2EnvEnd
95105

96106
final Thread thread1 = new Thread() {
97107
@Override
98108
void run() {
109+
t1VersionBeg = get(Tags.DD_VERSION)
110+
t1EnvBeg = get(Tags.DD_ENV)
99111
// no trace in scope
100112
thread1TraceId.set(get(CorrelationIdentifier.getTraceIdKey()))
113+
t1VersionEnd = get(Tags.DD_VERSION)
114+
t1EnvEnd = get(Tags.DD_ENV)
101115
}
102116
}
103117

104118
final Thread thread2 = new Thread() {
105119
@Override
106120
void run() {
121+
t2VersionBeg = get(Tags.DD_VERSION)
122+
t2EnvBeg = get(Tags.DD_ENV)
107123
// other trace in scope
108124
final AgentSpan thread2Span = startSpan("root2")
109125
final AgentScope thread2Scope = activateSpan(thread2Span)
@@ -113,6 +129,8 @@ abstract class LogContextInjectionTestBase extends AgentTestRunner {
113129
thread2Scope.close()
114130
thread2Span.finish()
115131
}
132+
t2VersionEnd = get(Tags.DD_VERSION)
133+
t2EnvEnd = get(Tags.DD_ENV)
116134
}
117135
}
118136
final AgentSpan mainSpan = startSpan("root")
@@ -131,19 +149,52 @@ abstract class LogContextInjectionTestBase extends AgentTestRunner {
131149
thread2TraceId.get() != null
132150
thread2TraceId.get() != mainThreadTraceId
133151
get(Tags.DD_SERVICE) == Config.DEFAULT_SERVICE_NAME
134-
get(Tags.DD_VERSION) == "0.1"
135-
get(Tags.DD_ENV) == "test"
152+
get(Tags.DD_VERSION) == TEST_VERSION
153+
get(Tags.DD_ENV) == TEST_ENV
154+
t1VersionBeg == TEST_VERSION
155+
t1VersionEnd == TEST_VERSION
156+
t2VersionBeg == TEST_VERSION
157+
t2VersionEnd == TEST_VERSION
158+
t1EnvBeg == TEST_ENV
159+
t1EnvEnd == TEST_ENV
160+
t2EnvBeg == TEST_ENV
161+
t2EnvEnd == TEST_ENV
136162

137163
cleanup:
138164
mainScope?.close()
139165
mainSpan?.finish()
140166
}
141167

142168
def "always log service, version, env"() {
169+
def mainThreadVersion = get(Tags.DD_VERSION)
170+
def t1VersionBeg
171+
def t1VersionEnd
172+
def t1EnvBeg
173+
def t1EnvEnd
174+
175+
final Thread thread1 = new Thread() {
176+
@Override
177+
void run() {
178+
t1VersionBeg = get(Tags.DD_VERSION)
179+
t1EnvBeg = get(Tags.DD_ENV)
180+
181+
println("something: " + this)
182+
183+
t1VersionEnd = get(Tags.DD_VERSION)
184+
t1EnvEnd = get(Tags.DD_ENV)
185+
}
186+
}
187+
thread1.start()
188+
thread1.join()
189+
143190
expect:
191+
mainThreadVersion == get(Tags.DD_VERSION)
144192
get(Tags.DD_SERVICE) == Config.DEFAULT_SERVICE_NAME
145-
get(Tags.DD_VERSION) == "0.1"
146-
get(Tags.DD_ENV) == "test"
147-
193+
get(Tags.DD_VERSION) == TEST_VERSION
194+
get(Tags.DD_ENV) == TEST_ENV
195+
t1VersionBeg == TEST_VERSION
196+
t1VersionEnd == TEST_VERSION
197+
t1EnvBeg == TEST_ENV
198+
t1EnvEnd == TEST_ENV
148199
}
149200
}

0 commit comments

Comments
 (0)