wip

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/jmx/MBeanServerRegistry.java

@@ -0,0 +1,13 @@
+package datadog.trace.bootstrap.instrumentation.jmx;
+
+import java.util.HashMap;
+import java.util.Map;
+import javax.management.MBeanServer;
+
+public class MBeanServerRegistry {
+  private static final Map<String, MBeanServer> MAP = new HashMap<>();
+
+  public static Map<String, MBeanServer> get() {
+    return MAP;
+  }
+}

dd-java-agent/agent-jmxfetch/build.gradle

@@ -9,17 +9,19 @@ plugins {
   id "com.github.johnrengelman.shadow"
 }
 apply from: "$rootDir/gradle/java.gradle"
-
+repositories {
+  mavenLocal()
+}
 dependencies {
-  api('com.datadoghq:jmxfetch:0.49.2') {
+  api('com.datadoghq:jmxfetch:0.48.0-SNAPSHOT') {
     exclude group: 'org.slf4j', module: 'slf4j-api'
     exclude group: 'org.slf4j', module: 'slf4j-jdk14'
     exclude group: 'com.beust', module: 'jcommander'
     exclude group: 'com.fasterxml.jackson.core', module: 'jackson-core'
     exclude group: 'com.fasterxml.jackson.jr', module: 'jackson-jr-objects'
   }
-  api deps.slf4j
   api project(':internal-api')
+  api project(':dd-java-agent:agent-bootstrap')
 }
 
 shadowJar {

dd-java-agent/agent-jmxfetch/src/main/java/datadog/trace/agent/jmxfetch/AgentConnectionFactory.java

@@ -0,0 +1,38 @@
+package datadog.trace.agent.jmxfetch;
+
+import java.io.IOException;
+import java.util.Map;
+import javax.management.MBeanServer;
+import org.datadog.jmxfetch.Connection;
+import org.datadog.jmxfetch.ConnectionFactory;
+import org.datadog.jmxfetch.DefaultConnectionFactory;
+import org.datadog.jmxfetch.JvmDirectConnection;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class AgentConnectionFactory implements ConnectionFactory {
+  private static final Logger LOGGER = LoggerFactory.getLogger(AgentConnectionFactory.class);
+
+  private final ConnectionFactory defaultConnectionFactory = new DefaultConnectionFactory();
+  private final Map<String, MBeanServer> instanceStore;
+
+  public AgentConnectionFactory(Map<String, MBeanServer> instanceStore) {
+    this.instanceStore = instanceStore;
+  }
+
+  @Override
+  public Connection createConnection(Map<String, Object> map) throws IOException {
+    Object mbeanServerClass = map.get("mbean_server_class");
+    if (mbeanServerClass != null) {
+      MBeanServer mBeanServer = instanceStore.get(mbeanServerClass.toString());
+      if (mBeanServer != null) {
+        return new InitialMBeanServerConnection(mBeanServer);
+      }
+      LOGGER.warn(
+          "Unable to provide a MBean server instance of {}. Falling back to platform default",
+          mbeanServerClass);
+      return new JvmDirectConnection();
+    }
+    return defaultConnectionFactory.createConnection(map);
+  }
+}

dd-java-agent/agent-jmxfetch/src/main/java/datadog/trace/agent/jmxfetch/InitialMBeanServerConnection.java

@@ -0,0 +1,16 @@
+package datadog.trace.agent.jmxfetch;
+
+import java.io.IOException;
+import javax.annotation.Nonnull;
+import javax.management.MBeanServerConnection;
+import org.datadog.jmxfetch.JvmDirectConnection;
+
+public class InitialMBeanServerConnection extends JvmDirectConnection {
+
+  public InitialMBeanServerConnection(@Nonnull final MBeanServerConnection mbs) throws IOException {
+    this.mbs = mbs;
+  }
+
+  @Override
+  protected void createConnection() throws IOException {}
+}

dd-java-agent/agent-jmxfetch/src/main/java/datadog/trace/agent/jmxfetch/JMXFetch.java

@@ -9,6 +9,7 @@
 import datadog.trace.api.StatsDClient;
 import datadog.trace.api.StatsDClientManager;
 import datadog.trace.api.flare.TracerFlare;
+import datadog.trace.bootstrap.instrumentation.jmx.MBeanServerRegistry;
 import de.thetaphi.forbiddenapis.SuppressForbidden;
 import java.io.IOException;
 import java.io.InputStream;
@@ -20,6 +21,7 @@
 import java.util.Scanner;
 import java.util.SortedSet;
 import java.util.TreeSet;
+import javax.management.MBeanServer;
 import org.datadog.jmxfetch.App;
 import org.datadog.jmxfetch.AppConfig;
 import org.slf4j.Logger;
@@ -108,7 +110,8 @@ private static void run(final StatsDClientManager statsDClientManager, final Con
             .initialRefreshBeansPeriod(initialRefreshBeansPeriod)
             .refreshBeansPeriod(refreshBeansPeriod)
             .globalTags(globalTags)
-            .reporter(reporter);
+            .reporter(reporter)
+            .connectionFactory(new AgentConnectionFactory(MBeanServerRegistry.get()));
 
     if (config.isJmxFetchMultipleRuntimeServicesEnabled()) {
       ServiceNameCollectingTraceInterceptor serviceNameProvider =

dd-java-agent/instrumentation/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/management/CustomMBeanBuilderInstrumentation.java

@@ -7,7 +7,7 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
-import datadog.trace.bootstrap.InstanceStore;
+import datadog.trace.bootstrap.instrumentation.jmx.MBeanServerRegistry;
 import javax.management.MBeanServer;
 import net.bytebuddy.asm.Advice;
 
@@ -43,7 +43,7 @@ public String configuredMatchingType() {
   public static class StoreMBeanServerAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void afterCreation(@Advice.Return final MBeanServer mBeanServer) {
-      InstanceStore.of(MBeanServer.class).putIfAbsent("", mBeanServer);
+      MBeanServerRegistry.get().putIfAbsent(mBeanServer.getClass().getName(), mBeanServer);
     }
   }
 }