Merge pull request #285 from DataDog/prognant/fix-RMI-socket-connecti…

…on-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);
+    }
+}