Skip to content

Commit 1e07b3d

Browse files
mccullsmcculls
committed
Delay starting JMXFetch when there's a custom JMX builder that's not on the system classpath (ie. it will be provisioned later)
1 parent 20473c0 commit 1e07b3d

4 files changed

Lines changed: 241 additions & 1 deletion

File tree

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java‎

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,8 @@ public static void start(final Instrumentation inst, final URL bootstrapURL) {
5757

5858
final boolean appUsingCustomLogManager = isAppUsingCustomLogManager();
5959

60+
final boolean appUsingCustomJMXBuilder = isAppUsingCustomJMXBuilder();
61+
6062
/*
6163
* java.util.logging.LogManager maintains a final static LogManager, which is created during class initialization.
6264
*
@@ -69,8 +71,15 @@ public static void start(final Instrumentation inst, final URL bootstrapURL) {
6971
*
7072
* Once we see the LogManager class loading, it's safe to start jmxfetch because the application is already setting
7173
* the global log manager and jmxfetch won't be able to touch it due to classloader locking.
74+
*
75+
* Likewise if a custom JMX builder is configured which is not on the system classpath then we delay starting
76+
* JMXFetch until we detect the custom MBeanServerBuilder is being used. This takes precedence over the custom
77+
* log manager check because any custom log manager will be installed before any custom MBeanServerBuilder.
7278
*/
73-
if (appUsingCustomLogManager) {
79+
if (appUsingCustomJMXBuilder) {
80+
log.debug("Custom JMX builder detected. Delaying JMXFetch initialization.");
81+
registerMBeanServerBuilderCallback(new StartJmxCallback(bootstrapURL));
82+
} else if (appUsingCustomLogManager) {
7483
log.debug("Custom logger detected. Delaying JMXFetch initialization.");
7584
registerLogManagerCallback(new StartJmxCallback(bootstrapURL));
7685
} else {
@@ -114,6 +123,18 @@ private static void registerLogManagerCallback(final ClassLoadCallBack callback)
114123
}
115124
}
116125

126+
private static void registerMBeanServerBuilderCallback(final ClassLoadCallBack callback) {
127+
try {
128+
final Class<?> agentInstallerClass =
129+
AGENT_CLASSLOADER.loadClass("datadog.trace.agent.tooling.AgentInstaller");
130+
final Method registerCallbackMethod =
131+
agentInstallerClass.getMethod("registerClassLoadCallback", String.class, Runnable.class);
132+
registerCallbackMethod.invoke(null, "javax.management.MBeanServerBuilder", callback);
133+
} catch (final Exception ex) {
134+
log.error("Error registering callback for " + callback.getName(), ex);
135+
}
136+
}
137+
117138
protected abstract static class ClassLoadCallBack implements Runnable {
118139

119140
final URL bootstrapURL;
@@ -508,6 +529,41 @@ private static boolean isAppUsingCustomLogManager() {
508529
return false;
509530
}
510531

532+
/**
533+
* Search for java or datadog-tracer sysprops which indicate that a custom JMX builder will be
534+
* used.
535+
*
536+
* @return true if we detect a custom JMX builder being used.
537+
*/
538+
private static boolean isAppUsingCustomJMXBuilder() {
539+
final String tracerCustomJMXBuilderSysprop = "dd.app.customjmxbuilder";
540+
final String customJMXBuilderProp = System.getProperty(tracerCustomJMXBuilderSysprop);
541+
final String customJMXBuilderEnv =
542+
System.getenv(tracerCustomJMXBuilderSysprop.replace('.', '_').toUpperCase());
543+
544+
if (customJMXBuilderProp != null || customJMXBuilderEnv != null) {
545+
log.debug("Prop - customjmxbuilder: " + customJMXBuilderProp);
546+
log.debug("Env - customjmxbuilder: " + customJMXBuilderEnv);
547+
// Allow setting to skip these automatic checks:
548+
return Boolean.parseBoolean(customJMXBuilderProp)
549+
|| Boolean.parseBoolean(customJMXBuilderEnv);
550+
}
551+
552+
final String jmxBuilderProp = System.getProperty("javax.management.builder.initial");
553+
if (jmxBuilderProp != null) {
554+
final boolean onSysClasspath =
555+
ClassLoader.getSystemResource(jmxBuilderProp.replaceAll("\\.", "/") + ".class") != null;
556+
log.debug("Prop - javax.management.builder.initial: " + jmxBuilderProp);
557+
log.debug("javax.management.builder.initial on system classpath: " + onSysClasspath);
558+
// Some applications set javax.management.builder.initial but never actually initialize JMX.
559+
// Check to see if the configured JMX builder is on the system classpath.
560+
// If so, it should be safe to initialize jmxfetch which will setup JMX.
561+
return !onSysClasspath;
562+
}
563+
564+
return false;
565+
}
566+
511567
private static boolean isJavaBefore9() {
512568
return System.getProperty("java.version").startsWith("1.");
513569
}

‎dd-java-agent/src/test/groovy/datadog/trace/agent/CustomMBeanServerBuilderTest.groovy‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
package datadog.trace.agent
2+
3+
import datadog.trace.agent.test.IntegrationTestUtils
4+
import jvmbootstraptest.MBeanServerBuilderSetter
5+
import spock.lang.Specification
6+
import spock.lang.Timeout
7+
8+
@Timeout(30)
9+
class CustomMBeanServerBuilderTest extends Specification {
10+
11+
private static final String DEFAULT_LOG_LEVEL = "debug"
12+
private static final String API_KEY = "01234567890abcdef123456789ABCDEF"
13+
14+
// Run all tests using forked jvm so we try different JMX settings
15+
def "JMXFetch starts up in premain with no custom MBeanServerBuilder set"() {
16+
expect:
17+
IntegrationTestUtils.runOnSeparateJvm(MBeanServerBuilderSetter.getName()
18+
, ["-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=$DEFAULT_LOG_LEVEL"] as String[]
19+
, "" as String[]
20+
, ["DD_API_KEY": API_KEY]
21+
, true) == 0
22+
}
23+
24+
def "JMXFetch starts up in premain if configured MBeanServerBuilder on system classpath"() {
25+
expect:
26+
IntegrationTestUtils.runOnSeparateJvm(MBeanServerBuilderSetter.getName()
27+
, ["-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=$DEFAULT_LOG_LEVEL", "-Djavax.management.builder.initial=jvmbootstraptest.CustomMBeanServerBuilder"] as String[]
28+
, "" as String[]
29+
, ["DD_API_KEY": API_KEY]
30+
, true) == 0
31+
}
32+
33+
def "JMXFetch startup is delayed with javax.management.builder.initial sysprop"() {
34+
expect:
35+
IntegrationTestUtils.runOnSeparateJvm(MBeanServerBuilderSetter.getName()
36+
, ["-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=$DEFAULT_LOG_LEVEL", "-Djavax.management.builder.initial=jvmbootstraptest.MissingMBeanServerBuilder"] as String[]
37+
, "" as String[]
38+
, ["DD_API_KEY": API_KEY]
39+
, true) == 0
40+
}
41+
42+
def "JMXFetch startup is delayed with tracer custom JMX builder setting"() {
43+
expect:
44+
IntegrationTestUtils.runOnSeparateJvm(MBeanServerBuilderSetter.getName()
45+
, ["-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=$DEFAULT_LOG_LEVEL", "-Ddd.app.customjmxbuilder=true"] as String[]
46+
, "" as String[]
47+
, ["DD_API_KEY": API_KEY]
48+
, true) == 0
49+
}
50+
51+
def "JMXFetch starts up in premain forced by customjmxbuilder=false"() {
52+
expect:
53+
IntegrationTestUtils.runOnSeparateJvm(MBeanServerBuilderSetter.getName()
54+
, ["-Ddd.jmxfetch.enabled=true", "-Ddd.jmxfetch.refresh-beans-period=1", "-Ddatadog.slf4j.simpleLogger.defaultLogLevel=$DEFAULT_LOG_LEVEL", "-Ddd.app.customjmxbuilder=false", "-Djavax.management.builder.initial=jvmbootstraptest.CustomMBeanServerBuilder"] as String[]
55+
, "" as String[]
56+
, ["DD_API_KEY": API_KEY]
57+
, true) == 0
58+
}
59+
}

‎dd-java-agent/src/test/java/jvmbootstraptest/CustomMBeanServerBuilder.java‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
package jvmbootstraptest;
2+
3+
import javax.management.MBeanServer;
4+
import javax.management.MBeanServerBuilder;
5+
import javax.management.MBeanServerDelegate;
6+
import javax.management.ObjectName;
7+
import javax.management.StandardMBean;
8+
9+
public class CustomMBeanServerBuilder extends MBeanServerBuilder {
10+
public interface TestMBean {}
11+
12+
@Override
13+
public MBeanServer newMBeanServer(
14+
final String defaultDomain, final MBeanServer outer, final MBeanServerDelegate delegate) {
15+
final MBeanServer mBeanServer = super.newMBeanServer(defaultDomain, outer, delegate);
16+
try {
17+
mBeanServer.registerMBean(
18+
new StandardMBean(new TestMBean() {}, TestMBean.class),
19+
new ObjectName("test:name=custom"));
20+
} catch (final Exception e) {
21+
throw new IllegalStateException(e);
22+
}
23+
return mBeanServer;
24+
}
25+
}

‎dd-java-agent/src/test/java/jvmbootstraptest/MBeanServerBuilderSetter.java‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
package jvmbootstraptest;
2+
3+
import java.lang.management.ManagementFactory;
4+
import java.util.Objects;
5+
import javax.management.MalformedObjectNameException;
6+
import javax.management.ObjectName;
7+
8+
public class MBeanServerBuilderSetter {
9+
public static void main(final String... args) throws Exception {
10+
if (System.getProperty("dd.app.customjmxbuilder") != null) {
11+
System.out.println("dd.app.customjmxbuilder != null");
12+
13+
if (Boolean.parseBoolean(System.getProperty("dd.app.customjmxbuilder"))) {
14+
System.setProperty(
15+
"javax.management.builder.initial", "jvmbootstraptest.CustomMBeanServerBuilder");
16+
customAssert(
17+
isCustomMBeanRegistered(),
18+
true,
19+
"Javaagent should not prevent setting a custom MBeanServerBuilder");
20+
} else {
21+
customAssert(
22+
isJmxfetchStarted(false),
23+
true,
24+
"jmxfetch should start in premain when customjmxbuilder=false.");
25+
}
26+
} else if (System.getProperty("javax.management.builder.initial") != null) {
27+
System.out.println("javax.management.builder.initial != null");
28+
29+
if (ClassLoader.getSystemResource(
30+
System.getProperty("javax.management.builder.initial").replaceAll("\\.", "/")
31+
+ ".class")
32+
== null) {
33+
customAssert(
34+
isJmxfetchStarted(false),
35+
false,
36+
"jmxfetch startup must be delayed when management builder system property is present.");
37+
// Change back to a valid MBeanServerBuilder.
38+
System.setProperty(
39+
"javax.management.builder.initial", "jvmbootstraptest.CustomMBeanServerBuilder");
40+
customAssert(
41+
isCustomMBeanRegistered(),
42+
true,
43+
"Javaagent should not prevent setting a custom MBeanServerBuilder");
44+
customAssert(
45+
isJmxfetchStarted(true),
46+
true,
47+
"jmxfetch should start after loading MBeanServerBuilder.");
48+
} else {
49+
customAssert(
50+
isJmxfetchStarted(false),
51+
true,
52+
"jmxfetch should start in premain when custom MBeanServerBuilder found on classpath.");
53+
}
54+
} else {
55+
System.out.println("No custom MBeanServerBuilder");
56+
57+
customAssert(
58+
isJmxfetchStarted(false),
59+
true,
60+
"jmxfetch should start in premain when no custom MBeanServerBuilder is set.");
61+
}
62+
}
63+
64+
private static boolean isCustomMBeanRegistered() throws MalformedObjectNameException {
65+
return ManagementFactory.getPlatformMBeanServer()
66+
.isRegistered(new ObjectName("test:name=custom"));
67+
}
68+
69+
private static void customAssert(
70+
final Object got, final Object expected, final String assertionMessage) {
71+
if (!Objects.equals(got, expected)) {
72+
throw new RuntimeException(
73+
"Assertion failed. Expected <" + expected + "> got <" + got + "> " + assertionMessage);
74+
}
75+
}
76+
77+
private static boolean isThreadStarted(final String name, final boolean wait) {
78+
// Wait up to 10 seconds for thread to appear
79+
for (int i = 0; i < 20; i++) {
80+
for (final Thread thread : Thread.getAllStackTraces().keySet()) {
81+
if (name.equals(thread.getName())) {
82+
return true;
83+
}
84+
}
85+
if (!wait) {
86+
break;
87+
}
88+
try {
89+
Thread.sleep(500);
90+
} catch (final InterruptedException e) {
91+
e.printStackTrace();
92+
}
93+
}
94+
return false;
95+
}
96+
97+
private static boolean isJmxfetchStarted(final boolean wait) {
98+
return isThreadStarted("dd-jmx-collector", wait);
99+
}
100+
}

0 commit comments

Comments
 (0)