Support Keep alive probes in Java 11+

JAVA-3656

Author: jstewart148

driver-core/src/main/com/mongodb/internal/connection/SocketStreamHelper.java

@@ -19,24 +19,39 @@
 import com.mongodb.MongoInternalException;
 import com.mongodb.connection.SocketSettings;
 import com.mongodb.connection.SslSettings;
+import jdk.net.ExtendedSocketOptions;
 
 import javax.net.ssl.SSLParameters;
 import javax.net.ssl.SSLSocket;
 import java.io.IOException;
+import java.lang.reflect.Method;
 import java.net.InetSocketAddress;
 import java.net.Socket;
+import java.net.SocketOption;
+import java.util.Arrays;
 
 import static com.mongodb.internal.connection.SslHelper.enableHostNameVerification;
 import static com.mongodb.internal.connection.SslHelper.enableSni;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
 final class SocketStreamHelper {
+    // Keep alive options and their values for Java 11+
+    private static final String TCP_KEEPIDLE = "TCP_KEEPIDLE";
+    private static final int TCP_KEEPIDLE_DURATION = 300;
+    private static final String TCP_KEEPCOUNT = "TCP_KEEPCOUNT";
+    private static final int TCP_KEEPCOUNT_LIMIT = 9;
+    private static final String TCP_KEEPINTERVAL = "TCP_KEEPINTERVAL";
+    private static final int TCP_KEEPINTERVAL_DURATION = 10;
 
     static void initialize(final Socket socket, final InetSocketAddress inetSocketAddress, final SocketSettings settings,
                            final SslSettings sslSettings) throws IOException {
         socket.setTcpNoDelay(true);
         socket.setSoTimeout(settings.getReadTimeout(MILLISECONDS));
         socket.setKeepAlive(true);
+
+        // Adding keep alive options for users of Java 11+. These options will be ignored for older Java versions.
+        setExtendedSocketOptions(socket);
+
         if (settings.getReceiveBufferSize() > 0) {
             socket.setReceiveBufferSize(settings.getReceiveBufferSize());
         }
@@ -63,6 +78,22 @@ static void initialize(final Socket socket, final InetSocketAddress inetSocketAd
         socket.connect(inetSocketAddress, settings.getConnectTimeout(MILLISECONDS));
     }
 
+    @SuppressWarnings("unchecked")
+    private static void setExtendedSocketOptions(final Socket socket) {
+        if (Arrays.stream(ExtendedSocketOptions.class.getDeclaredFields()).anyMatch(f -> f.getName().equals(TCP_KEEPCOUNT))) {
+            try {
+                Method setOptionMethod = Socket.class.getMethod("setOption", SocketOption.class, Object.class);
+                setOptionMethod.invoke(socket, ExtendedSocketOptions.class.getDeclaredField(TCP_KEEPCOUNT).get(null),
+                        TCP_KEEPCOUNT_LIMIT);
+                setOptionMethod.invoke(socket, ExtendedSocketOptions.class.getDeclaredField(TCP_KEEPIDLE).get(null),
+                        TCP_KEEPIDLE_DURATION);
+                setOptionMethod.invoke(socket, ExtendedSocketOptions.class.getDeclaredField(TCP_KEEPINTERVAL).get(null),
+                        TCP_KEEPINTERVAL_DURATION);
+            } catch (Throwable t) {
+            }
+        }
+    }
+
     private SocketStreamHelper() {
     }
 }

driver-core/src/test/functional/com/mongodb/internal/connection/SocketStreamHelperSpecification.groovy

@@ -20,13 +20,15 @@ import com.mongodb.ClusterFixture
 import com.mongodb.MongoInternalException
 import com.mongodb.connection.SocketSettings
 import com.mongodb.connection.SslSettings
+import jdk.net.ExtendedSocketOptions
 import spock.lang.IgnoreIf
 import spock.lang.Specification
 
 import javax.net.SocketFactory
 import javax.net.ssl.SNIHostName
 import javax.net.ssl.SSLSocket
 import javax.net.ssl.SSLSocketFactory
+import java.lang.reflect.Method
 
 import static com.mongodb.ClusterFixture.getPrimary
 import static java.util.concurrent.TimeUnit.MILLISECONDS
@@ -49,6 +51,14 @@ class SocketStreamHelperSpecification extends Specification {
         socket.getKeepAlive()
         socket.getSoTimeout() == socketSettings.getReadTimeout(MILLISECONDS)
 
+        // If the Java 11+ extended socket options for keep alive probes are available, check those values.
+        if (Arrays.stream(ExtendedSocketOptions.getDeclaredFields()).anyMatch{ f -> f.getName().equals('TCP_KEEPCOUNT') }) {
+            Method getOptionMethod = Socket.getMethod('getOption', SocketOption);
+            getOptionMethod.invoke(socket, ExtendedSocketOptions.getDeclaredField('TCP_KEEPCOUNT').get(null)) == 9
+            getOptionMethod.invoke(socket, ExtendedSocketOptions.getDeclaredField('TCP_KEEPIDLE').get(null)) == 300
+            getOptionMethod.invoke(socket, ExtendedSocketOptions.getDeclaredField('TCP_KEEPINTERVAL').get(null)) == 10
+        }
+
         cleanup:
         socket?.close()
     }