Delay jmxfetch init when custom log manager is used by client's app

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java

@@ -10,7 +10,9 @@ import static net.bytebuddy.matcher.ElementMatchers.not;
 
 import datadog.trace.bootstrap.WeakMap;
 import java.lang.instrument.Instrumentation;
+import java.util.Map;
 import java.util.ServiceLoader;
+import java.util.concurrent.ConcurrentHashMap;
 import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.agent.builder.ResettableClassFileTransformer;
@@ -20,6 +22,8 @@ import net.bytebuddy.utility.JavaModule;
 
 @Slf4j
 public class AgentInstaller {
+  private static final Map<String, Runnable> classLoadCallbacks = new ConcurrentHashMap<>();
+
   static {
     // WeakMap is used by other classes below, so we need to register the provider first.
     registerWeakMapProvider();
@@ -54,6 +58,7 @@ public class AgentInstaller {
             .with(AgentBuilder.DescriptionStrategy.Default.POOL_ONLY)
             .with(POOL_STRATEGY)
             .with(new LoggingListener())
+            .with(new ClassLoadListener())
             .with(LOCATION_STRATEGY)
             // FIXME: we cannot enable it yet due to BB/JVM bug, see
             // https://github.com/raphw/byte-buddy/issues/558
@@ -173,5 +178,53 @@ public class AgentInstaller {
         final boolean loaded) {}
   }
 
+  /**
+   * Register a callback to run when a class is loading.
+   *
+   * <p>Caveats: 1: This callback will be invoked by a jvm class transformer. 2: Classes filtered
+   * out by {@link AgentInstaller}'s skip list will not be matched.
+   *
+   * @param className name of the class to match against
+   * @param classLoadCallback runnable to invoke when class name matches
+   */
+  public static void registerClassLoadCallback(
+      final String className, final Runnable classLoadCallback) {
+    classLoadCallbacks.put(className, classLoadCallback);
+  }
+
+  private static class ClassLoadListener implements AgentBuilder.Listener {
+    @Override
+    public void onDiscovery(
+        String typeName, ClassLoader classLoader, JavaModule javaModule, boolean b) {
+      for (Map.Entry<String, Runnable> entry : classLoadCallbacks.entrySet()) {
+        if (entry.getKey().equals(typeName)) {
+          entry.getValue().run();
+        }
+      }
+    }
+
+    @Override
+    public void onTransformation(
+        TypeDescription typeDescription,
+        ClassLoader classLoader,
+        JavaModule javaModule,
+        boolean b,
+        DynamicType dynamicType) {}
+
+    @Override
+    public void onIgnored(
+        TypeDescription typeDescription,
+        ClassLoader classLoader,
+        JavaModule javaModule,
+        boolean b) {}
+
+    @Override
+    public void onError(
+        String s, ClassLoader classLoader, JavaModule javaModule, boolean b, Throwable throwable) {}
+
+    @Override
+    public void onComplete(String s, ClassLoader classLoader, JavaModule javaModule, boolean b) {}
+  }
+
   private AgentInstaller() {}
 }

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

@@ -31,45 +31,55 @@ import java.util.jar.JarFile;
 
 /** Entry point for initializing the agent. */
 public class TracingAgent {
-  private static boolean inited = false;
+  // fields must be managed under class lock
   private static ClassLoader AGENT_CLASSLOADER = null;
   private static ClassLoader JMXFETCH_CLASSLOADER = null;
+  private static File bootstrapJar = null;
+  private static File toolingJar = null;
+  private static File jmxFetchJar = null;
 
   public static void premain(final String agentArgs, final Instrumentation inst) throws Exception {
-    startAgent(agentArgs, inst);
+    agentmain(agentArgs, inst);
   }
 
   public static void agentmain(final String agentArgs, final Instrumentation inst)
       throws Exception {
-    startAgent(agentArgs, inst);
+    startDatadogAgent(agentArgs, inst);
+    if (isAppUsingCustomLogManager()) {
+      // Application is setting a custom log mananger. JMXFetch touches the global log manager and
+      // must not be started until the app has initialized the global log manager.
+      final Class<?> agentInstallerClass =
+          AGENT_CLASSLOADER.loadClass("datadog.trace.agent.tooling.AgentInstaller");
+      final Method registerCallbackMethod =
+          agentInstallerClass.getMethod("registerClassLoadCallback", String.class, Runnable.class);
+      registerCallbackMethod.invoke(
+          null,
+          "java.util.logging.LogManager",
+          new Runnable() {
+            @Override
+            public void run() {
+              try {
+                startJmxFetch();
+              } catch (Exception e) {
+                throw new RuntimeException(e);
+              }
+            }
+          });
+    } else {
+      startJmxFetch();
+    }
   }
 
-  private static synchronized void startAgent(final String agentArgs, final Instrumentation inst)
+  public static synchronized void startDatadogAgent(
-      throws Exception {
+      final String agentArgs, final Instrumentation inst) throws Exception {
-    if (!inited) {
+    initializeJars();
-      final File toolingJar =
+    if (AGENT_CLASSLOADER == null) {
-          extractToTmpFile(
-              TracingAgent.class.getClassLoader(),
-              "agent-tooling-and-instrumentation.jar.zip",
-              "agent-tooling-and-instrumentation.jar");
-      final File jmxFetchJar =
-          extractToTmpFile(
-              TracingAgent.class.getClassLoader(), "agent-jmxfetch.jar.zip", "agent-jmxfetch.jar");
-      final File bootstrapJar =
-          extractToTmpFile(
-              TracingAgent.class.getClassLoader(),
-              "agent-bootstrap.jar.zip",
-              "agent-bootstrap.jar");
-
       // bootstrap jar must be appended before agent classloader is created.
       inst.appendToBootstrapClassLoaderSearch(new JarFile(bootstrapJar));
       final ClassLoader agentClassLoader = createDatadogClassLoader(bootstrapJar, toolingJar);
-      final ClassLoader jmxFetchClassLoader = createDatadogClassLoader(bootstrapJar, jmxFetchJar);
-
       final ClassLoader contextLoader = Thread.currentThread().getContextClassLoader();
       try {
         Thread.currentThread().setContextClassLoader(agentClassLoader);
-
         { // install agent
           final Class<?> agentInstallerClass =
               agentClassLoader.loadClass("datadog.trace.agent.tooling.AgentInstaller");
@@ -86,23 +96,55 @@ public class TracingAgent {
           final Method logVersionInfoMethod = tracerInstallerClass.getMethod("logVersionInfo");
           logVersionInfoMethod.invoke(null);
         }
+        AGENT_CLASSLOADER = agentClassLoader;
+      } finally {
+        Thread.currentThread().setContextClassLoader(contextLoader);
+      }
+    }
+  }
 
+  public static synchronized void startJmxFetch() throws Exception {
+    initializeJars();
+    if (JMXFETCH_CLASSLOADER == null) {
+      final ClassLoader jmxFetchClassLoader = createDatadogClassLoader(bootstrapJar, jmxFetchJar);
+      final ClassLoader contextLoader = Thread.currentThread().getContextClassLoader();
+      try {
         Thread.currentThread().setContextClassLoader(jmxFetchClassLoader);
-        { // install jmxfetch tracer
-          // We would like jmxfetch to be loaded after APM agent to avoid needing to retransform
-          // classes
         final Class<?> jmxFetchAgentClass =
             jmxFetchClassLoader.loadClass("datadog.trace.agent.jmxfetch.JMXFetch");
         final Method jmxFetchInstallerMethod = jmxFetchAgentClass.getMethod("run");
         jmxFetchInstallerMethod.invoke(null);
-        }
-
-        AGENT_CLASSLOADER = agentClassLoader;
         JMXFETCH_CLASSLOADER = jmxFetchClassLoader;
       } finally {
         Thread.currentThread().setContextClassLoader(contextLoader);
       }
-      inited = true;
+    }
+  }
+
+  /**
+   * Extract embeded jars out of the dd-java-agent to a temporary location.
+   *
+   * <p>Has no effect if jars are already extracted.
+   */
+  private static synchronized void initializeJars() throws Exception {
+    if (bootstrapJar == null) {
+      bootstrapJar =
+          extractToTmpFile(
+              TracingAgent.class.getClassLoader(),
+              "agent-bootstrap.jar.zip",
+              "agent-bootstrap.jar");
+    }
+    if (toolingJar == null) {
+      toolingJar =
+          extractToTmpFile(
+              TracingAgent.class.getClassLoader(),
+              "agent-tooling-and-instrumentation.jar.zip",
+              "agent-tooling-and-instrumentation.jar");
+    }
+    if (jmxFetchJar == null) {
+      jmxFetchJar =
+          extractToTmpFile(
+              TracingAgent.class.getClassLoader(), "agent-jmxfetch.jar.zip", "agent-jmxfetch.jar");
     }
   }
 
@@ -186,6 +228,16 @@ public class TracingAgent {
     return (ClassLoader) method.invoke(null);
   }
 
+  private static boolean isAppUsingCustomLogManager() {
+    final String tracerCustomLogManSysprop = "dd.app.customlogmanager";
+    return System.getProperty("java.util.logging.manager") != null
+        || Boolean.parseBoolean(System.getProperty(tracerCustomLogManSysprop))
+        || Boolean.parseBoolean(
+            System.getenv(tracerCustomLogManSysprop.replace('.', '_').toUpperCase()))
+        // Classes known to set a custom log mananger
+        || ClassLoader.getSystemResource("org/jboss/modules/Main.class") != null;
+  }
+
   /**
    * Main entry point.
    *

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

@@ -9,9 +9,10 @@ import java.lang.management.ManagementFactory
 import java.lang.management.RuntimeMXBean
 
 class CustomLogManagerTest extends Specification {
-  // Tests using forked jvm because groovy has already set the global log manager
+  // Run all tests using forked jvm because groovy has already set the global log manager
+
   @Shared
-  private def agentArg
+  private agentArg
 
   def setupSpec() {
     final RuntimeMXBean runtimeMxBean = ManagementFactory.getRuntimeMXBean()
@@ -24,7 +25,7 @@ class CustomLogManagerTest extends Specification {
     assert agentArg != null
   }
 
-  def "jmxfetch starts up in premain when no custom log manager is set"() {
+  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[]
@@ -32,7 +33,7 @@ class CustomLogManagerTest extends Specification {
       , true) == 0
   }
 
-  def "jmxfetch startup is delayed when java.util.logging.manager sysprop is present"() {
+  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.CustomLogManager"] as String[]
@@ -40,10 +41,10 @@ class CustomLogManagerTest extends Specification {
       , true) == 0
   }
 
-  def "jmxfetch startup is delayed when tracer custom log manager setting is present"() {
+  def "jmxfetch startup delayed with tracer custom log manager setting"() {
     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[]
+      , [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

@@ -11,7 +11,7 @@ public class LogManagerSetter {
           "jmxfetch startup must be delayed when log manager system property is present.");
       customAssert(
           LogManager.getLogManager().getClass(),
-          LogManager.class
+          LogManagerSetter.class
               .getClassLoader()
               .loadClass(System.getProperty("java.util.logging.manager")),
           "Javaagent should not prevent setting a custom log manager");
@@ -20,7 +20,7 @@ public class LogManagerSetter {
       System.setProperty("java.util.logging.manager", CustomLogManager.class.getName());
       customAssert(
           LogManager.getLogManager().getClass(),
-          LogManager.class
+          LogManagerSetter.class
               .getClassLoader()
               .loadClass(System.getProperty("java.util.logging.manager")),
           "Javaagent should not prevent setting a custom log manager");

dd-trace-api/src/main/java/datadog/trace/api/Config.java

@@ -49,6 +49,7 @@ public class Config {
   public static final String JMX_FETCH_REFRESH_BEANS_PERIOD = "jmxfetch.refresh-beans-period";
   public static final String JMX_FETCH_STATSD_HOST = "jmxfetch.statsd.host";
   public static final String JMX_FETCH_STATSD_PORT = "jmxfetch.statsd.port";
+  public static final String APP_CUSTOM_LOG_MANAGER = "app.customlogmanager";
 
   public static final String RUNTIME_ID_TAG = "runtime-id";
   public static final String LANGUAGE_TAG_KEY = "language";
@@ -73,6 +74,8 @@ public class Config {
 
   public static final int DEFAULT_JMX_FETCH_STATSD_PORT = 8125;
 
+  private static final boolean DEFAULT_APP_CUSTOM_LOG_MANAGER = false;
+
   /**
    * this is a random UUID that gets generated on JVM start up and is attached to every root span
    * and every JMX metric that is sent out.
@@ -98,6 +101,7 @@ public class Config {
   @Getter private final String jmxFetchStatsdHost;
   @Getter private final Integer jmxFetchStatsdPort;
   @Getter private final boolean logsInjectionEnabled;
+  @Getter private final boolean appCustomLogManager;
 
   // Read order: System Properties -> Env Variables, [-> default value]
   // Visible for testing
@@ -138,6 +142,9 @@ public class Config {
 
     logsInjectionEnabled =
         getBooleanSettingFromEnvironment(LOGS_INJECTION_ENABLED, DEFAULT_LOGS_INJECTION_ENABLED);
+
+    appCustomLogManager =
+        getBooleanSettingFromEnvironment(APP_CUSTOM_LOG_MANAGER, DEFAULT_APP_CUSTOM_LOG_MANAGER);
   }
 
   // Read order: Properties -> Parent
@@ -182,6 +189,9 @@ public class Config {
 
     logsInjectionEnabled =
         getBooleanSettingFromEnvironment(LOGS_INJECTION_ENABLED, DEFAULT_LOGS_INJECTION_ENABLED);
+
+    appCustomLogManager =
+        getBooleanSettingFromEnvironment(APP_CUSTOM_LOG_MANAGER, DEFAULT_APP_CUSTOM_LOG_MANAGER);
   }
 
   public Map<String, String> getMergedSpanTags() {