Change agent jar inclusion in tests that launch a process

The tests were the main problem. By using a different approach to pass in the agent jar, the TracingAgent code can be much simpler

dd-java-agent/dd-java-agent.gradle

@@ -78,7 +78,7 @@ jar {
 
 shadowJar {
   configurations = [project.configurations.shadowInclude]
-  
+
   classifier null
 
   mergeServiceFiles()
@@ -114,22 +114,18 @@ dependencies {
   testCompile deps.opentracingMock
   testCompile deps.testLogging
   testCompile deps.guava
-  
+
   shadowInclude project(':dd-java-agent:agent-bootstrap')
 }
 
 tasks.withType(Test).configureEach {
   jvmArgs "-Ddd.service.name=java-agent-tests"
   jvmArgs "-Ddd.writer.type=LoggingWriter"
+  jvmArgs "-Dagent.jar.to.forward=${shadowJar.archivePath}"
   // Multi-threaded logging seems to be causing deadlocks with Gradle's log capture.
 //  jvmArgs "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=debug"
 //  jvmArgs "-Dorg.slf4j.simpleLogger.defaultLogLevel=debug"
 
-  doFirst {
-    // Defining here to allow jacoco to be first on the command line.
-    jvmArgs "-javaagent:${shadowJar.archivePath}"
-  }
-
   testLogging {
     events "started"
   }

dd-java-agent/src/main/java/datadog/trace/agent/TracingAgent.java

@@ -152,32 +152,16 @@ public class TracingAgent {
   private static synchronized URL installBootstrapJar(
       final Instrumentation inst, final boolean usingCustomLogManager)
       throws IOException, URISyntaxException, ClassNotFoundException {
-    URL bootstrapURL = null;
-    File bootstrapFile = null;
-
     final CodeSource codeSource = TracingAgent.class.getProtectionDomain().getCodeSource();
+
     if (codeSource != null) {
-      bootstrapURL = codeSource.getLocation();
+      final URL bootstrapURL = codeSource.getLocation();
 
-      bootstrapFile = new File(bootstrapURL.toURI());
+      inst.appendToBootstrapClassLoaderSearch(new JarFile(new File(bootstrapURL.toURI())));
+      return bootstrapURL;
+    } else {
+      throw new RuntimeException("Unable to install bootstrap jar.  Code source not found");
     }
-
-    if (bootstrapFile == null || bootstrapFile.isDirectory()) {
-      // At most one of ( DatadogClassLoader.class , TracingAgent.class ) has a CodeSource that
-      // doesn't lead to the agent jar
-
-      bootstrapURL =
-          ClassLoader.getSystemClassLoader()
-              .loadClass("datadog.trace.bootstrap.DatadogClassLoader")
-              .getProtectionDomain()
-              .getCodeSource()
-              .getLocation();
-
-      bootstrapFile = new File(bootstrapURL.toURI());
-    }
-
-    inst.appendToBootstrapClassLoaderSearch(new JarFile(bootstrapFile));
-    return bootstrapURL;
   }
 
   /**
@@ -186,11 +170,11 @@ public class TracingAgent {
    *
    * @param innerJarFilename Filename of internal jar to use for the classpath of the datadog
    *     classloader
-   * @param bootStrapURL
+   * @param bootstrapURL
    * @return Datadog Classloader
    */
   private static ClassLoader createDatadogClassLoader(
-      final String innerJarFilename, final URL bootStrapURL) throws Exception {
+      final String innerJarFilename, final URL bootstrapURL) throws Exception {
     final ClassLoader agentParent;
     final String javaVersion = System.getProperty("java.version");
     if (javaVersion.startsWith("1.7") || javaVersion.startsWith("1.8")) {
@@ -204,7 +188,7 @@ public class TracingAgent {
         ClassLoader.getSystemClassLoader().loadClass("datadog.trace.bootstrap.DatadogClassLoader");
     final Constructor constructor =
         loaderClass.getDeclaredConstructor(URL.class, String.class, ClassLoader.class);
-    return (ClassLoader) constructor.newInstance(bootStrapURL, innerJarFilename, agentParent);
+    return (ClassLoader) constructor.newInstance(bootstrapURL, innerJarFilename, agentParent);
   }
 
   private static ClassLoader getPlatformClassLoader()

dd-java-agent/src/test/groovy/datadog/trace/agent/AgentLoadedIntoBootstrapTest.groovy

@@ -2,31 +2,15 @@ package datadog.trace.agent
 
 import datadog.trace.agent.test.IntegrationTestUtils
 import jvmbootstraptest.AgentLoadedChecker
-import spock.lang.Shared
 import spock.lang.Specification
 
-import java.lang.management.ManagementFactory
-import java.lang.management.RuntimeMXBean
-
 class AgentLoadedIntoBootstrapTest extends Specification {
-  @Shared
-  private agentArg
 
-  def setupSpec() {
-    final RuntimeMXBean runtimeMxBean = ManagementFactory.getRuntimeMXBean()
-    for (String arg : runtimeMxBean.getInputArguments()) {
-      if (arg.startsWith("-javaagent")) {
-        agentArg = arg
-        break
-      }
-    }
-    assert agentArg != null
-  }
 
   def "Agent loads in when separate jvm is launched"() {
     expect:
     IntegrationTestUtils.runOnSeparateJvm(AgentLoadedChecker.getName()
-      , [agentArg] as String[]
+      , "" as String[]
       , "" as String[]
       , [:]
       , true) == 0

dd-java-agent/src/test/groovy/datadog/trace/agent/CustomLogManagerTest.groovy

@@ -4,38 +4,19 @@ import datadog.trace.agent.test.IntegrationTestUtils
 import jvmbootstraptest.LogManagerSetter
 import spock.lang.Requires
 import spock.lang.Retry
-import spock.lang.Shared
 import spock.lang.Specification
 import spock.lang.Timeout
 
-import java.lang.management.ManagementFactory
-import java.lang.management.RuntimeMXBean
-
 // Note: this test is fails on IBM JVM, we would need to investigate this at some point
 @Requires({ !System.getProperty("java.vm.name").contains("IBM J9 VM") })
 @Retry
 @Timeout(30)
 class CustomLogManagerTest extends Specification {
   // Run all tests using forked jvm because groovy has already set the global log manager
-
-  @Shared
-  private agentArg
-
-  def setupSpec() {
-    final RuntimeMXBean runtimeMxBean = ManagementFactory.getRuntimeMXBean()
-    for (String arg : runtimeMxBean.getInputArguments()) {
-      if (arg.startsWith("-javaagent")) {
-        agentArg = arg
-        break
-      }
-    }
-    assert agentArg != null
-  }
-
   def "jmxfetch starts up in premain with no custom log manager set"() {
     expect:
     IntegrationTestUtils.runOnSeparateJvm(LogManagerSetter.getName()
-      , [agentArg, "-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off"] as String[]
+      , ["-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off"] as String[]
       , "" as String[]
       , [:]
       , true) == 0
@@ -44,7 +25,7 @@ class CustomLogManagerTest extends Specification {
   def "jmxfetch starts up in premain if configured log manager on system classpath"() {
     expect:
     IntegrationTestUtils.runOnSeparateJvm(LogManagerSetter.getName()
-      , [agentArg, "-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off", "-Djava.util.logging.manager=jvmbootstraptest.CustomLogManager"] as String[]
+      , ["-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off", "-Djava.util.logging.manager=jvmbootstraptest.CustomLogManager"] as String[]
       , "" as String[]
       , [:]
       , true) == 0
@@ -53,7 +34,7 @@ class CustomLogManagerTest extends Specification {
   def "jmxfetch startup is delayed with java.util.logging.manager sysprop"() {
     expect:
     IntegrationTestUtils.runOnSeparateJvm(LogManagerSetter.getName()
-      , [agentArg, "-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off", "-Djava.util.logging.manager=jvmbootstraptest.MissingLogManager"] as String[]
+      , ["-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off", "-Djava.util.logging.manager=jvmbootstraptest.MissingLogManager"] as String[]
       , "" as String[]
       , [:]
       , true) == 0
@@ -62,7 +43,7 @@ class CustomLogManagerTest extends Specification {
   def "jmxfetch startup delayed with tracer custom log manager setting"() {
     expect:
     IntegrationTestUtils.runOnSeparateJvm(LogManagerSetter.getName()
-      , [agentArg, "-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off", "-Ddd.app.customlogmanager=true"] as String[]
+      , ["-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off", "-Ddd.app.customlogmanager=true"] as String[]
       , "" as String[]
       , [:]
       , true) == 0
@@ -71,7 +52,7 @@ class CustomLogManagerTest extends Specification {
   def "jmxfetch startup delayed with JBOSS_HOME environment variable"() {
     expect:
     IntegrationTestUtils.runOnSeparateJvm(LogManagerSetter.getName()
-      , [agentArg, "-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off", "-Ddd.app.customlogmanager=true"] as String[]
+      , ["-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off", "-Ddd.app.customlogmanager=true"] as String[]
       , "" as String[]
       , ["JBOSS_HOME": "/"]
       , true) == 0
@@ -80,7 +61,7 @@ class CustomLogManagerTest extends Specification {
   def "jmxfetch startup in premain forced by customlogmanager=false"() {
     expect:
     IntegrationTestUtils.runOnSeparateJvm(LogManagerSetter.getName()
-      , [agentArg, "-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off", "-Ddd.app.customlogmanager=false", "-Djava.util.logging.manager=jvmbootstraptest.CustomLogManager"] as String[]
+      , ["-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off", "-Ddd.app.customlogmanager=false", "-Djava.util.logging.manager=jvmbootstraptest.CustomLogManager"] as String[]
       , "" as String[]
       , ["JBOSS_HOME": "/"]
       , true) == 0

dd-java-agent/src/test/java/datadog/trace/agent/test/IntegrationTestUtils.java

@@ -190,16 +190,29 @@ public class IntegrationTestUtils {
       final Map<String, String> envVars,
       final boolean printOutputStreams)
       throws Exception {
+
     final String separator = System.getProperty("file.separator");
-    final String classpath = System.getProperty("java.class.path");
+    final String agentJar = System.getProperty("agent.jar.to.forward");
+
+    if (agentJar == null) {
+      throw new RuntimeException("agent.jar.to.forward property must be set");
+    }
+
+    final String classpath =
+        agentJar + System.getProperty("path.separator") + System.getProperty("java.class.path");
     final String path = System.getProperty("java.home") + separator + "bin" + separator + "java";
+
+    final List<String> vmArgsList = new ArrayList<>(Arrays.asList(jvmArgs));
+    vmArgsList.add("-javaagent:" + agentJar);
+
     final List<String> commands = new ArrayList<>();
     commands.add(path);
-    commands.addAll(Arrays.asList(jvmArgs));
+    commands.addAll(vmArgsList);
     commands.add("-cp");
     commands.add(classpath);
     commands.add(mainClassName);
     commands.addAll(Arrays.asList(mainMethodArgs));
+
     final ProcessBuilder processBuilder = new ProcessBuilder(commands.toArray(new String[0]));
     processBuilder.environment().putAll(envVars);
     final Process process = processBuilder.start();

dd-java-agent/src/test/java/jvmbootstraptest/AgentLoadedChecker.java

@@ -2,12 +2,9 @@ package jvmbootstraptest;
 
 import java.net.URL;
 import java.net.URLClassLoader;
-import java.util.Arrays;
 
 public class AgentLoadedChecker {
   public static void main(final String[] args) throws ClassNotFoundException {
-    System.out.println(Arrays.toString(args));
-
     // Empty classloader that delegates to bootstrap
     final URLClassLoader emptyClassLoader = new URLClassLoader(new URL[] {}, null);
     final Class agentClass = emptyClassLoader.loadClass("datadog.trace.agent.TracingAgent");