Minor cleanups

Author: nikita-tkachenko-datadog

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/BootstrapClasspathSetup.java

@@ -11,12 +11,9 @@
 import de.thetaphi.forbiddenapis.SuppressForbidden;
 import java.io.File;
 import java.io.IOException;
-import java.lang.reflect.Field;
-import java.lang.reflect.Method;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.net.URLClassLoader;
-import java.util.Arrays;
 import java.util.HashSet;
 import java.util.Set;
 import java.util.TreeSet;
@@ -29,16 +26,15 @@ public class BootstrapClasspathSetup implements LauncherSessionListener {
 
   @Override
   public void launcherSessionOpened(LauncherSession session) {
-    // do nothing
+    // this method is only needed to trigger this class' static initializer before JUnit does
+    // classpath scanning
   }
 
-  /**
-   * An exact copy of {@link datadog.trace.bootstrap.Constants#BOOTSTRAP_PACKAGE_PREFIXES}.
-   *
-   * <p>This list is needed to initialize the bootstrap classpath because Utils' static initializer
-   * references bootstrap classes (e.g. DatadogClassLoader).
-   */
-  public static final String[] BOOTSTRAP_PACKAGE_PREFIXES_COPY = {
+  private static final String[] TEST_EXCLUDED_BOOTSTRAP_PACKAGE_PREFIXES = {
+    "ch.qos.logback.classic.servlet", // this draws javax.servlet deps that are not needed
+  };
+
+  private static final String[] TEST_BOOTSTRAP_PREFIXES = {
     "datadog.slf4j",
     "datadog.context",
     "datadog.environment",
@@ -51,32 +47,14 @@ public void launcherSessionOpened(LauncherSession session) {
     "datadog.trace.instrumentation.api",
     "datadog.trace.logging",
     "datadog.trace.util",
+    "org.slf4j",
+    "ch.qos.logback",
   };
 
-  private static final String[] TEST_EXCLUDED_BOOTSTRAP_PACKAGE_PREFIXES = {
-    "ch.qos.logback.classic.servlet", // this draws javax.servlet deps that are not needed
-  };
-
-  private static final String[] TEST_BOOTSTRAP_PREFIXES;
-
-  public static final ClassPath TEST_CLASSPATH;
+  public static final ClassPath TEST_CLASSPATH = computeTestClasspath();
 
   static {
     ByteBuddyAgent.install();
-    final String[] testBS = {
-      "org.slf4j", "ch.qos.logback",
-    };
-
-    TEST_BOOTSTRAP_PREFIXES =
-        Arrays.copyOf(
-            BOOTSTRAP_PACKAGE_PREFIXES_COPY,
-            BOOTSTRAP_PACKAGE_PREFIXES_COPY.length + testBS.length);
-    for (int i = 0; i < testBS.length; ++i) {
-      TEST_BOOTSTRAP_PREFIXES[i + BOOTSTRAP_PACKAGE_PREFIXES_COPY.length] = testBS[i];
-    }
-
-    TEST_CLASSPATH = computeTestClasspath();
-
     setupBootstrapClasspath();
   }
 
@@ -116,41 +94,6 @@ private static ClassLoader buildJavaClassPathClassLoader() {
     return new URLClassLoader(urls.build().toArray(new URL[0]), null);
   }
 
-  public static void assertNoBootstrapClassesInTestClass(final Class<?> testClass) {
-    for (final Field field : testClass.getDeclaredFields()) {
-      assertNotBootstrapClass(testClass, field.getType());
-    }
-    for (final Method method : testClass.getDeclaredMethods()) {
-      assertNotBootstrapClass(testClass, method.getReturnType());
-      for (final Class<?> paramType : method.getParameterTypes()) {
-        assertNotBootstrapClass(testClass, paramType);
-      }
-    }
-  }
-
-  private static void assertNotBootstrapClass(final Class<?> testClass, final Class<?> clazz) {
-    if (!clazz.isPrimitive() && isBootstrapClass(clazz.getName())) {
-      throw new IllegalStateException(
-          testClass.getName()
-              + ": Bootstrap classes are not allowed in test class field or method signatures. Offending class: "
-              + clazz.getName());
-    }
-  }
-
-  private static boolean isBootstrapClass(final String name) {
-    for (String prefix : TEST_BOOTSTRAP_PREFIXES) {
-      if (name.startsWith(prefix)) {
-        for (String excluded : TEST_EXCLUDED_BOOTSTRAP_PACKAGE_PREFIXES) {
-          if (name.startsWith(excluded)) {
-            return false;
-          }
-        }
-        return true;
-      }
-    }
-    return false;
-  }
-
   private static void setupBootstrapClasspath() {
     // Ensure there weren't any bootstrap classes loaded prematurely.
     Set<String> prematureBootstrapClasses = new TreeSet<>();
@@ -172,8 +115,6 @@ private static void setupBootstrapClasspath() {
       final File bootstrapJar = createBootstrapJar();
       ByteBuddyAgent.getInstrumentation()
           .appendToBootstrapClassLoaderSearch(new JarFile(bootstrapJar));
-      // Utils cannot be referenced before this line, as its static initializers load bootstrap
-      // classes (for example, the bootstrap proxy).
       BootstrapProxy.addBootstrapResource(bootstrapJar.toURI().toURL());
     } catch (final IOException e) {
       throw new RuntimeException(e);
@@ -192,4 +133,18 @@ private static File createBootstrapJar() throws IOException {
             SpockExtension.class.getClassLoader(), bootstrapClasses.toArray(new String[0]));
     return new File(jar.getFile());
   }
+
+  public static boolean isBootstrapClass(final String name) {
+    for (String prefix : TEST_BOOTSTRAP_PREFIXES) {
+      if (name.startsWith(prefix)) {
+        for (String excluded : TEST_EXCLUDED_BOOTSTRAP_PACKAGE_PREFIXES) {
+          if (name.startsWith(excluded)) {
+            return false;
+          }
+        }
+        return true;
+      }
+    }
+    return false;
+  }
 }

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/SpockExtension.java

@@ -1,6 +1,8 @@
 package datadog.trace.agent.test;
 
 import java.io.IOException;
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
 import java.util.List;
 import net.bytebuddy.dynamic.ClassFileLocator;
 import org.junit.jupiter.api.extension.AfterAllCallback;
@@ -20,16 +22,40 @@ public final class SpockExtension
         TestExecutionExceptionHandler {
 
   private static final ExtensionContext.Namespace NAMESPACE =
-      ExtensionContext.Namespace.create("dd", "spock", "loader");
+      ExtensionContext.Namespace.create("dd", "spock");
+
+  private static final String INSTRUMENTATION_CLASSLOADER = "instrumentation-class-loader";
+  private static final String ORIGINAL_CLASSLOADER = "original-thread-context-class-loader";
 
   @Override
   public void beforeAll(ExtensionContext ctx) {
     Class<?> testClass = ctx.getRequiredTestClass();
-    BootstrapClasspathSetup.assertNoBootstrapClassesInTestClass(testClass);
+    assertNoBootstrapClassesInTestClass(testClass);
 
     InstrumentationClassLoader custom =
         new InstrumentationClassLoader(testClass.getClassLoader(), testClass.getName());
-    ctx.getStore(NAMESPACE).put("loader", custom);
+    ctx.getStore(NAMESPACE).put(INSTRUMENTATION_CLASSLOADER, custom);
+  }
+
+  public static void assertNoBootstrapClassesInTestClass(final Class<?> testClass) {
+    for (final Field field : testClass.getDeclaredFields()) {
+      assertNotBootstrapClass(testClass, field.getType());
+    }
+    for (final Method method : testClass.getDeclaredMethods()) {
+      assertNotBootstrapClass(testClass, method.getReturnType());
+      for (final Class<?> paramType : method.getParameterTypes()) {
+        assertNotBootstrapClass(testClass, paramType);
+      }
+    }
+  }
+
+  private static void assertNotBootstrapClass(final Class<?> testClass, final Class<?> clazz) {
+    if (!clazz.isPrimitive() && BootstrapClasspathSetup.isBootstrapClass(clazz.getName())) {
+      throw new IllegalStateException(
+          testClass.getName()
+              + ": Bootstrap classes are not allowed in test class field or method signatures. Offending class: "
+              + clazz.getName());
+    }
   }
 
   @Override
@@ -39,14 +65,16 @@ public void afterAll(ExtensionContext ctx) {
 
   @Override
   public void beforeEach(ExtensionContext ctx) {
-    ClassLoader custom = ctx.getStore(NAMESPACE).get("loader", ClassLoader.class);
-    ctx.getStore(NAMESPACE).put("prevTCCL", Thread.currentThread().getContextClassLoader());
-    Thread.currentThread().setContextClassLoader(custom);
+    ClassLoader instrumentationClassloader =
+        ctx.getStore(NAMESPACE).get(INSTRUMENTATION_CLASSLOADER, ClassLoader.class);
+    ctx.getStore(NAMESPACE)
+        .put(ORIGINAL_CLASSLOADER, Thread.currentThread().getContextClassLoader());
+    Thread.currentThread().setContextClassLoader(instrumentationClassloader);
   }
 
   @Override
   public void afterEach(ExtensionContext ctx) {
-    ClassLoader prev = ctx.getStore(NAMESPACE).remove("prevTCCL", ClassLoader.class);
+    ClassLoader prev = ctx.getStore(NAMESPACE).remove(ORIGINAL_CLASSLOADER, ClassLoader.class);
     if (prev != null) {
       Thread.currentThread().setContextClassLoader(prev);
     }
@@ -86,11 +114,7 @@ static void fixTooManyInvocationsError(final TooManyInvocationsError error) {
     }
   }
 
-  /**
-   * Class‑loader that <em>shadows</em> the test class, so that any re‑definitions via ByteBuddy
-   * operate on a loader‑private copy instead of the one used by the build toolʼs class‑path
-   * scanning. Delegation order = child‑first for the testʼs own package, parent‑first otherwise.
-   */
+  /** Run test classes in a classloader which loads test classes before delegating. */
   private static class InstrumentationClassLoader extends ClassLoader {
     private final ClassLoader parent;
     private final String shadowPrefix;
@@ -101,7 +125,7 @@ private static class InstrumentationClassLoader extends ClassLoader {
       this.shadowPrefix = shadowPrefix;
     }
 
-    /** Inject the bytes of {@code clazz} into <b>this</b> loader, producing a shadow copy. */
+    /** Forcefully inject the bytes of clazz into this classloader. */
     Class<?> shadow(Class<?> clazz) throws IOException {
       Class<?> loaded = findLoadedClass(clazz.getName());
       if (loaded != null && loaded.getClassLoader() == this) {