Add config options for TCP nodelay and keepalive

mp911de · mp911de · commit 233f16e7ef02 · 2020-07-02T11:59:49.000+02:00
[resolves #296]
diff --git a/README.md b/README.md
@@ -88,6 +88,8 @@ Mono<Connection> connectionMono = Mono.from(connectionFactory.create());
 | `sslCert`         | Path to SSL certificate for TLS authentication in PEM format. _(Optional)_
 | `sslPassword`     | Key password to decrypt SSL key. _(Optional)_
 | `sslHostnameVerifier` | `javax.net.ssl.HostnameVerifier` implementation. _(Optional)_
+| `tcpNoDelay`      | Enabled/disable TCP NoDelay. Disabled by default. _(Optional)_
+| `tcpKeepAlive`    | Enabled/disable TCP KeepAlive. Disabled by default _(Optional)_
 
 **Programmatic Configuration**
 
diff --git a/src/main/java/io/r2dbc/postgresql/PostgresqlConnectionConfiguration.java b/src/main/java/io/r2dbc/postgresql/PostgresqlConnectionConfiguration.java
@@ -30,6 +30,7 @@
 
 import javax.net.ssl.HostnameVerifier;
 import java.io.File;
+import java.net.Socket;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -82,14 +83,16 @@ public final class PostgresqlConnectionConfiguration {
 
     private final SSLConfig sslConfig;
 
+    private final boolean tcpKeepAlive;
+
+    private final boolean tcpNoDelay;
+
     private final int preparedStatementCacheQueries;
 
-    private PostgresqlConnectionConfiguration(String applicationName, boolean autodetectExtensions,
-                                              @Nullable Duration connectTimeout, @Nullable String database, List<Extension> extensions,
-                                              ToIntFunction<String> fetchSize, boolean forceBinary,
-                                              @Nullable String host,
-                                              @Nullable Map<String, String> options, @Nullable CharSequence password, int port, @Nullable String schema, @Nullable String socket, String username,
-                                              SSLConfig sslConfig, int preparedStatementCacheQueries) {
+    private PostgresqlConnectionConfiguration(String applicationName, boolean autodetectExtensions, @Nullable Duration connectTimeout, @Nullable String database, List<Extension> extensions,
+                                              ToIntFunction<String> fetchSize, boolean forceBinary, @Nullable String host, @Nullable Map<String, String> options, @Nullable CharSequence password,
+                                              int port, @Nullable String schema, @Nullable String socket, boolean tcpKeepAlive, boolean tcpNoDelay, String username, SSLConfig sslConfig,
+                                              int preparedStatementCacheQueries) {
         this.applicationName = Assert.requireNonNull(applicationName, "applicationName must not be null");
         this.autodetectExtensions = autodetectExtensions;
         this.connectTimeout = connectTimeout;
@@ -109,6 +112,8 @@ private PostgresqlConnectionConfiguration(String applicationName, boolean autode
         this.socket = socket;
         this.username = Assert.requireNonNull(username, "username must not be null");
         this.sslConfig = sslConfig;
+        this.tcpKeepAlive = tcpKeepAlive;
+        this.tcpNoDelay = tcpNoDelay;
         this.preparedStatementCacheQueries = preparedStatementCacheQueries;
     }
 
@@ -135,6 +140,8 @@ public String toString() {
             ", options='" + this.options + '\'' +
             ", password='" + obfuscate(this.password != null ? this.password.length() : 0) + '\'' +
             ", port=" + this.port +
+            ", tcpKeepAlive=" + this.tcpKeepAlive +
+            ", tcpNoDelay=" + this.tcpNoDelay +
             ", username='" + this.username + '\'' +
             '}';
     }
@@ -222,6 +229,14 @@ boolean isForceBinary() {
         return this.forceBinary;
     }
 
+    boolean isTcpKeepAlive() {
+        return this.tcpKeepAlive;
+    }
+
+    boolean isTcpNoDelay() {
+        return this.tcpNoDelay;
+    }
+
     boolean isUseSocket() {
         return getSocket() != null;
     }
@@ -302,6 +317,10 @@ public static final class Builder {
 
         private Function<SslContextBuilder, SslContextBuilder> sslContextBuilderCustomizer = Function.identity();
 
+        private boolean tcpKeepAlive;
+
+        private boolean tcpNoDelay;
+
         @Nullable
         private String username;
 
@@ -353,8 +372,8 @@ public PostgresqlConnectionConfiguration build() {
             }
 
             return new PostgresqlConnectionConfiguration(this.applicationName, this.autodetectExtensions, this.connectTimeout, this.database, this.extensions, this.fetchSize, this.forceBinary,
-                this.host,
-                this.options, this.password, this.port, this.schema, this.socket, this.username, this.createSslConfig(), this.preparedStatementCacheQueries);
+                this.host, this.options, this.password, this.port, this.schema, this.socket, this.tcpKeepAlive, this.tcpNoDelay, this.username, this.createSslConfig(),
+                this.preparedStatementCacheQueries);
         }
 
         /**
@@ -610,6 +629,32 @@ public Builder sslRootCert(String sslRootCert) {
             return this;
         }
 
+        /**
+         * Configure TCP KeepAlive.
+         *
+         * @param enabled whether to enable TCP KeepAlive
+         * @return this {@link Builder}
+         * @see Socket#setKeepAlive(boolean)
+         * @since 0.8.4
+         */
+        public Builder tcpKeepAlive(boolean enabled) {
+            this.tcpKeepAlive = enabled;
+            return this;
+        }
+
+        /**
+         * Configure TCP NoDelay.
+         *
+         * @param enabled whether to enable TCP NoDelay
+         * @return this {@link Builder}
+         * @see Socket#setTcpNoDelay(boolean)
+         * @since 0.8.4
+         */
+        public Builder tcpNoDelay(boolean enabled) {
+            this.tcpNoDelay = enabled;
+            return this;
+        }
+
         /**
          * Configure the username.
          *
@@ -658,6 +703,8 @@ public String toString() {
                 ", sslCert='" + this.sslCert + '\'' +
                 ", sslKey='" + this.sslKey + '\'' +
                 ", sslHostnameVerifier='" + this.sslHostnameVerifier + '\'' +
+                ", tcpKeepAlive='" + this.tcpKeepAlive + '\'' +
+                ", tcpNoDelay='" + this.tcpNoDelay + '\'' +
                 ", preparedStatementCacheQueries='" + this.preparedStatementCacheQueries + '\'' +
                 '}';
         }
diff --git a/src/main/java/io/r2dbc/postgresql/PostgresqlConnectionFactory.java b/src/main/java/io/r2dbc/postgresql/PostgresqlConnectionFactory.java
@@ -21,6 +21,7 @@
 import io.r2dbc.postgresql.authentication.AuthenticationHandler;
 import io.r2dbc.postgresql.authentication.PasswordAuthenticationHandler;
 import io.r2dbc.postgresql.authentication.SASLAuthenticationHandler;
+import io.r2dbc.postgresql.client.ConnectionSettings;
 import io.r2dbc.postgresql.client.Client;
 import io.r2dbc.postgresql.client.ReactorNettyClient;
 import io.r2dbc.postgresql.client.SSLConfig;
@@ -77,7 +78,9 @@ public final class PostgresqlConnectionFactory implements ConnectionFactory {
     public PostgresqlConnectionFactory(PostgresqlConnectionConfiguration configuration) {
         this.configuration = Assert.requireNonNull(configuration, "configuration must not be null");
         this.endpoint = createSocketAddress(configuration);
-        this.clientFactory = sslConfig -> ReactorNettyClient.connect(ConnectionProvider.newConnection(), this.endpoint, configuration.getConnectTimeout(), sslConfig).cast(Client.class);
+
+        ConnectionSettings options = new ConnectionSettings(configuration.getConnectTimeout(), configuration.isTcpKeepAlive(), configuration.isTcpNoDelay());
+        this.clientFactory = sslConfig -> ReactorNettyClient.connect(ConnectionProvider.newConnection(), this.endpoint, options, sslConfig).cast(Client.class);
         this.extensions = getExtensions(configuration);
     }
 
diff --git a/src/main/java/io/r2dbc/postgresql/PostgresqlConnectionFactoryProvider.java b/src/main/java/io/r2dbc/postgresql/PostgresqlConnectionFactoryProvider.java
@@ -125,6 +125,20 @@ public final class PostgresqlConnectionFactoryProvider implements ConnectionFact
      */
     public static final Option<String> SSL_ROOT_CERT = Option.valueOf("sslRootCert");
 
+    /**
+     * Enable TCP KeepAlive.
+     *
+     * @since 0.8.4
+     */
+    public static final Option<Boolean> TCP_KEEPALIVE = Option.valueOf("tcpKeepAlive");
+
+    /**
+     * Enable TCP NoDelay.
+     *
+     * @since 0.8.4
+     */
+    public static final Option<Boolean> TCP_NODELAY = Option.valueOf("tcpNoDelay");
+
     /**
      * Determine the number of queries that are cached in each connection.
      * The default is {@code -1}, meaning there's no limit. The value of {@code 0} disables the cache. Any other value specifies the cache size.
@@ -195,6 +209,8 @@ private static PostgresqlConnectionConfiguration.Builder fromConnectionFactoryOp
             builder.host(options.getRequiredValue(HOST));
             setupSsl(builder, mapper);
         });
+        mapper.from(TCP_KEEPALIVE).map(OptionMapper::toBoolean).to(builder::tcpKeepAlive);
+        mapper.from(TCP_NODELAY).map(OptionMapper::toBoolean).to(builder::tcpNoDelay);
         builder.username(options.getRequiredValue(USER));
 
         return builder;
diff --git a/src/main/java/io/r2dbc/postgresql/client/ConnectionSettings.java b/src/main/java/io/r2dbc/postgresql/client/ConnectionSettings.java
@@ -0,0 +1,56 @@
+/*
+ * Copyright 2020 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.r2dbc.postgresql.client;
+
+import reactor.util.annotation.Nullable;
+
+import java.time.Duration;
+
+/**
+ * Value object capturing common connection settings.
+ *
+ * @since 0.8.4
+ */
+public final class ConnectionSettings {
+
+    private final @Nullable
+    Duration connectTimeout;
+
+    private final boolean tcpKeepAlive;
+
+    private final boolean tcpNoDelay;
+
+    public ConnectionSettings(@Nullable Duration connectTimeout, boolean tcpKeepAlive, boolean tcpNoDelay) {
+        this.tcpKeepAlive = tcpKeepAlive;
+        this.tcpNoDelay = tcpNoDelay;
+        this.connectTimeout = connectTimeout;
+    }
+
+    @Nullable
+    Duration getConnectTimeout() {
+        return this.connectTimeout;
+    }
+
+    boolean isTcpKeepAlive() {
+        return this.tcpKeepAlive;
+    }
+
+    boolean isTcpNoDelay() {
+        return this.tcpNoDelay;
+    }
+
+}
diff --git a/src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient.java b/src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient.java
@@ -336,30 +336,34 @@ public static Mono<ReactorNettyClient> connect(String host, int port) {
      * @throws IllegalArgumentException if {@code host} is {@code null}
      */
     public static Mono<ReactorNettyClient> connect(String host, int port, @Nullable Duration connectTimeout, SSLConfig sslConfig) {
-        return connect(ConnectionProvider.newConnection(), InetSocketAddress.createUnresolved(host, port), connectTimeout, sslConfig);
+        return connect(ConnectionProvider.newConnection(), InetSocketAddress.createUnresolved(host, port), new ConnectionSettings(connectTimeout, false, false), sslConfig);
     }
 
     /**
-     * Creates a new frame processor connected to a given host.
+     * Creates a new frame processor connected to a given {@link SocketAddress}.
      *
      * @param connectionProvider the connection provider resources
      * @param socketAddress      the socketAddress to connect to
-     * @param connectTimeout     connect timeout
+     * @param connectionSettings channel options
      * @param sslConfig          SSL configuration
      * @throws IllegalArgumentException if {@code host} is {@code null}
      */
-    public static Mono<ReactorNettyClient> connect(ConnectionProvider connectionProvider, SocketAddress socketAddress, @Nullable Duration connectTimeout, SSLConfig sslConfig) {
+    public static Mono<ReactorNettyClient> connect(ConnectionProvider connectionProvider, SocketAddress socketAddress, ConnectionSettings connectionSettings, SSLConfig sslConfig) {
         Assert.requireNonNull(connectionProvider, "connectionProvider must not be null");
         Assert.requireNonNull(socketAddress, "socketAddress must not be null");
+        Assert.requireNonNull(connectionSettings, "channelOptions must not be null");
 
         TcpClient tcpClient = TcpClient.create(connectionProvider).remoteAddress(() -> socketAddress);
 
         if (!(socketAddress instanceof InetSocketAddress)) {
             tcpClient = tcpClient.runOn(new SocketLoopResources(), true);
+        } else {
+            tcpClient = tcpClient.option(ChannelOption.SO_KEEPALIVE, connectionSettings.isTcpKeepAlive());
+            tcpClient = tcpClient.option(ChannelOption.TCP_NODELAY, connectionSettings.isTcpNoDelay());
         }
 
-        if (connectTimeout != null) {
-            tcpClient = tcpClient.option(ChannelOption.CONNECT_TIMEOUT_MILLIS, Math.toIntExact(connectTimeout.toMillis()));
+        if (connectionSettings.getConnectTimeout() != null) {
+            tcpClient = tcpClient.option(ChannelOption.CONNECT_TIMEOUT_MILLIS, Math.toIntExact(connectionSettings.getConnectTimeout().toMillis()));
         }
 
         return tcpClient.connect().flatMap(it -> {
diff --git a/src/test/java/io/r2dbc/postgresql/PostgresqlConnectionFactoryProviderUnitTests.java b/src/test/java/io/r2dbc/postgresql/PostgresqlConnectionFactoryProviderUnitTests.java
@@ -35,6 +35,8 @@
 import static io.r2dbc.postgresql.PostgresqlConnectionFactoryProvider.SOCKET;
 import static io.r2dbc.postgresql.PostgresqlConnectionFactoryProvider.SSL_CONTEXT_BUILDER_CUSTOMIZER;
 import static io.r2dbc.postgresql.PostgresqlConnectionFactoryProvider.SSL_MODE;
+import static io.r2dbc.postgresql.PostgresqlConnectionFactoryProvider.TCP_KEEPALIVE;
+import static io.r2dbc.postgresql.PostgresqlConnectionFactoryProvider.TCP_NODELAY;
 import static io.r2dbc.spi.ConnectionFactoryOptions.DRIVER;
 import static io.r2dbc.spi.ConnectionFactoryOptions.HOST;
 import static io.r2dbc.spi.ConnectionFactoryOptions.PASSWORD;
@@ -308,6 +310,34 @@ void shouldApplySslContextBuilderCustomizer() {
         assertThatIllegalStateException().isThrownBy(() -> factory.getConfiguration().getSslConfig().getSslProvider().get()).withMessageContaining("Works!");
     }
 
+    @Test
+    void shouldConfigureTcpKeepAlive() {
+
+        PostgresqlConnectionFactory factory = this.provider.create(builder()
+            .option(DRIVER, POSTGRESQL_DRIVER)
+            .option(HOST, "test-host")
+            .option(PASSWORD, "test-password")
+            .option(USER, "test-user")
+            .option(TCP_KEEPALIVE, true)
+            .build());
+
+        assertThat(factory.getConfiguration().isTcpKeepAlive()).isTrue();
+    }
+
+    @Test
+    void shouldConfigureTcpNoDelay() {
+
+        PostgresqlConnectionFactory factory = this.provider.create(builder()
+            .option(DRIVER, POSTGRESQL_DRIVER)
+            .option(HOST, "test-host")
+            .option(PASSWORD, "test-password")
+            .option(USER, "test-user")
+            .option(TCP_NODELAY, true)
+            .build());
+
+        assertThat(factory.getConfiguration().isTcpNoDelay()).isTrue();
+    }
+
     @Test
     void shouldConnectUsingUnixDomainSocket() {
 
