Share tooling loader across muzzle runs
MODIFY
buildSrc/src/main/groovy/MuzzlePlugin.groovy
buildSrc/src/main/groovy/MuzzlePlugin.groovy
@@ -22,6 +22,7 @@ import org.gradle.api.Task
import org.gradle.api.model.ObjectFactory
import java.lang.reflect.Method
+import java.util.concurrent.atomic.AtomicReference
/**
* muzzle task plugin which runs muzzle validation against a range of dependencies.
@@ -31,6 +32,7 @@ class MuzzlePlugin implements Plugin<Project> {
* Remote repositories used to query version ranges and fetch dependencies
*/
private static final List<RemoteRepository> MUZZLE_REPOS
+ private static final AtomicReference<ClassLoader> TOOLING_LOADER = new AtomicReference<>()
static {
RemoteRepository central = new RemoteRepository.Builder("central", "default", "http://central.maven.org/maven2/").build()
MUZZLE_REPOS = new ArrayList<RemoteRepository>(Arrays.asList(central))
@@ -52,10 +54,10 @@ class MuzzlePlugin implements Plugin<Project> {
if (!project.muzzle.directives.any { it.assertPass }) {
project.getLogger().info('No muzzle pass directives configured. Asserting pass against instrumentation compile-time dependencies')
final ClassLoader userCL = createCompileDepsClassLoader(project, bootstrapProject)
- final ClassLoader agentCL = createDDClassloader(project, toolingProject)
- Method assertionMethod = agentCL.loadClass('datadog.trace.agent.tooling.muzzle.MuzzleVersionScanPlugin')
- .getMethod('assertInstrumentationMuzzled', ClassLoader.class, boolean.class)
- assertionMethod.invoke(null, userCL, true)
+ final ClassLoader instrumentationCL = createInstrumentationClassloader(project, toolingProject)
+ Method assertionMethod = instrumentationCL.loadClass('datadog.trace.agent.tooling.muzzle.MuzzleVersionScanPlugin')
+ .getMethod('assertInstrumentationMuzzled', ClassLoader.class, ClassLoader.class, boolean.class)
+ assertionMethod.invoke(null, instrumentationCL, userCL, true)
}
}
}
@@ -63,10 +65,10 @@ class MuzzlePlugin implements Plugin<Project> {
group = 'Muzzle'
description = "Print references created by instrumentation muzzle"
doLast {
- final ClassLoader agentCL = createDDClassloader(project, toolingProject)
- Method assertionMethod = agentCL.loadClass('datadog.trace.agent.tooling.muzzle.MuzzleVersionScanPlugin')
- .getMethod('printMuzzleReferences')
- assertionMethod.invoke(null)
+ final ClassLoader instrumentationCL = createInstrumentationClassloader(project, toolingProject)
+ Method assertionMethod = instrumentationCL.loadClass('datadog.trace.agent.tooling.muzzle.MuzzleVersionScanPlugin')
+ .getMethod('printMuzzleReferences', ClassLoader.class)
+ assertionMethod.invoke(null, instrumentationCL)
}
}
project.tasks.compileMuzzle.dependsOn(bootstrapProject.tasks.compileJava)
@@ -91,6 +93,18 @@ class MuzzlePlugin implements Plugin<Project> {
return
}
+ if (!project.rootProject.tasks.getNames().contains('muzzleCleanup')) {
+ def muzzleCleanup = project.rootProject.task('muzzleCleanup') {
+ group = 'Muzzle'
+ doLast {
+ project.rootProject.getLogger().info("Cleaning up global tooling loader")
+ TOOLING_LOADER.set(null)
+ }
+ }
+ muzzleCleanup.outputs.upToDateWhen { false }
+ }
+ project.tasks.muzzle.finalizedBy(project.rootProject.tasks.muzzleCleanup)
+
final RepositorySystem system = newRepositorySystem()
final RepositorySystemSession session = newRepositorySystemSession(system)
@@ -115,22 +129,34 @@ class MuzzlePlugin implements Plugin<Project> {
}
}
- /**
- * Create a classloader with core agent classes and project instrumentation on the classpath.
- */
- private static ClassLoader createDDClassloader(Project project, Project toolingProject) {
- project.getLogger().info("Creating dd classpath for: " + project.getName())
+ private static ClassLoader getOrCreateToolingLoader(Project toolingProject) {
+ final ClassLoader toolingLoader = TOOLING_LOADER.get()
+ if (toolingLoader == null) {
Set<URL> ddUrls = new HashSet<>()
+ toolingProject.getLogger().info('creating classpath for agent-tooling')
for (File f : toolingProject.sourceSets.main.runtimeClasspath.getFiles()) {
- project.getLogger().info('--' + f)
+ toolingProject.getLogger().info('--' + f)
ddUrls.add(f.toURI().toURL())
}
+ TOOLING_LOADER.compareAndSet(null, new URLClassLoader(ddUrls.toArray(new URL[0]), (ClassLoader) null))
+ return TOOLING_LOADER.get()
+ } else {
+ return toolingLoader
+ }
+ }
+
+ /**
+ * Create a classloader with core agent classes and project instrumentation on the classpath.
+ */
+ private static ClassLoader createInstrumentationClassloader(Project project, Project toolingProject) {
+ project.getLogger().info("Creating instrumentation classpath for: " + project.getName())
+ Set<URL> ddUrls = new HashSet<>()
for (File f : project.sourceSets.main.runtimeClasspath.getFiles()) {
project.getLogger().info('--' + f)
ddUrls.add(f.toURI().toURL())
}
- return new URLClassLoader(ddUrls.toArray(new URL[0]), (ClassLoader) null)
+ return new URLClassLoader(ddUrls.toArray(new URL[0]), getOrCreateToolingLoader(toolingProject))
}
/**
@@ -248,11 +274,11 @@ class MuzzlePlugin implements Plugin<Project> {
def muzzleTask = instrumentationProject.task(taskName) {
doLast {
final ClassLoader userCL = createClassLoaderForTask(instrumentationProject, bootstrapProject, taskName)
- final ClassLoader agentCL = createDDClassloader(instrumentationProject, toolingProject)
+ final ClassLoader instrumentationCL = createInstrumentationClassloader(instrumentationProject, toolingProject)
// find all instrumenters, get muzzle, and assert
- Method assertionMethod = agentCL.loadClass('datadog.trace.agent.tooling.muzzle.MuzzleVersionScanPlugin')
- .getMethod('assertInstrumentationMuzzled', ClassLoader.class, boolean.class)
- assertionMethod.invoke(null, userCL, muzzleDirective.assertPass)
+ Method assertionMethod = instrumentationCL.loadClass('datadog.trace.agent.tooling.muzzle.MuzzleVersionScanPlugin')
+ .getMethod('assertInstrumentationMuzzled', ClassLoader.class, ClassLoader.class, boolean.class)
+ assertionMethod.invoke(null, instrumentationCL, userCL, muzzleDirective.assertPass)
for (Thread thread : Thread.getThreads()) {
if (thread.getName().startsWith("dd-")) {

MODIFY
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java
dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java
@@ -16,18 +16,21 @@ public interface WeakMap<K, V> {
void put(K key, V value);
+ @Slf4j
class Provider {
private static final AtomicReference<Supplier> provider =
new AtomicReference<>(Supplier.DEFAULT);
public static void registerIfAbsent(final Supplier provider) {
if (provider != null && provider != Supplier.DEFAULT) {
- Provider.provider.compareAndSet(Supplier.DEFAULT, provider);
+ if (Provider.provider.compareAndSet(Supplier.DEFAULT, provider)) {
+ log.debug("Weak map provider set to {}", provider);
+ }
}
}
public static boolean isProviderRegistered() {
- return Provider.provider.get() != null;
+ return provider.get() != Supplier.DEFAULT;
}
public static <K, V> WeakMap<K, V> newWeakMap() {

MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPlugin.java
@@ -3,11 +3,15 @@ package datadog.trace.agent.tooling.muzzle;
import datadog.trace.agent.tooling.HelperInjector;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.bootstrap.WeakMap;
+import java.io.IOException;
import java.lang.reflect.Method;
import java.util.Collections;
+import java.util.LinkedHashMap;
import java.util.List;
+import java.util.Map;
import java.util.ServiceLoader;
import java.util.WeakHashMap;
+import net.bytebuddy.dynamic.ClassFileLocator;
/**
* Entry point for muzzle version scan gradle plugin.
@@ -29,18 +33,20 @@ public class MuzzleVersionScanPlugin {
});
}
- public static void assertInstrumentationMuzzled(final ClassLoader cl, final boolean assertPass)
+ public static void assertInstrumentationMuzzled(
+ final ClassLoader instrumentationLoader,
+ final ClassLoader userClassLoader,
+ final boolean assertPass)
throws Exception {
// muzzle validate all instrumenters
for (Instrumenter instrumenter :
- ServiceLoader.load(Instrumenter.class, MuzzleGradlePlugin.class.getClassLoader())) {
+ ServiceLoader.load(Instrumenter.class, instrumentationLoader)) {
if (instrumenter.getClass().getName().endsWith("TraceConfigInstrumentation")) {
// TraceConfigInstrumentation doesn't do muzzle checks
// check on TracerClassInstrumentation instead
instrumenter =
(Instrumenter)
- MuzzleGradlePlugin.class
- .getClassLoader()
+ instrumentationLoader
.loadClass(instrumenter.getClass().getName() + "$TracerClassInstrumentation")
.getDeclaredConstructor()
.newInstance();
@@ -50,7 +56,8 @@ public class MuzzleVersionScanPlugin {
m = instrumenter.getClass().getDeclaredMethod("getInstrumentationMuzzle");
m.setAccessible(true);
final ReferenceMatcher muzzle = (ReferenceMatcher) m.invoke(instrumenter);
- final List<Reference.Mismatch> mismatches = muzzle.getMismatchedReferenceSources(cl);
+ final List<Reference.Mismatch> mismatches =
+ muzzle.getMismatchedReferenceSources(userClassLoader);
final boolean passed = mismatches.size() == 0;
if (mismatches.size() > 0) {}
if (passed && !assertPass) {
@@ -76,14 +83,13 @@ public class MuzzleVersionScanPlugin {
// run helper injector on all instrumenters
if (assertPass) {
for (Instrumenter instrumenter :
- ServiceLoader.load(Instrumenter.class, MuzzleGradlePlugin.class.getClassLoader())) {
+ ServiceLoader.load(Instrumenter.class, instrumentationLoader)) {
if (instrumenter.getClass().getName().endsWith("TraceConfigInstrumentation")) {
// TraceConfigInstrumentation doesn't do muzzle checks
// check on TracerClassInstrumentation instead
instrumenter =
(Instrumenter)
- MuzzleGradlePlugin.class
- .getClassLoader()
+ instrumentationLoader
.loadClass(instrumenter.getClass().getName() + "$TracerClassInstrumentation")
.getDeclaredConstructor()
.newInstance();
@@ -92,7 +98,8 @@ public class MuzzleVersionScanPlugin {
// verify helper injector works
final String[] helperClassNames = instrumenter.helperClassNames();
if (helperClassNames.length > 0) {
- new HelperInjector(helperClassNames).transform(null, null, cl, null);
+ new HelperInjector(createHelperMap(instrumenter))
+ .transform(null, null, userClassLoader, null);
}
} catch (final Exception e) {
System.err.println(
@@ -103,9 +110,21 @@ public class MuzzleVersionScanPlugin {
}
}
- public static void printMuzzleReferences() {
+ private static Map<String, byte[]> createHelperMap(Instrumenter instrumenter) throws IOException {
+ final Map<String, byte[]> helperMap =
+ new LinkedHashMap<>(instrumenter.helperClassNames().length);
+ for (final String helperName : instrumenter.helperClassNames()) {
+ final ClassFileLocator locator =
+ ClassFileLocator.ForClassLoader.of(instrumenter.getClass().getClassLoader());
+ final byte[] classBytes = locator.locate(helperName).resolve();
+ helperMap.put(helperName, classBytes);
+ }
+ return helperMap;
+ }
+
+ public static void printMuzzleReferences(final ClassLoader instrumentationLoader) {
for (final Instrumenter instrumenter :
- ServiceLoader.load(Instrumenter.class, MuzzleGradlePlugin.class.getClassLoader())) {
+ ServiceLoader.load(Instrumenter.class, instrumentationLoader)) {
if (instrumenter instanceof Instrumenter.Default) {
try {
final Method getMuzzleMethod =

Share tooling loader across muzzle runs