From 220239a4c330e55641e2b3732b06adcbf3d53e76 Mon Sep 17 00:00:00 2001 From: Stuart McCulloch Date: Mon, 24 Jun 2024 09:41:36 +0100 Subject: [PATCH] Support custom ConnectionFactory --- .../java/org/datadog/jmxfetch/AppConfig.java | 9 +++++ .../datadog/jmxfetch/ConnectionFactory.java | 35 +--------------- .../jmxfetch/DefaultConnectionFactory.java | 40 +++++++++++++++++++ .../java/org/datadog/jmxfetch/Instance.java | 4 +- 4 files changed, 53 insertions(+), 35 deletions(-) create mode 100644 src/main/java/org/datadog/jmxfetch/DefaultConnectionFactory.java diff --git a/src/main/java/org/datadog/jmxfetch/AppConfig.java b/src/main/java/org/datadog/jmxfetch/AppConfig.java index c76a265d1..2e2daeb93 100644 --- a/src/main/java/org/datadog/jmxfetch/AppConfig.java +++ b/src/main/java/org/datadog/jmxfetch/AppConfig.java @@ -295,6 +295,11 @@ public class AppConfig { */ private ServiceNameProvider serviceNameProvider; + /** + * Controls how JMX connections are created. + */ + private ConnectionFactory connectionFactory; + @Builder.Default private Status status = new Status(); @@ -475,6 +480,10 @@ public ServiceNameProvider getServiceNameProvider() { return serviceNameProvider; } + public ConnectionFactory getConnectionFactory() { + return connectionFactory; + } + /** * @return Whether or not internal threads will be run as daemon. */ diff --git a/src/main/java/org/datadog/jmxfetch/ConnectionFactory.java b/src/main/java/org/datadog/jmxfetch/ConnectionFactory.java index 61b7e32db..39d3c1874 100644 --- a/src/main/java/org/datadog/jmxfetch/ConnectionFactory.java +++ b/src/main/java/org/datadog/jmxfetch/ConnectionFactory.java @@ -1,40 +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 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 { - 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 connectionParams) throws IOException; } diff --git a/src/main/java/org/datadog/jmxfetch/DefaultConnectionFactory.java b/src/main/java/org/datadog/jmxfetch/DefaultConnectionFactory.java new file mode 100644 index 000000000..b3f9b82c8 --- /dev/null +++ b/src/main/java/org/datadog/jmxfetch/DefaultConnectionFactory.java @@ -0,0 +1,40 @@ +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 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 { + 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); + } +} diff --git a/src/main/java/org/datadog/jmxfetch/Instance.java b/src/main/java/org/datadog/jmxfetch/Instance.java index e924284b3..dfb752b5d 100644 --- a/src/main/java/org/datadog/jmxfetch/Instance.java +++ b/src/main/java/org/datadog/jmxfetch/Instance.java @@ -403,11 +403,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; }