Run tests in a custom classloader with bootstrap classpath
MODIFY
dd-java-agent/testing/src/main/java/datadog/trace/agent/test/AgentTestRunner.java
dd-java-agent/testing/src/main/java/datadog/trace/agent/test/AgentTestRunner.java
package datadog.trace.agent.test;
-import static datadog.trace.agent.tooling.Utils.AGENT_PACKAGE_PREFIXES;
-
import ch.qos.logback.classic.Level;
import ch.qos.logback.classic.Logger;
import datadog.opentracing.DDSpan;
@@ -10,18 +8,14 @@ import datadog.trace.agent.tooling.AgentInstaller;
import datadog.trace.agent.tooling.Instrumenter;
import datadog.trace.common.writer.ListWriter;
import io.opentracing.Tracer;
-import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
-import java.lang.reflect.Field;
-import java.util.List;
+import java.util.*;
import java.util.concurrent.Phaser;
import java.util.concurrent.atomic.AtomicInteger;
-import lombok.extern.slf4j.Slf4j;
import net.bytebuddy.agent.ByteBuddyAgent;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.description.type.TypeDescription;
-import net.bytebuddy.dynamic.ClassFileLocator;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.utility.JavaModule;
import org.junit.After;
@@ -29,10 +23,7 @@ import org.junit.AfterClass;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.runner.RunWith;
-import org.junit.runner.notification.RunNotifier;
-import org.junit.runners.model.InitializationError;
import org.slf4j.LoggerFactory;
-import org.spockframework.runtime.Sputnik;
import org.spockframework.runtime.model.SpecMetadata;
import spock.lang.Specification;
@@ -51,8 +42,7 @@ import spock.lang.Specification;
* in an initialized state.
* </ul>
*/
-@Slf4j
-@RunWith(AgentTestRunner.SpockRunner.class)
+@RunWith(SpockRunner.class)
@SpecMetadata(filename = "AgentTestRunner.java", line = 0)
public abstract class AgentTestRunner extends Specification {
/**
@@ -66,11 +56,13 @@ public abstract class AgentTestRunner extends Specification {
private static final AtomicInteger INSTRUMENTATION_ERROR_COUNT = new AtomicInteger();
private static final Instrumentation instrumentation;
- private static ClassFileTransformer activeTransformer = null;
+ private static volatile ClassFileTransformer activeTransformer = null;
protected static final Phaser WRITER_PHASER = new Phaser();
static {
+ instrumentation = ByteBuddyAgent.getInstrumentation();
+
((Logger) LoggerFactory.getLogger(Logger.ROOT_LOGGER_NAME)).setLevel(Level.WARN);
((Logger) LoggerFactory.getLogger("datadog")).setLevel(Level.DEBUG);
@@ -85,27 +77,24 @@ public abstract class AgentTestRunner extends Specification {
}
};
TEST_TRACER = new DDTracer(TEST_WRITER);
-
- ByteBuddyAgent.install();
- instrumentation = ByteBuddyAgent.getInstrumentation();
+ TestUtils.registerOrReplaceGlobalTracer(TEST_TRACER);
}
@BeforeClass
- public static synchronized void agentSetup() {
+ public static synchronized void agentSetup() throws Exception {
if (null != activeTransformer) {
throw new IllegalStateException("transformer already in place: " + activeTransformer);
}
activeTransformer =
AgentInstaller.installBytebuddyAgent(instrumentation, new ErrorCountingListener());
- TestUtils.registerOrReplaceGlobalTracer(TEST_TRACER);
}
@Before
public void beforeTest() {
TEST_WRITER.start();
INSTRUMENTATION_ERROR_COUNT.set(0);
- assert TEST_TRACER.activeSpan() == null;
+ assert (TEST_TRACER).activeSpan() == null;
}
@After
@@ -115,11 +104,13 @@ public abstract class AgentTestRunner extends Specification {
@AfterClass
public static synchronized void agentClenup() {
+ if (null != activeTransformer) {
instrumentation.removeTransformer(activeTransformer);
activeTransformer = null;
}
+ }
- private static class ErrorCountingListener implements AgentBuilder.Listener {
+ public static class ErrorCountingListener implements AgentBuilder.Listener {
@Override
public void onDiscovery(
final String typeName,
@@ -160,77 +151,94 @@ public abstract class AgentTestRunner extends Specification {
final boolean loaded) {}
}
- public static class SpockRunner extends Sputnik {
- private final InstrumentationClassLoader customLoader;
+ // TODO: remove in a separate commit
+ // Notes for running agent on isolated classloader.
+ /*
+ final File bootstrapJar = createBootstrapJar();
+ final File agentJar = createAgentJar();
- public SpockRunner(Class<?> clazz)
- throws InitializationError, NoSuchFieldException, SecurityException,
- IllegalArgumentException, IllegalAccessException {
- super(shadowTestClass(clazz));
- // access the classloader created in shadowTestClass above
- Field clazzField = Sputnik.class.getDeclaredField("clazz");
- try {
- clazzField.setAccessible(true);
- customLoader =
- (InstrumentationClassLoader) ((Class<?>) clazzField.get(this)).getClassLoader();
- } finally {
- clazzField.setAccessible(false);
- }
- }
-
- // Shadow the test class with bytes loaded by InstrumentationClassLoader
- private static Class<?> shadowTestClass(final Class<?> clazz) {
- try {
- InstrumentationClassLoader customLoader =
- new InstrumentationClassLoader(SpockRunner.class.getClassLoader());
- return customLoader.shadow(clazz);
- } catch (IOException e) {
- throw new IllegalStateException(e);
- }
- }
+ instrumentation.appendToBootstrapClassLoaderSearch(new JarFile(bootstrapJar));
- // Replace the context class loader for each test with InstrumentationClassLoader
- @Override
- public void run(final RunNotifier notifier) {
+ final ClassLoader agentClassLoader = createDatadogClassLoader(bootstrapJar, agentJar);
final ClassLoader contextLoader = Thread.currentThread().getContextClassLoader();
try {
- Thread.currentThread().setContextClassLoader(customLoader);
- super.run(notifier);
+ Thread.currentThread().setContextClassLoader(agentClassLoader);
+
+ { // install agent
+ final Class<?> agentInstallerClass =
+ agentClassLoader.loadClass("datadog.trace.agent.tooling.AgentInstaller");
+ final Class<?> listenerArrayClass =
+ Array.newInstance(agentClassLoader.loadClass("net.bytebuddy.agent.builder.AgentBuilder$Listener"), 0).getClass();
+ final Class<?> errorListenerClass =
+ agentClassLoader.loadClass(AgentTestRunner.class.getName()+ "$ErrorCountingListener");
+ final Method agentInstallerMethod =
+ agentInstallerClass.getMethod("installBytebuddyAgent", Instrumentation.class, listenerArrayClass);
+ final Object listeners = Array.newInstance(errorListenerClass, 1);
+ Array.set(listeners, 0, errorListenerClass.newInstance());
+ activeTransformer = (ClassFileTransformer) agentInstallerMethod.invoke(null, instrumentation, listeners);
+ }
} finally {
Thread.currentThread().setContextClassLoader(contextLoader);
}
- }
- }
-
- /** TODO: Doc */
- private static class InstrumentationClassLoader extends java.lang.ClassLoader {
- final ClassLoader parent;
-
- public InstrumentationClassLoader(ClassLoader parent) {
- super(parent);
- this.parent = parent;
- }
-
- /** Forcefully inject the bytes of clazz into this classloader. */
- public Class<?> shadow(Class<?> clazz) throws IOException {
- final ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(clazz.getClassLoader());
- final byte[] classBytes = locator.locate(clazz.getName()).resolve();
-
- Class<?> shadowed = this.defineClass(clazz.getName(), classBytes, 0, classBytes.length);
- return shadowed;
- }
+ */
- @Override
- protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
- if (!name.startsWith("datadog.trace.agent.test.")) {
- for (int i = 0; i < AGENT_PACKAGE_PREFIXES.length; ++i) {
- if (name.startsWith(AGENT_PACKAGE_PREFIXES[i])) {
- throw new ClassNotFoundException(
- "refusing to load agent class" + name + " on test classloader.");
- }
- }
- }
- return parent.loadClass(name);
- }
+ /*
+ private static File createAgentJar() throws IOException {
+ final ClassLoader loader = AgentTestRunner.class.getClassLoader();
+ final ClassPath testCP = ClassPath.from(loader);
+ Set<String> agentClasses = new HashSet<String>();
+ for (ClassPath.ClassInfo info : testCP.getAllClasses()) {
+ boolean isAgentClass = true;
+ for (int i = 0; i < TEST_BOOTSTRAP.length; ++i) {
+ if (info.getName().startsWith(TEST_BOOTSTRAP[i])) {
+ isAgentClass = false;
+ break;
+ }
+ }
+ if (info.getName().startsWith("org.junit")
+ || info.getName().startsWith("junit")
+ || info.getName().startsWith("org.mockito")
+ || info.getName().startsWith("org.assertj")
+ || info.getName().startsWith("org.omg")
+ || info.getName().startsWith("lombok")
+ || info.getName().startsWith("org.spockframework")
+ || info.getName().startsWith("spock")
+ || info.getName().startsWith("java")
+ || info.getName().startsWith("sun")
+ || info.getName().startsWith("jdk")
+ || info.getName().startsWith("com.intellij")
+ || info.getName().startsWith("org.jetbrains")
+ || info.getName().startsWith("com.oracle")
+ || info.getName().startsWith("com.sun")
+ || info.getName().startsWith("ratpack")
+ || info.getName().startsWith("org.codehaus.groovy")
+ || info.getName().startsWith("org.groovy")
+ || info.getName().startsWith("groovy")) {
+ isAgentClass = false;
+ }
+ if (isAgentClass) {
+ agentClasses.add(info.getResourceName());
+ }
+ }
+ for (ClassPath.ResourceInfo resource : testCP.getResources()) {
+ if (resource.getResourceName().startsWith("META-INF/services/")) {
+ agentClasses.add(resource.getResourceName());
+ }
+ }
+ final File file = new File(TestUtils.createJarWithClasses(loader, agentClasses.toArray(new String[0])).getFile());
+ return file;
+ }
+
+ private static ClassLoader createDatadogClassLoader(File bootstrapJar, File toolingJar)
+ throws Exception {
+ final ClassLoader agentParent = ClassLoader.getSystemClassLoader().getParent();
+ Class<?> loaderClass =
+ ClassLoader.getSystemClassLoader().loadClass("datadog.trace.bootstrap.DatadogClassLoader");
+ Constructor constructor =
+ loaderClass.getDeclaredConstructor(URL.class, URL.class, ClassLoader.class);
+ return (ClassLoader)
+ constructor.newInstance(
+ bootstrapJar.toURI().toURL(), toolingJar.toURI().toURL(), agentParent);
}
+ */
}

ADD
dd-java-agent/testing/src/main/java/datadog/trace/agent/test/SpockRunner.java
+package datadog.trace.agent.test;
+
+import static datadog.trace.agent.tooling.Utils.BOOTSTRAP_PACKAGE_PREFIXES;
+
+import com.google.common.reflect.ClassPath;
+import java.io.File;
+import java.io.IOException;
+import java.lang.reflect.Field;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.jar.JarFile;
+import net.bytebuddy.agent.ByteBuddyAgent;
+import net.bytebuddy.dynamic.ClassFileLocator;
+import org.junit.runner.notification.RunNotifier;
+import org.junit.runners.model.InitializationError;
+import org.spockframework.runtime.Sputnik;
+
+/**
+ * Runs a spock test in an agent-friendly way.
+ *
+ * <ul>
+ * <li> Adds agent bootstrap classes to bootstrap classpath.
+ * <li> Runs tests in a custom classloader which cannot see core agent classes.
+ * </ul>
+ */
+public class SpockRunner extends Sputnik {
+ private static final String[] TEST_BOOTSTRAP_PREFIXES;
+
+ static {
+ ByteBuddyAgent.install();
+ final String[] testBS = {
+ "io.opentracing",
+ // Putting the logger on the bootstrap breaks some tests
+ // we can get away with keeping the logger in the system loader because we control the classloading env
+ // "org.slf4j",
+ // "ch.qos.logback"
+ };
+ TEST_BOOTSTRAP_PREFIXES =
+ Arrays.copyOf(
+ BOOTSTRAP_PACKAGE_PREFIXES, BOOTSTRAP_PACKAGE_PREFIXES.length + testBS.length);
+ for (int i = 0; i < testBS.length; ++i) {
+ TEST_BOOTSTRAP_PREFIXES[i + BOOTSTRAP_PACKAGE_PREFIXES.length] = testBS[i];
+ }
+
+ setupBootstrapClasspath();
+ }
+
+ private final InstrumentationClassLoader customLoader;
+
+ public SpockRunner(Class<?> clazz)
+ throws InitializationError, NoSuchFieldException, SecurityException, IllegalArgumentException,
+ IllegalAccessException {
+ super(shadowTestClass(clazz));
+ // access the classloader created in shadowTestClass above
+ Field clazzField = Sputnik.class.getDeclaredField("clazz");
+ try {
+ clazzField.setAccessible(true);
+ customLoader =
+ (InstrumentationClassLoader) ((Class<?>) clazzField.get(this)).getClassLoader();
+ } finally {
+ clazzField.setAccessible(false);
+ }
+ }
+
+ // Shadow the test class with bytes loaded by InstrumentationClassLoader
+ private static Class<?> shadowTestClass(final Class<?> clazz) {
+ try {
+ InstrumentationClassLoader customLoader =
+ new InstrumentationClassLoader(
+ datadog.trace.agent.test.SpockRunner.class.getClassLoader(), clazz.getName());
+ return customLoader.shadow(clazz);
+ } catch (Exception e) {
+ throw new IllegalStateException(e);
+ }
+ }
+
+ @Override
+ public void run(final RunNotifier notifier) {
+ final ClassLoader contextLoader = Thread.currentThread().getContextClassLoader();
+ try {
+ Thread.currentThread().setContextClassLoader(customLoader);
+ super.run(notifier);
+ } finally {
+ Thread.currentThread().setContextClassLoader(contextLoader);
+ }
+ }
+
+ private static void setupBootstrapClasspath() {
+ try {
+ final File bootstrapJar = createBootstrapJar();
+ ByteBuddyAgent.getInstrumentation()
+ .appendToBootstrapClassLoaderSearch(new JarFile(bootstrapJar));
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ private static File createBootstrapJar() throws IOException {
+ final ClassLoader loader = AgentTestRunner.class.getClassLoader();
+ final ClassPath testCP = ClassPath.from(loader);
+ Set<String> bootstrapClasses = new HashSet<String>();
+ for (ClassPath.ClassInfo info : testCP.getAllClasses()) {
+ // if info starts with bootstrap prefix: add to bootstrap jar
+ for (int i = 0; i < TEST_BOOTSTRAP_PREFIXES.length; ++i) {
+ if (info.getName().startsWith(TEST_BOOTSTRAP_PREFIXES[i])) {
+ bootstrapClasses.add(info.getResourceName());
+ break;
+ }
+ }
+ }
+ return new File(
+ TestUtils.createJarWithClasses(loader, bootstrapClasses.toArray(new String[0])).getFile());
+ }
+
+ /** TODO: Doc */
+ private static class InstrumentationClassLoader extends java.lang.ClassLoader {
+ final ClassLoader parent;
+ final String shadowPrefix;
+
+ public InstrumentationClassLoader(ClassLoader parent, String shadowPrefix) {
+ super(parent);
+ this.parent = parent;
+ this.shadowPrefix = shadowPrefix;
+ }
+
+ /** Forcefully inject the bytes of clazz into this classloader. */
+ public Class<?> shadow(Class<?> clazz) throws IOException {
+ Class<?> loaded = this.findLoadedClass(clazz.getName());
+ if (loaded != null && loaded.getClassLoader() == this) {
+ return loaded;
+ }
+ final ClassFileLocator locator = ClassFileLocator.ForClassLoader.of(clazz.getClassLoader());
+ final byte[] classBytes = locator.locate(clazz.getName()).resolve();
+
+ return this.defineClass(clazz.getName(), classBytes, 0, classBytes.length);
+ }
+
+ @Override
+ protected Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
+ Class c = this.findLoadedClass(name);
+ if (c != null) {
+ return c;
+ }
+ if (name.startsWith(shadowPrefix)) {
+ try {
+ return shadow(super.loadClass(name, resolve));
+ } catch (Exception e) {
+ }
+ }
+
+ /*
+ if (!name.startsWith("datadog.trace.agent.test.")) {
+ for (int i = 0; i < AGENT_PACKAGE_PREFIXES.length; ++i) {
+ if (name.startsWith(AGENT_PACKAGE_PREFIXES[i])) {
+ throw new ClassNotFoundException(
+ "refusing to load agent class" + name + " on test classloader.");
+ }
+ }
+ }
+ */
+ return parent.loadClass(name);
+ }
+ }
+}

MODIFY
dd-java-agent/testing/src/main/java/datadog/trace/agent/test/TestUtils.java
dd-java-agent/testing/src/main/java/datadog/trace/agent/test/TestUtils.java
@@ -4,10 +4,7 @@ import datadog.trace.agent.tooling.Utils;
import io.opentracing.Scope;
import io.opentracing.Tracer;
import io.opentracing.util.GlobalTracer;
-import java.io.File;
-import java.io.FileOutputStream;
-import java.io.IOException;
-import java.io.InputStream;
+import java.io.*;
import java.lang.reflect.Field;
import java.net.URL;
import java.util.UUID;
@@ -42,6 +39,22 @@ public class TestUtils {
}
}
+ /** Get the tracer implementation out of the GlobalTracer */
+ public static Tracer getUnderlyingGlobalTracer() {
+ Field field = null;
+ try {
+ field = GlobalTracer.class.getDeclaredField("tracer");
+ field.setAccessible(true);
+ return (Tracer) field.get(GlobalTracer.get());
+ } catch (final Exception e2) {
+ throw new IllegalStateException(e2);
+ } finally {
+ if (null != field) {
+ field.setAccessible(false);
+ }
+ }
+ }
+
public static <T extends Object> Object runUnderTrace(
final String rootOperationName, final Callable<T> r) {
final Scope scope = GlobalTracer.get().buildSpan(rootOperationName).startActive(true);
@@ -54,6 +67,63 @@ public class TestUtils {
}
}
+ public static byte[] convertToByteArray(InputStream resource) throws IOException {
+ ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+ int bytesRead;
+ byte[] data = new byte[1024];
+ while ((bytesRead = resource.read(data, 0, data.length)) != -1) {
+ buffer.write(data, 0, bytesRead);
+ }
+ buffer.flush();
+ return buffer.toByteArray();
+ }
+
+ public static byte[] convertToByteArray(Class<?> clazz) throws IOException {
+ InputStream inputStream = null;
+ try {
+ inputStream =
+ clazz.getClassLoader().getResourceAsStream(Utils.getResourceName(clazz.getName()));
+ return convertToByteArray(inputStream);
+ } finally {
+ if (inputStream != null) {
+ inputStream.close();
+ }
+ }
+ }
+
+ /**
+ * Create a temporary jar on the filesystem with the bytes of the given classes.
+ *
+ * <p>The jar file will be removed when the jvm exits.
+ *
+ * @param loader classloader used to load bytes
+ * @param resourceNames names of resources to copy into the new jar
+ * @return the location of the newly created jar.
+ * @throws IOException
+ */
+ public static URL createJarWithClasses(final ClassLoader loader, final String... resourceNames)
+ throws IOException {
+ final File tmpJar = File.createTempFile(UUID.randomUUID().toString() + "", ".jar");
+ tmpJar.deleteOnExit();
+
+ final Manifest manifest = new Manifest();
+ final JarOutputStream target = new JarOutputStream(new FileOutputStream(tmpJar), manifest);
+ for (final String resourceName : resourceNames) {
+ InputStream is = null;
+ try {
+ is = loader.getResourceAsStream(resourceName);
+ addToJar(resourceName, convertToByteArray(is), target);
+ } finally {
+ if (null != is) {
+ is.close();
+ }
+ }
+ }
+ target.close();
+
+ return tmpJar.toURI().toURL();
+ }
+
/**
* Create a temporary jar on the filesystem with the bytes of the given classes.
*
@@ -70,35 +140,24 @@ public class TestUtils {
final Manifest manifest = new Manifest();
final JarOutputStream target = new JarOutputStream(new FileOutputStream(tmpJar), manifest);
for (final Class<?> clazz : classes) {
- addToJar(clazz, target);
+ addToJar(Utils.getResourceName(clazz.getName()), convertToByteArray(clazz), target);
}
target.close();
return tmpJar.toURI().toURL();
}
- private static void addToJar(final Class<?> clazz, final JarOutputStream jarOutputStream)
- throws IOException {
- InputStream inputStream = null;
- try {
- final JarEntry entry = new JarEntry(Utils.getResourceName(clazz.getName()));
- jarOutputStream.putNextEntry(entry);
- inputStream =
- clazz.getClassLoader().getResourceAsStream(Utils.getResourceName(clazz.getName()));
+ public static URL createJarWithClasses() {
- final byte[] buffer = new byte[1024];
- while (true) {
- final int count = inputStream.read(buffer);
- if (count == -1) {
- break;
- }
- jarOutputStream.write(buffer, 0, count);
+ return null;
}
+
+ private static void addToJar(
+ final String resourceName, final byte[] bytes, final JarOutputStream jarOutputStream)
+ throws IOException {
+ final JarEntry entry = new JarEntry(resourceName);
+ jarOutputStream.putNextEntry(entry);
+ jarOutputStream.write(bytes, 0, bytes.length);
jarOutputStream.closeEntry();
- } finally {
- if (inputStream != null) {
- inputStream.close();
- }
- }
}
}

ADD
dd-java-agent/testing/src/test/groovy/AgentTestRunnerTest.groovy
+import datadog.trace.agent.test.TestUtils
+
+import java.lang.reflect.Field
+
+import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER
+
+import datadog.trace.agent.test.AgentTestRunner
+
+class AgentTestRunnerTest extends AgentTestRunner {
+ static {
+ // when test class initializes, opentracing should be set up, but not the agent.
+ assert io.opentracing.Tracer.getClassLoader() == BOOTSTRAP_CLASSLOADER
+ assert getAgentTransformer() == null
+ }
+
+ def "classpath setup"() {
+ expect:
+ io.opentracing.Tracer.getClassLoader() == BOOTSTRAP_CLASSLOADER
+ TEST_TRACER == TestUtils.getUnderlyingGlobalTracer()
+ getAgentTransformer() != null
+ }
+
+ def "logging works"() {
+ when:
+ org.slf4j.LoggerFactory.getLogger(AgentTestRunnerTest).debug("hello")
+ then:
+ noExceptionThrown()
+ }
+
+ def "can't see agent classes"() {
+ // TODO
+ }
+
+ private static getAgentTransformer() {
+ Field f
+ try {
+ f = AgentTestRunner.getDeclaredField("activeTransformer")
+ f.setAccessible(true)
+ return f.get(null)
+ } finally {
+ f.setAccessible(false)
+ }
+ }
+}

Run tests in a custom classloader with bootstrap classpath