Fix rmi socket connection timeout

.gitignore

@@ -18,3 +18,4 @@ target/*
 
 # jenv
 .java_version
+.factorypath

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

@@ -1,5 +1,7 @@
 package org.datadog.jmxfetch;
 
+import static java.rmi.server.RMISocketFactory.setSocketFactory;
+
 import lombok.extern.slf4j.Slf4j;
 
 import java.io.IOException;
@@ -27,6 +29,8 @@ public class RemoteConnection extends Connection {
  private static final String KEY_STORE_PASSWORD_KEY = "key_store_password";
  private static final String DEFAULT_RMI_RESPONSE_TIMEOUT =
  "15000"; // Match the collection period default
+ // The default behaviour is to wait indefinitely, we change that and only wait for 15sec.
+ private static final Integer DEFAULT_RMI_CONNECTION_TIMEOUT = Integer.valueOf(15000);
 
  /** RemoteConnection constructor for specified remote connection parameters. */
  public RemoteConnection(Map<String, Object> connectionParams) throws IOException {
@@ -89,6 +93,18 @@ public RemoteConnection(Map<String, Object> connectionParams) throws IOException
  // Set an RMI timeout so we don't get stuck waiting for a bean to report a value
  System.setProperty("sun.rmi.transport.tcp.responseTimeout", rmiTimeout);
 
+ Integer rmiConnectionTimeout = DEFAULT_RMI_CONNECTION_TIMEOUT;
+ try {
+ rmiConnectionTimeout = (Integer) connectionParams.get("rmi_connection_timeout");
+ } catch (final ClassCastException e) {
+ rmiConnectionTimeout = 
+ Integer.parseInt((String) connectionParams.get("rmi_connection_timeout"));
+ }
+ if (rmiConnectionTimeout == null) {
+ rmiConnectionTimeout = DEFAULT_RMI_CONNECTION_TIMEOUT;
+ }
+ setSocketFactory(new SocketFactory(rmiConnectionTimeout));
+
  createConnection();
  }
 

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

@@ -0,0 +1,27 @@
+package org.datadog.jmxfetch;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.net.ServerSocket;
+import java.net.Socket;
+import java.rmi.server.RMISocketFactory;
+
+public class SocketFactory extends RMISocketFactory {
+ private final int timeout;
+
+ SocketFactory(final int timeout) {
+ this.timeout = timeout;
+ }
+
+ @Override
+ public Socket createSocket(final String host, final int port) throws IOException {
+ final Socket socket = new Socket();
+ socket.connect(new InetSocketAddress(host, port), timeout);
+ return socket;
+ }
+
+ @Override
+ public ServerSocket createServerSocket(final int port) throws IOException {
+ return new ServerSocket(port);
+ }
+}