Expand custom logman tests to handle three different logman cases

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

@@ -2,31 +2,48 @@ package datadog.trace.agent
 
 import datadog.trace.agent.test.IntegrationTestUtils
 import jvmbootstraptest.LogManagerSetter
+import spock.lang.Shared
 import spock.lang.Specification
 
 import java.lang.management.ManagementFactory
 import java.lang.management.RuntimeMXBean
 
 class CustomLogManagerTest extends Specification {
-  /**
-   * Forked jvm test must run in a pure java runtime because groovy sets the global log manager.
-   */
-  def "javaagent setup does not set the global log manager"() {
-    setup:
+  // Tests using forked jvm because groovy has already set the global log manager
+  @Shared
+  private def agentArg
+
+  def setupSpec() {
     final RuntimeMXBean runtimeMxBean = ManagementFactory.getRuntimeMXBean()
-    String agentArg = null
     for (String arg : runtimeMxBean.getInputArguments()) {
       if (arg.startsWith("-javaagent")) {
         agentArg = arg
         break
       }
     }
-    final URL customAgent = IntegrationTestUtils.createJarWithClasses(LogManagerSetter.getName(), LogManagerSetter)
+    assert agentArg != null
+  }
 
+  def "jmxfetch starts up in premain when no custom log manager is set"() {
     expect:
-    agentArg != null
     IntegrationTestUtils.runOnSeparateJvm(LogManagerSetter.getName()
-      , ["-javaagent:" + customAgent.getPath(), agentArg, "-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off"] as String[]
+      , [agentArg, "-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=off"] as String[]
+      , "" as String[]
+      , true) == 0
+  }
+
+  def "jmxfetch startup is delayed when java.util.logging.manager sysprop is present"() {
+    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[]
+      , "" as String[]
+      , true) == 0
+  }
+
+  def "jmxfetch startup is delayed when tracer custom log manager setting is present"() {
+    expect:
+    IntegrationTestUtils.runOnSeparateJvm(LogManagerSetter.getName()
+      , ["-javaagent:" + customAgent.getPath(), agentArg, "-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
   }

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

@@ -1,44 +1,36 @@
 package jvmbootstraptest;
 
-import java.lang.instrument.Instrumentation;
-import java.net.DatagramPacket;
-import java.net.DatagramSocket;
-import java.net.SocketException;
 import java.util.logging.LogManager;
 
 public class LogManagerSetter {
-  private static final DatagramSocket socket;
-  private static final int localPort;
-
-  static {
-    try {
-      socket = new DatagramSocket(0);
-      localPort = socket.getLocalPort();
-    } catch (SocketException e) {
-      throw new RuntimeException(e);
-    }
-  }
-
-  public static void premain(final String agentArgs, final Instrumentation inst) throws Exception {
-    // set jmxfetch port in premain before tracer's premain runs
-    System.setProperty("dd.jmxfetch.statsd.port", Integer.toString(localPort));
-    System.setProperty("dd.jmxfetch.statsd.host", "localhost");
-  }
-
   public static void main(String... args) throws Exception {
-    try {
-      // block until jmxfetch sends data
-      final byte[] buf = new byte[1500];
-      final DatagramPacket packet = new DatagramPacket(buf, buf.length);
-      socket.receive(packet);
-    } finally {
-      socket.close();
+    if (System.getProperty("java.util.logging.manager") != null) {
+      customAssert(
+          isJmxfetchStarted(),
+          false,
+          "jmxfetch startup must be delayed when log manager system property is present.");
+      customAssert(
+          LogManager.getLogManager().getClass(),
+          LogManager.class
+              .getClassLoader()
+              .loadClass(System.getProperty("java.util.logging.manager")),
+          "Javaagent should not prevent setting a custom log manager");
+      customAssert(isJmxfetchStarted(), true, "jmxfetch should start after loading LogManager.");
+    } else if (System.getProperty("dd.app.customlogmanager") != null) {
+      System.setProperty("java.util.logging.manager", CustomLogManager.class.getName());
+      customAssert(
+          LogManager.getLogManager().getClass(),
+          LogManager.class
+              .getClassLoader()
+              .loadClass(System.getProperty("java.util.logging.manager")),
+          "Javaagent should not prevent setting a custom log manager");
+
+    } else {
+      customAssert(
+          isJmxfetchStarted(),
+          true,
+          "jmxfetch should start in premain when no custom log manager is set.");
     }
-    System.setProperty("java.util.logging.manager", CustomLogManager.class.getName());
-    customAssert(
-        LogManager.getLogManager().getClass(),
-        CustomLogManager.class,
-        "Javaagent should not prevent setting a custom log manager");
   }
 
   private static void customAssert(Object got, Object expected, String assertionMessage) {
@@ -47,4 +39,13 @@ public class LogManagerSetter {
           "Assertion failed. Expected <" + expected + "> got <" + got + "> " + assertionMessage);
     }
   }
+
+  private static boolean isJmxfetchStarted() {
+    for (Thread thread : Thread.getAllStackTraces().keySet()) {
+      if ("dd-jmx-collector".equals(thread.getName())) {
+        return true;
+      }
+    }
+    return false;
+  }
 }