Support custom ConnectionFactory (#530)

src/main/java/org/datadog/jmxfetch/AppConfig.java

@@ -306,6 +306,12 @@ public class AppConfig {
      */
     private ServiceNameProvider serviceNameProvider;
 
+    /**
+     * Controls how JMX connections are created.
+     */
+    @Builder.Default
+    private ConnectionFactory connectionFactory = new DefaultConnectionFactory();
+
     @Builder.Default
     private Status status = new Status();
 
@@ -494,6 +500,10 @@ public ServiceNameProvider getServiceNameProvider() {
         return serviceNameProvider;
     }
 
+    public ConnectionFactory getConnectionFactory() {
+        return connectionFactory;
+    }
+
     /**
      * @return Whether or not internal threads will be run as daemon.
      */

src/main/java/org/datadog/jmxfetch/ConnectionFactory.java

@@ -1,42 +1,9 @@
 package org.datadog.jmxfetch;
 
-import static org.datadog.jmxfetch.Instance.isDirectInstance;
-
-import lombok.extern.slf4j.Slf4j;
-
 import java.io.IOException;
 import java.util.Map;
 
-/** Singleton used to create connections to the MBeanServer. */
-@Slf4j
-public class ConnectionFactory {
-    public static final String PROCESS_NAME_REGEX = "process_name_regex";
-
+public interface ConnectionFactory {
     /** Factory method to create connections, both remote and local to the target JVM. */
-    public static Connection createConnection(Map<String, Object> connectionParams)
-            throws IOException {
-        // This is used by dd-java-agent to enable directly connecting to the mbean server.
-        // This works since jmxfetch is being run as a library inside the process being monitored.
-        if (isDirectInstance(connectionParams)) {
-            log.info("Connecting to JMX directly on the JVM");
-            return new JvmDirectConnection();
-        }
-
-        if (connectionParams.get(PROCESS_NAME_REGEX) != null) {
-            try {
-                // AttachNotSupportedException is accessible in java 7 and 8 through tools.jar 
-                // and java 9+ by default
-                Class.forName("com.sun.tools.attach.AttachNotSupportedException");
-            } catch (ClassNotFoundException e) {
-                throw new IOException(
-                        "Unable to find tools.jar."
-                                + " Are you using a JDK and did you set the path to tools.jar ?");
-            }
-            log.info("Connecting using Attach API");
-            return new AttachApiConnection(connectionParams);
-        }
-
-        log.info("Connecting using JMX Remote");
-        return new RemoteConnection(connectionParams);
-    }
+    Connection createConnection(Map<String, Object> connectionParams) throws IOException;
 }

src/main/java/org/datadog/jmxfetch/DefaultConnectionFactory.java

@@ -0,0 +1,42 @@
+package org.datadog.jmxfetch;
+
+import static org.datadog.jmxfetch.Instance.isDirectInstance;
+
+import lombok.extern.slf4j.Slf4j;
+
+import java.io.IOException;
+import java.util.Map;
+
+/** Singleton used to create connections to the MBeanServer. */
+@Slf4j
+public class DefaultConnectionFactory implements ConnectionFactory {
+    public static final String PROCESS_NAME_REGEX = "process_name_regex";
+
+    /** Factory method to create connections, both remote and local to the target JVM. */
+    public Connection createConnection(Map<String, Object> connectionParams)
+            throws IOException {
+        // This is used by dd-java-agent to enable directly connecting to the mbean server.
+        // This works since jmxfetch is being run as a library inside the process being monitored.
+        if (isDirectInstance(connectionParams)) {
+            log.info("Connecting to JMX directly on the JVM");
+            return new JvmDirectConnection();
+        }
+
+        if (connectionParams.get(PROCESS_NAME_REGEX) != null) {
+            try {
+                // AttachNotSupportedException is accessible in java 7 and 8 through tools.jar 
+                // and java 9+ by default
+                Class.forName("com.sun.tools.attach.AttachNotSupportedException");
+            } catch (ClassNotFoundException e) {
+                throw new IOException(
+                        "Unable to find tools.jar."
+                                + " Are you using a JDK and did you set the path to tools.jar ?");
+            }
+            log.info("Connecting using Attach API");
+            return new AttachApiConnection(connectionParams);
+        }
+
+        log.info("Connecting using JMX Remote");
+        return new RemoteConnection(connectionParams);
+    }
+}

src/main/java/org/datadog/jmxfetch/Instance.java

@@ -424,11 +424,11 @@ public Connection getConnection(
             log.info(
                     "Connection closed or does not exist. "
                     + "Attempting to create a new connection...");
-            return ConnectionFactory.createConnection(connectionParams);
+            return appConfig.getConnectionFactory().createConnection(connectionParams);
         } else if (forceNewConnection) {
             log.info("Forcing a new connection, attempting to create...");
             connection.closeConnector();
-            return ConnectionFactory.createConnection(connectionParams);
+            return appConfig.getConnectionFactory().createConnection(connectionParams);
         }
         return connection;
     }