DataDog · mcculls · Feb 12, 2021 · Feb 12, 2021
@@ -1,93 +1,94 @@
 package datadog.trace.agent
 
 import datadog.trace.agent.test.IntegrationTestUtils
-import datadog.trace.api.Config
 import jvmbootstraptest.AgentLoadedChecker
-import org.junit.Rule
-import org.junit.contrib.java.lang.system.RestoreSystemProperties
+import jvmbootstraptest.JmxStartedChecker
+import spock.lang.Shared
 import spock.lang.Specification
 import spock.lang.Timeout
 
-import java.lang.reflect.Method
-
 @Timeout(30)
 class JMXFetchTest extends Specification {
+  @Shared
+  DatagramSocket jmxStatsSocket
+
+  def setupSpec() {
+    jmxStatsSocket = new DatagramSocket(0)
+  }
 
-  @Rule
-  public final RestoreSystemProperties restoreSystemProperties = new RestoreSystemProperties()
+  def cleanupSpec() {
+    jmxStatsSocket.close()
+  }
 
   def "test jmxfetch"() {
     setup:
-    def currentContextLoader = Thread.currentThread().getContextClassLoader()
-    DatagramSocket socket = new DatagramSocket(0)
-
-    System.setProperty("dd.jmxfetch.enabled", "true")
-    System.setProperty("dd.jmxfetch.start-delay", "0")
-    System.setProperty("dd.jmxfetch.statsd.port", Integer.toString(socket.localPort))
-    // Overwrite writer type to disable console jmxfetch reporter
-    System.setProperty("dd.writer.type", "DDAgentWriter")
-
-    def classLoader = IntegrationTestUtils.getJmxFetchClassLoader()
-    // Have to set this so JMXFetch knows where to find resources
-    Thread.currentThread().setContextClassLoader(classLoader)
-    final Class<?> jmxFetchAgentClass =
-      classLoader.loadClass("datadog.trace.agent.jmxfetch.JMXFetch")
-    final Method jmxFetchInstallerMethod = jmxFetchAgentClass.getDeclaredMethod("run", Config)
-    jmxFetchInstallerMethod.setAccessible(true)
-    jmxFetchInstallerMethod.invoke(null, new Config())
+    // verify that JMX starts and reports metrics through the given socket.
+    def returnCode = IntegrationTestUtils.runOnSeparateJvm(JmxStartedChecker.getName()
+      , ["-Ddd.jmxfetch.enabled=true",
+         "-Ddd.jmxfetch.start-delay=0",
+         "-Ddd.jmxfetch.statsd.port=${jmxStatsSocket.localPort}",
+         "-Ddd.writer.type=DDAgentWriter"] as String[]
+      , "" as String[]
+      , [:]
+      , true)
 
     byte[] buf = new byte[1500]
     DatagramPacket packet = new DatagramPacket(buf, buf.length)
-    socket.receive(packet)
+    jmxStatsSocket.receive(packet)
     String received = new String(packet.getData(), 0, packet.getLength())
 
-    Set<String> threads = Thread.getAllStackTraces().keySet().collect { it.name }
-
     expect:
-    threads.contains("dd-jmx-collector")
-    received.contains("jvm.")
-
-    cleanup:
-    jmxFetchInstallerMethod.setAccessible(false)
-    socket.close()
-    Thread.currentThread().setContextClassLoader(currentContextLoader)
+    returnCode == 0
+    received.contains("#service:${JmxStartedChecker.getName()}")
   }
 
   def "Agent loads when JmxFetch is misconfigured"() {
+    setup:
     // verify the agent starts up correctly with a bogus address.
-    expect:
-    IntegrationTestUtils.runOnSeparateJvm(AgentLoadedChecker.getName()
+    def returnCode = IntegrationTestUtils.runOnSeparateJvm(AgentLoadedChecker.getName()
       , ["-Ddd.jmxfetch.enabled=true",
          "-Ddd.jmxfetch.start-delay=0",
-         "-Ddd.jmxfetch.statsd.host=example.local"] as String[]
+         "-Ddd.jmxfetch.statsd.host=example.local",
+         "-Ddd.writer.type=DDAgentWriter"] as String[]
       , "" as String[]
       , [:]
-      , true) == 0
+      , true)
+
+    expect:
+    returnCode == 0
   }
 
   def "test jmxfetch config"() {
     setup:
-    names.each {
-      System.setProperty("dd.jmxfetch.${it}.enabled", "$enable")
+    def configSettings = names.collect {
+      "-Ddd.jmxfetch.${it}.enabled=${enable}"
+    }
+    def testOutput = new ByteArrayOutputStream()
+    def returnCode = IntegrationTestUtils.runOnSeparateJvm(JmxStartedChecker.getName()
+      , ["-Ddd.jmxfetch.enabled=true",
+         "-Ddd.jmxfetch.start-delay=0",
+         "-Ddd.jmxfetch.statsd.port=${jmxStatsSocket.localPort}",
+         "-Ddd.trace.debug=true",
+         "-Ddd.writer.type=DDAgentWriter"] + configSettings as String[]
+      , "" as String[]
+      , [:]
+      , new PrintStream(testOutput))
+
+    def actualConfig = []
+    new ByteArrayInputStream((testOutput.toByteArray())).eachLine {
+      System.out.println(it)
+      def match = (it =~ 'Reading metric config resource (.*)')
+      if (match) {
+        actualConfig += match[0][1]
+      }
     }
-    def classLoader = IntegrationTestUtils.getJmxFetchClassLoader()
-    // Have to set this so JMXFetch knows where to find resources
-    Thread.currentThread().setContextClassLoader(classLoader)
-    final Class<?> jmxFetchAgentClass =
-      classLoader.loadClass("datadog.trace.agent.jmxfetch.JMXFetch")
-    final Method jmxFetchInstallerMethod = jmxFetchAgentClass.getDeclaredMethod("getInternalMetricFiles")
-    jmxFetchInstallerMethod.setAccessible(true)
 
     expect:
-    jmxFetchInstallerMethod.invoke(null).sort() == result.sort()
-
-    cleanup:
-    names.each {
-      System.clearProperty("dd.jmxfetch.${it}.enabled")
-    }
+    returnCode == 0
+    actualConfig as Set == expectedConfig as Set
 
     where:
-    names               | enable | result
+    names               | enable | expectedConfig
     []                  | true   | []
     ["tomcat"]          | false  | []
     ["tomcat"]          | true   | ["datadog/trace/agent/jmxfetch/metricconfigs/tomcat.yaml"]

@@ -9,6 +9,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
+import java.io.PrintStream;
 import java.lang.management.ManagementFactory;
 import java.lang.management.RuntimeMXBean;
 import java.lang.reflect.Field;
@@ -33,11 +34,6 @@ public static ClassLoader getAgentClassLoader() {
     return getAgentFieldClassloader("AGENT_CLASSLOADER");
   }
 
-  /** Returns the classloader the jmxfetch is running on. */
-  public static ClassLoader getJmxFetchClassLoader() {
-    return getAgentFieldClassloader("JMXFETCH_CLASSLOADER");
-  }
-
   private static ClassLoader getAgentFieldClassloader(final String fieldName) {
     Field classloaderField = null;
     try {
@@ -157,16 +153,43 @@ public static int runOnSeparateJvm(
       final Map<String, String> envVars,
       final boolean printOutputStreams)
       throws Exception {
+    return runOnSeparateJvm(
+        mainClassName, jvmArgs, mainMethodArgs, envVars, printOutputStreams ? System.out : null);
+  }
+
+  public static int runOnSeparateJvm(
+      final String mainClassName,
+      final String[] jvmArgs,
+      final String[] mainMethodArgs,
+      final Map<String, String> envVars,
+      final PrintStream out)
+      throws Exception {
     final String classPath = System.getProperty("java.class.path");
+    return runOnSeparateJvm(mainClassName, jvmArgs, mainMethodArgs, envVars, classPath, out);
+  }
+
+  public static int runOnSeparateJvm(
+      final String mainClassName,
+      final String[] jvmArgs,
+      final String[] mainMethodArgs,
+      final Map<String, String> envVars,
+      final String classpath,
+      final boolean printOutputStreams)
+      throws Exception {
     return runOnSeparateJvm(
-        mainClassName, jvmArgs, mainMethodArgs, envVars, classPath, printOutputStreams);
+        mainClassName,
+        jvmArgs,
+        mainMethodArgs,
+        envVars,
+        classpath,
+        printOutputStreams ? System.out : null);
   }
 
   /**
    * On a separate JVM, run the main method for a given class.
    *
    * @param mainClassName The name of the entry point class. Must declare a main method.
-   * @param printOutputStreams if true, print stdout and stderr of the child jvm
+   * @param out Optional stream to print the stdout and stderr of the child jvm
    * @return the return code of the child jvm
    * @throws Exception
    */
@@ -176,7 +199,7 @@ public static int runOnSeparateJvm(
       final String[] mainMethodArgs,
       final Map<String, String> envVars,
       final String classpath,
-      final boolean printOutputStreams)
+      final PrintStream out)
       throws Exception {
 
     final String separator = System.getProperty("file.separator");
@@ -201,10 +224,8 @@ public static int runOnSeparateJvm(
 
     final Process process = processBuilder.start();
 
-    final StreamGobbler errorGobbler =
-        new StreamGobbler(process.getErrorStream(), "ERROR", printOutputStreams);
-    final StreamGobbler outputGobbler =
-        new StreamGobbler(process.getInputStream(), "OUTPUT", printOutputStreams);
+    final StreamGobbler errorGobbler = new StreamGobbler(process.getErrorStream(), "ERROR", out);
+    final StreamGobbler outputGobbler = new StreamGobbler(process.getInputStream(), "OUTPUT", out);
     outputGobbler.start();
     errorGobbler.start();
 
@@ -238,12 +259,12 @@ private static void waitFor(final Process process, final long timeout, final Tim
   private static class StreamGobbler extends Thread {
     InputStream stream;
     String type;
-    boolean print;
+    PrintStream out;
 
-    private StreamGobbler(final InputStream stream, final String type, final boolean print) {
+    private StreamGobbler(final InputStream stream, final String type, final PrintStream out) {
       this.stream = stream;
       this.type = type;
-      this.print = print;
+      this.out = out;
     }
 
     @Override
@@ -252,8 +273,8 @@ public void run() {
         final BufferedReader reader = new BufferedReader(new InputStreamReader(stream));
         String line = null;
         while ((line = reader.readLine()) != null) {
-          if (print) {
-            System.out.println(type + "> " + line);
+          if (null != out) {
+            out.println(type + "> " + line);
           }
         }
       } catch (final IOException e) {

@@ -0,0 +1,18 @@
+package jvmbootstraptest;
+
+public class JmxStartedChecker {
+  public static void main(final String[] args) throws Exception {
+    AgentLoadedChecker.main(args);
+
+    boolean jmxStarted = false;
+    for (Thread t : Thread.getAllStackTraces().keySet()) {
+      if ("dd-jmx-collector".equals(t.getName())) {
+        jmxStarted = true;
+      }
+    }
+
+    if (!jmxStarted) {
+      throw new IllegalStateException("JMXFetch did not start");
+    }
+  }
+}