Skip to content

Commit 36646ce

Browse files
bantonssonbantonsson
committed
Make FieldBackedProvider respect class loader matchers
1 parent 2cc3d3e commit 36646ce

41 files changed

Lines changed: 323 additions & 109 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

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

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,8 @@ private static boolean loadsExpectedClass(
116116
private static class ClassLoaderHasClassesNamedMatcher
117117
extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
118118

119-
private final WeakCache<ClassLoader, Boolean> cache = AgentTooling.newWeakCache(25);
119+
// Initialize this lazily because of startup ordering and muzzle plugin usage patterns
120+
private volatile WeakCache<ClassLoader, Boolean> cacheHolder = null;
120121

121122
private final String[] resources;
122123

@@ -127,6 +128,17 @@ private ClassLoaderHasClassesNamedMatcher(final String... classNames) {
127128
}
128129
}
129130

131+
private WeakCache<ClassLoader, Boolean> getCache() {
132+
if (cacheHolder == null) {
133+
synchronized (this) {
134+
if (cacheHolder == null) {
135+
cacheHolder = AgentTooling.newWeakCache(25);
136+
}
137+
}
138+
}
139+
return cacheHolder;
140+
}
141+
130142
private boolean hasResources(final ClassLoader cl) {
131143
for (final String resource : resources) {
132144
if (cl.getResource(resource) == null) {
@@ -143,6 +155,7 @@ public boolean matches(final ClassLoader cl) {
143155
return false;
144156
}
145157
final Boolean cached;
158+
WeakCache<ClassLoader, Boolean> cache = getCache();
146159
if ((cached = cache.getIfPresent(cl)) != null) {
147160
return cached;
148161
}

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

Lines changed: 57 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package datadog.trace.agent.tooling;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers.failSafe;
4+
import static java.util.Collections.singletonMap;
45
import static net.bytebuddy.matcher.ElementMatchers.any;
56
import static net.bytebuddy.matcher.ElementMatchers.isAnnotatedWith;
67
import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -17,6 +18,7 @@
1718
import java.security.ProtectionDomain;
1819
import java.util.Arrays;
1920
import java.util.Collections;
21+
import java.util.HashMap;
2022
import java.util.List;
2123
import java.util.Map;
2224
import java.util.SortedSet;
@@ -47,6 +49,7 @@ public interface Instrumenter {
4749

4850
@Slf4j
4951
abstract class Default implements Instrumenter {
52+
private static final ElementMatcher<ClassLoader> ANY_CLASS_LOADER = any();
5053

5154
// Added here instead of AgentInstaller's ignores because it's relatively
5255
// expensive. https://github.com/DataDog/dd-trace-java/pull/1045
@@ -55,7 +58,8 @@ abstract class Default implements Instrumenter {
5558

5659
private final SortedSet<String> instrumentationNames;
5760
private final String instrumentationPrimaryName;
58-
private final InstrumentationContextProvider contextProvider;
61+
private InstrumentationContextProvider contextProvider;
62+
private boolean initialized;
5963
protected final boolean enabled;
6064

6165
protected final String packageName =
@@ -67,11 +71,35 @@ public Default(final String instrumentationName, final String... additionalNames
6771
instrumentationPrimaryName = instrumentationName;
6872

6973
enabled = Config.get().isIntegrationEnabled(instrumentationNames, defaultEnabled());
70-
Map<String, String> contextStore = contextStore();
71-
if (!contextStore.isEmpty()) {
72-
contextProvider = new FieldBackedProvider(this, contextStore);
73-
} else {
74-
contextProvider = NoopContextProvider.INSTANCE;
74+
}
75+
76+
// Since the super(...) call is first in the constructor, we can't really rely on things
77+
// being properly initialized in the Instrumentation until the super(...) call has finished
78+
// so do the rest of the initialization lazily
79+
private void lazyInit() {
80+
synchronized (this) {
81+
if (!initialized) {
82+
Map<ElementMatcher<ClassLoader>, Map<String, String>> contextStores;
83+
Map<String, String> allClassLoaderContextStores = contextStoreForAll();
84+
Map<String, String> matchedContextStores = contextStore();
85+
if (!allClassLoaderContextStores.isEmpty()) {
86+
if (contextStore().isEmpty()) {
87+
contextStores = singletonMap(ANY_CLASS_LOADER, allClassLoaderContextStores);
88+
} else {
89+
contextStores = new HashMap<>();
90+
contextStores.put(classLoaderMatcher(), matchedContextStores);
91+
contextStores.put(ANY_CLASS_LOADER, allClassLoaderContextStores);
92+
}
93+
} else {
94+
contextStores = singletonMap(classLoaderMatcher(), matchedContextStores);
95+
}
96+
if (!contextStores.isEmpty()) {
97+
contextProvider = new FieldBackedProvider(this, contextStores);
98+
} else {
99+
contextProvider = NoopContextProvider.INSTANCE;
100+
}
101+
initialized = true;
102+
}
75103
}
76104
}
77105

@@ -82,6 +110,8 @@ public final AgentBuilder instrument(final AgentBuilder parentAgentBuilder) {
82110
return parentAgentBuilder;
83111
}
84112

113+
lazyInit();
114+
85115
AgentBuilder.Identified.Extendable agentBuilder =
86116
parentAgentBuilder
87117
.type(
@@ -200,9 +230,17 @@ public String[] helperClassNames() {
200230
return new String[0];
201231
}
202232

203-
/** @return A type matcher used to match the classloader under transform */
233+
/**
234+
* A type matcher used to match the classloader under transform.
235+
*
236+
* <p>This matcher needs to either implement equality checks or be the same for different
237+
* instrumentations that share context stores to avoid enabling the context store
238+
* instrumentations multiple times.
239+
*
240+
* @return A type matcher used to match the classloader under transform.
241+
*/
204242
public ElementMatcher<ClassLoader> classLoaderMatcher() {
205-
return any();
243+
return ANY_CLASS_LOADER;
206244
}
207245

208246
/** @return A type matcher used to match the class under transform. */
@@ -231,7 +269,7 @@ public void postMatch(
231269
public abstract Map<? extends ElementMatcher<? super MethodDescription>, String> transformers();
232270

233271
/**
234-
* Context stores to define for this instrumentation.
272+
* Context stores to define for this instrumentation. Are added to matching class loaders.
235273
*
236274
* <p>A map of {class-name -> context-class-name}. Keys (and their subclasses) will be
237275
* associated with a context of the value.
@@ -240,6 +278,16 @@ public Map<String, String> contextStore() {
240278
return Collections.emptyMap();
241279
}
242280

281+
/**
282+
* Context stores to define for this instrumentation. Are added to all class loaders.
283+
*
284+
* <p>A map of {class-name -> context-class-name}. Keys (and their subclasses) will be
285+
* associated with a context of the value.
286+
*/
287+
public Map<String, String> contextStoreForAll() {
288+
return Collections.emptyMap();
289+
}
290+
243291
protected boolean defaultEnabled() {
244292
return Config.get().isIntegrationsEnabled();
245293
}

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

Lines changed: 111 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
import net.bytebuddy.jar.asm.MethodVisitor;
4444
import net.bytebuddy.jar.asm.Opcodes;
4545
import net.bytebuddy.jar.asm.Type;
46+
import net.bytebuddy.matcher.ElementMatcher;
4647
import net.bytebuddy.pool.TypePool;
4748
import net.bytebuddy.utility.JavaModule;
4849

@@ -100,6 +101,7 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
100101
private final Instrumenter.Default instrumenter;
101102
private final ByteBuddy byteBuddy;
102103
private final Map<String, String> contextStore;
104+
private final Map<ElementMatcher<ClassLoader>, Map<String, String>> matchedContextStores;
103105

104106
/** fields-accessor-interface-name -> fields-accessor-interface-dynamic-type */
105107
private final Map<String, DynamicType.Unloaded<?>> fieldAccessorInterfaces;
@@ -114,9 +116,21 @@ public class FieldBackedProvider implements InstrumentationContextProvider {
114116
private final boolean fieldInjectionEnabled;
115117

116118
public FieldBackedProvider(
117-
final Instrumenter.Default instrumenter, Map<String, String> contextStore) {
119+
final Instrumenter.Default instrumenter,
120+
Map<ElementMatcher<ClassLoader>, Map<String, String>> matchedContextStores) {
118121
this.instrumenter = instrumenter;
119-
this.contextStore = contextStore;
122+
this.matchedContextStores = matchedContextStores;
123+
if (matchedContextStores.isEmpty()) {
124+
this.contextStore = Collections.emptyMap();
125+
} else if (matchedContextStores.size() == 1) {
126+
this.contextStore = matchedContextStores.entrySet().iterator().next().getValue();
127+
} else {
128+
this.contextStore = new HashMap<>();
129+
for (Map.Entry<ElementMatcher<ClassLoader>, Map<String, String>> matcherAndStores :
130+
matchedContextStores.entrySet()) {
131+
this.contextStore.putAll(matcherAndStores.getValue());
132+
}
133+
}
120134
byteBuddy = new ByteBuddy();
121135
fieldAccessorInterfaces = generateFieldAccessorInterfaces();
122136
fieldAccessorInterfacesInjector = bootstrapHelperInjector(fieldAccessorInterfaces.values());
@@ -349,16 +363,20 @@ public DynamicType.Builder<?> transform(
349363
}
350364

351365
/*
352-
Set of pairs (context holder, context class) for which we have matchers installed.
353-
We use this to make sure we do not install matchers repeatedly for cases when same
354-
context class is used by multiple instrumentations.
366+
* HashMap from the instrumentations contextClassLoaderMatcher to a set of pairs (context holder, context class)
367+
* for which we have matchers installed. We use this to make sure we do not install matchers repeatedly for cases
368+
* when same context class is used by multiple instrumentations.
355369
*/
356-
private static final Set<Map.Entry<String, String>> INSTALLED_CONTEXT_MATCHERS = new HashSet<>();
370+
private static final HashMap<ElementMatcher<ClassLoader>, Set<Map.Entry<String, String>>>
371+
INSTALLED_CONTEXT_MATCHERS = new HashMap<>();
372+
373+
private static final Set<Map.Entry<String, String>> INSTALLED_HELPERS = new HashSet<>();
357374

358375
/** Clear set that prevents multiple matchers for same context class */
359376
public static void resetContextMatchers() {
360377
synchronized (INSTALLED_CONTEXT_MATCHERS) {
361378
INSTALLED_CONTEXT_MATCHERS.clear();
379+
INSTALLED_HELPERS.clear();
362380
}
363381
}
364382

@@ -367,56 +385,81 @@ public AgentBuilder.Identified.Extendable additionalInstrumentation(
367385
AgentBuilder.Identified.Extendable builder) {
368386

369387
if (fieldInjectionEnabled) {
370-
for (final Map.Entry<String, String> entry : contextStore.entrySet()) {
371-
/*
372-
* For each context store defined in a current instrumentation we create an agent builder
373-
* that injects necessary fields.
374-
* Note: this synchronization should not have any impact on performance
375-
* since this is done when agent builder is being made, it doesn't affect actual
376-
* class transformation.
377-
*/
378-
synchronized (INSTALLED_CONTEXT_MATCHERS) {
379-
// FIXME: This makes an assumption that class loader matchers for instrumenters that use
380-
// same context classes should be the same - which seems reasonable, but is not checked.
381-
// Addressing this properly requires some notion of 'compound intrumenters' which we
382-
// currently do not have.
383-
if (INSTALLED_CONTEXT_MATCHERS.contains(entry)) {
384-
log.debug("Skipping builder for {} {}", instrumenter.getClass().getName(), entry);
385-
continue;
386-
}
387-
388-
log.debug("Making builder for {} {}", instrumenter.getClass().getName(), entry);
389-
INSTALLED_CONTEXT_MATCHERS.add(entry);
390-
388+
for (Map.Entry<ElementMatcher<ClassLoader>, Map<String, String>> matcherAndStores :
389+
matchedContextStores.entrySet()) {
390+
ElementMatcher<ClassLoader> classLoaderMatcher = matcherAndStores.getKey();
391+
for (Map.Entry<String, String> entry : matcherAndStores.getValue().entrySet()) {
391392
/*
392393
* For each context store defined in a current instrumentation we create an agent builder
393394
* that injects necessary fields.
395+
* Note: this synchronization should not have any impact on performance
396+
* since this is done when agent builder is being made, it doesn't affect actual
397+
* class transformation.
394398
*/
395-
builder =
396-
builder
397-
.type(safeHasSuperType(named(entry.getKey())), instrumenter.classLoaderMatcher())
398-
.and(safeToInjectFieldsMatcher())
399-
.and(Default.NOT_DECORATOR_MATCHER)
400-
.transform(NoOpTransformer.INSTANCE);
399+
Set<Map.Entry<String, String>> installedContextMatchers;
400+
synchronized (INSTALLED_CONTEXT_MATCHERS) {
401+
installedContextMatchers = INSTALLED_CONTEXT_MATCHERS.get(classLoaderMatcher);
402+
if (installedContextMatchers == null) {
403+
installedContextMatchers = new HashSet<>();
404+
INSTALLED_CONTEXT_MATCHERS.put(classLoaderMatcher, installedContextMatchers);
405+
}
406+
}
407+
synchronized (installedContextMatchers) {
408+
if (installedContextMatchers.contains(entry)) {
409+
if (log.isDebugEnabled()) {
410+
log.debug(
411+
"Skipping builder for {} with matcher {}: {} -> {}",
412+
instrumenter.getClass().getName(),
413+
classLoaderMatcher,
414+
entry.getKey(),
415+
entry.getValue());
416+
}
417+
continue;
418+
}
401419

402-
/*
403-
* We inject helpers here as well as when instrumentation is applied to ensure that
404-
* helpers are present even if instrumented classes are not loaded, but classes with state
405-
* fields added are loaded (e.g. sun.net.www.protocol.https.HttpsURLConnectionImpl).
406-
*/
407-
builder = injectHelpersIntoBootstrapClassloader(builder);
420+
if (log.isDebugEnabled()) {
421+
log.debug(
422+
"Making builder for {} with matcher {}: {} -> {}",
423+
instrumenter.getClass().getName(),
424+
classLoaderMatcher,
425+
entry.getKey(),
426+
entry.getValue());
427+
}
428+
installedContextMatchers.add(entry);
429+
430+
/*
431+
* For each context store defined in a current instrumentation we create an agent builder
432+
* that injects necessary fields.
433+
*/
434+
builder =
435+
builder
436+
.type(safeHasSuperType(named(entry.getKey())), classLoaderMatcher)
437+
.and(safeToInjectFieldsMatcher(entry.getKey(), entry.getValue()))
438+
.and(Default.NOT_DECORATOR_MATCHER)
439+
.transform(
440+
getTransformerForASMVisitor(
441+
getFieldInjectionVisitor(entry.getKey(), entry.getValue())));
442+
}
408443

409-
builder =
410-
builder.transform(
411-
getTransformerForASMVisitor(
412-
getFieldInjectionVisitor(entry.getKey(), entry.getValue())));
444+
synchronized (INSTALLED_HELPERS) {
445+
if (!INSTALLED_HELPERS.contains(entry)) {
446+
/*
447+
* We inject helpers here as well as when instrumentation is applied to ensure that
448+
* helpers are present even if instrumented classes are not loaded, but classes with state
449+
* fields added are loaded (e.g. sun.net.www.protocol.https.HttpsURLConnectionImpl).
450+
*/
451+
builder = injectHelpersIntoBootstrapClassloader(builder);
452+
INSTALLED_HELPERS.add(entry);
453+
}
454+
}
413455
}
414456
}
415457
}
416458
return builder;
417459
}
418460

419-
private static AgentBuilder.RawMatcher safeToInjectFieldsMatcher() {
461+
private static AgentBuilder.RawMatcher safeToInjectFieldsMatcher(
462+
final String keyType, final String valueType) {
420463
return new AgentBuilder.RawMatcher() {
421464
@Override
422465
public boolean matches(
@@ -433,9 +476,30 @@ public boolean matches(
433476
* parents. It looks like current JVM implementation does exactly this but javadoc is not
434477
* explicit about that.
435478
*/
436-
return classBeingRedefined == null
437-
|| Arrays.asList(classBeingRedefined.getInterfaces())
438-
.contains(FieldBackedContextStoreAppliedMarker.class);
479+
boolean result =
480+
classBeingRedefined == null
481+
|| Arrays.asList(classBeingRedefined.getInterfaces())
482+
.contains(FieldBackedContextStoreAppliedMarker.class);
483+
if (log.isDebugEnabled()) {
484+
if (result) {
485+
// Only log success the first time we add it to the class
486+
if (classBeingRedefined == null) {
487+
log.debug(
488+
"Added context-store field to {}: {} -> {}",
489+
typeDescription.getName(),
490+
keyType,
491+
valueType);
492+
}
493+
} else {
494+
// This will log for every failed redefine
495+
log.debug(
496+
"Failed to add context-store field to {}: {} -> {}",
497+
typeDescription.getName(),
498+
keyType,
499+
valueType);
500+
}
501+
}
502+
return result;
439503
}
440504
};
441505
}

‎dd-java-agent/instrumentation/akka-concurrent/src/main/java/datadog/trace/instrumentation/akka/concurrent/AkkaForkJoinPoolInstrumentation.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,11 @@ public AkkaForkJoinPoolInstrumentation() {
3131
super("java_concurrent", "akka_concurrent");
3232
}
3333

34+
@Override
35+
public ElementMatcher<ClassLoader> classLoaderMatcher() {
36+
return AkkaForkJoinTaskInstrumentation.CLASS_LOADER_MATCHER;
37+
}
38+
3439
@Override
3540
public ElementMatcher<TypeDescription> typeMatcher() {
3641
// This might need to be an extendsClass matcher...

0 commit comments

Comments
 (0)