Skip to content

Commit ca17f11

Browse files
mccullsmcculls
authored
Merge pull request #1958 from DataDog/mcculls/supportCustomManagementBuilders
Delay starting JMXFetch when there's a custom JMX builder
2 parents f06e14b + f919ffd commit ca17f11

File tree

4 files changed

+245
-1
lines changed

4 files changed

+245
-1
lines changed

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

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package datadog.trace.bootstrap;
22

33
import java.lang.instrument.Instrumentation;
4+
import java.lang.management.ManagementFactory;
45
import java.lang.reflect.Constructor;
56
import java.lang.reflect.InvocationTargetException;
67
import java.lang.reflect.Method;
@@ -57,6 +58,8 @@ public static void start(final Instrumentation inst, final URL bootstrapURL) {
5758

5859
final boolean appUsingCustomLogManager = isAppUsingCustomLogManager();
5960

61+
final boolean appUsingCustomJMXBuilder = isAppUsingCustomJMXBuilder();
62+
6063
/*
6164
* java.util.logging.LogManager maintains a final static LogManager, which is created during class initialization.
6265
*
@@ -69,8 +72,15 @@ public static void start(final Instrumentation inst, final URL bootstrapURL) {
6972
*
7073
* Once we see the LogManager class loading, it's safe to start jmxfetch because the application is already setting
7174
* the global log manager and jmxfetch won't be able to touch it due to classloader locking.
75+
*
76+
* Likewise if a custom JMX builder is configured which is not on the system classpath then we delay starting
77+
* JMXFetch until we detect the custom MBeanServerBuilder is being used. This takes precedence over the custom
78+
* log manager check because any custom log manager will be installed before any custom MBeanServerBuilder.
7279
*/
73-
if (appUsingCustomLogManager) {
80+
if (appUsingCustomJMXBuilder) {
81+
log.debug("Custom JMX builder detected. Delaying JMXFetch initialization.");
82+
registerMBeanServerBuilderCallback(new StartJmxCallback(bootstrapURL));
83+
} else if (appUsingCustomLogManager) {
7484
log.debug("Custom logger detected. Delaying JMXFetch initialization.");
7585
registerLogManagerCallback(new StartJmxCallback(bootstrapURL));
7686
} else {
@@ -114,6 +124,18 @@ private static void registerLogManagerCallback(final ClassLoadCallBack callback)
114124
}
115125
}
116126

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

119141
final URL bootstrapURL;
@@ -260,6 +282,9 @@ private static synchronized void installDatadogTracer() {
260282
}
261283

262284
private static synchronized void startJmx(final URL bootstrapURL) {
285+
// load core JMX using the inherited thread-context-classloader
286+
ManagementFactory.getPlatformMBeanServer();
287+
263288
startJmxFetch(bootstrapURL);
264289

265290
if (AGENT_CLASSLOADER == null) {
@@ -508,6 +533,41 @@ private static boolean isAppUsingCustomLogManager() {
508533
return false;
509534
}
510535

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

‎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", "-Ddd.profiling.enabled=true", "-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", "-Ddd.profiling.enabled=true", "-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", "-Ddd.profiling.enabled=true", "-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", "-Ddd.profiling.enabled=true", "-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", "-Ddd.profiling.enabled=true", "-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)