Merge branch '66-connection-timeout'

nebhale · nebhale · commit 95bb6ea56776 · 2019-01-31T10:23:49.000-08:00
diff --git a/src/main/java/io/r2dbc/postgresql/PostgresqlConnectionConfiguration.java b/src/main/java/io/r2dbc/postgresql/PostgresqlConnectionConfiguration.java
@@ -20,6 +20,8 @@
 import io.r2dbc.postgresql.util.Assert;
 import reactor.util.annotation.Nullable;
 
+import java.time.Duration;
+
 /**
  * Connection configuration information for connecting to a PostgreSQL database.
  */
@@ -32,6 +34,8 @@ public final class PostgresqlConnectionConfiguration {
 
     private final String applicationName;
 
+    private final Duration connectTimeout;
+
     private final String database;
 
     private final String host;
@@ -44,8 +48,11 @@ public final class PostgresqlConnectionConfiguration {
 
     private final String username;
 
-    private PostgresqlConnectionConfiguration(String applicationName, @Nullable String database, String host, String password, int port, @Nullable String schema, String username) {
+    private PostgresqlConnectionConfiguration(String applicationName, @Nullable Duration connectTimeout, @Nullable String database, String host, String password, int port, @Nullable String schema,
+                                              String username) {
+
         this.applicationName = Assert.requireNonNull(applicationName, "applicationName must not be null");
+        this.connectTimeout = connectTimeout;
         this.database = database;
         this.host = Assert.requireNonNull(host, "host must not be null");
         this.password = Assert.requireNonNull(password, "password must not be null");
@@ -65,8 +72,9 @@ public static Builder builder() {
 
     @Override
     public String toString() {
-        return "PostgresConnectionConfiguration{" +
+        return "PostgresqlConnectionConfiguration{" +
             "applicationName='" + this.applicationName + '\'' +
+            ", connectTimeout=" + this.connectTimeout +
             ", database='" + this.database + '\'' +
             ", host='" + this.host + '\'' +
             ", password='" + this.password + '\'' +
@@ -80,6 +88,11 @@ String getApplicationName() {
         return this.applicationName;
     }
 
+    @Nullable
+    Duration getConnectTimeout() {
+        return this.connectTimeout;
+    }
+
     @Nullable
     String getDatabase() {
         return this.database;
@@ -115,6 +128,8 @@ public static final class Builder {
 
         private String applicationName = "r2dbc-postgresql";
 
+        private Duration connectTimeout;
+
         private String database;
 
         private String host;
@@ -148,7 +163,12 @@ public Builder applicationName(String applicationName) {
          * @return a configured {@link PostgresqlConnectionConfiguration}
          */
         public PostgresqlConnectionConfiguration build() {
-            return new PostgresqlConnectionConfiguration(this.applicationName, this.database, this.host, this.password, this.port, this.schema, this.username);
+            return new PostgresqlConnectionConfiguration(this.applicationName, this.connectTimeout, this.database, this.host, this.password, this.port, this.schema, this.username);
+        }
+
+        public Builder connectTimeout(@Nullable Duration connectTimeout) {
+            this.connectTimeout = connectTimeout;
+            return this;
         }
 
         /**
@@ -212,6 +232,7 @@ public Builder schema(@Nullable String schema) {
         public String toString() {
             return "Builder{" +
                 "applicationName='" + this.applicationName + '\'' +
+                ", connectTimeout='" + this.connectTimeout + '\'' +
                 ", database='" + this.database + '\'' +
                 ", host='" + this.host + '\'' +
                 ", password='" + this.password + '\'' +
diff --git a/src/main/java/io/r2dbc/postgresql/PostgresqlConnectionFactory.java b/src/main/java/io/r2dbc/postgresql/PostgresqlConnectionFactory.java
@@ -47,7 +47,7 @@ public PostgresqlConnectionFactory(PostgresqlConnectionConfiguration configurati
         this(Mono.defer(() -> {
             Assert.requireNonNull(configuration, "configuration must not be null");
 
-            return ReactorNettyClient.connect(configuration.getHost(), configuration.getPort()).cast(Client.class);
+            return ReactorNettyClient.connect(configuration.getHost(), configuration.getPort(), configuration.getConnectTimeout()).cast(Client.class);
         }), configuration);
     }
 
diff --git a/src/main/java/io/r2dbc/postgresql/PostgresqlConnectionFactoryProvider.java b/src/main/java/io/r2dbc/postgresql/PostgresqlConnectionFactoryProvider.java
@@ -21,6 +21,7 @@
 import io.r2dbc.spi.ConnectionFactoryProvider;
 import io.r2dbc.spi.Option;
 
+import static io.r2dbc.spi.ConnectionFactoryOptions.CONNECT_TIMEOUT;
 import static io.r2dbc.spi.ConnectionFactoryOptions.DATABASE;
 import static io.r2dbc.spi.ConnectionFactoryOptions.HOST;
 import static io.r2dbc.spi.ConnectionFactoryOptions.PASSWORD;
@@ -58,6 +59,7 @@ public PostgresqlConnectionFactory create(ConnectionFactoryOptions connectionFac
             builder.applicationName(applicationName);
         }
 
+        builder.connectTimeout(connectionFactoryOptions.getValue(CONNECT_TIMEOUT));
         builder.database(connectionFactoryOptions.getValue(DATABASE));
         builder.host(connectionFactoryOptions.getRequiredValue(HOST));
         builder.password(connectionFactoryOptions.getRequiredValue(PASSWORD).toString());
diff --git a/src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient.java b/src/main/java/io/r2dbc/postgresql/client/ReactorNettyClient.java
@@ -19,6 +19,7 @@
 import io.netty.buffer.ByteBufAllocator;
 import io.netty.channel.ChannelDuplexHandler;
 import io.netty.channel.ChannelHandlerContext;
+import io.netty.channel.ChannelOption;
 import io.r2dbc.postgresql.message.backend.BackendKeyData;
 import io.r2dbc.postgresql.message.backend.BackendMessage;
 import io.r2dbc.postgresql.message.backend.BackendMessageDecoder;
@@ -41,8 +42,10 @@
 import reactor.netty.Connection;
 import reactor.netty.resources.ConnectionProvider;
 import reactor.netty.tcp.TcpClient;
+import reactor.util.annotation.Nullable;
 import reactor.util.concurrent.Queues;
 
+import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
 import java.util.Queue;
@@ -168,7 +171,21 @@ private ReactorNettyClient(Connection connection) {
     public static Mono<ReactorNettyClient> connect(String host, int port) {
         Assert.requireNonNull(host, "host must not be null");
 
-        return connect(ConnectionProvider.newConnection(), host, port);
+        return connect(ConnectionProvider.newConnection(), host, port, null);
+    }
+
+    /**
+     * Creates a new frame processor connected to a given host.
+     *
+     * @param host           the host to connect to
+     * @param port           the port to connect to
+     * @param connectTimeout connect timeout
+     * @throws IllegalArgumentException if {@code host} is {@code null}
+     */
+    public static Mono<ReactorNettyClient> connect(String host, int port, @Nullable Duration connectTimeout) {
+        Assert.requireNonNull(host, "host must not be null");
+
+        return connect(ConnectionProvider.newConnection(), host, port, connectTimeout);
     }
 
     /**
@@ -177,15 +194,21 @@ public static Mono<ReactorNettyClient> connect(String host, int port) {
      * @param connectionProvider the connection provider resources
      * @param host               the host to connect to
      * @param port               the port to connect to
+     * @param connectTimeout     connect timeout
      * @throws IllegalArgumentException if {@code host} is {@code null}
      */
-    public static Mono<ReactorNettyClient> connect(ConnectionProvider connectionProvider, String host, int port) {
+    public static Mono<ReactorNettyClient> connect(ConnectionProvider connectionProvider, String host, int port, @Nullable Duration connectTimeout) {
         Assert.requireNonNull(connectionProvider, "connectionProvider must not be null");
         Assert.requireNonNull(host, "host must not be null");
 
-        Mono<? extends Connection> connection = TcpClient.create(connectionProvider)
-            .host(host).port(port)
-            .connect();
+        TcpClient tcpClient = TcpClient.create(connectionProvider)
+            .host(host).port(port);
+
+        if (connectTimeout != null) {
+            tcpClient = tcpClient.option(ChannelOption.CONNECT_TIMEOUT_MILLIS, Math.toIntExact(connectTimeout.toMillis()));
+        }
+
+        Mono<? extends Connection> connection = tcpClient.connect();
 
         return connection.map(ReactorNettyClient::new);
     }
diff --git a/src/test/java/io/r2dbc/postgresql/PostgresqlConnectionConfigurationTest.java b/src/test/java/io/r2dbc/postgresql/PostgresqlConnectionConfigurationTest.java
@@ -18,6 +18,8 @@
 
 import org.junit.jupiter.api.Test;
 
+import java.time.Duration;
+
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
 
@@ -51,6 +53,7 @@ void builderNoUsername() {
     void configuration() {
         PostgresqlConnectionConfiguration configuration = PostgresqlConnectionConfiguration.builder()
             .applicationName("test-application-name")
+            .connectTimeout(Duration.ofMillis(1000))
             .database("test-database")
             .host("test-host")
             .password("test-password")
@@ -61,6 +64,7 @@ void configuration() {
 
         assertThat(configuration)
             .hasFieldOrPropertyWithValue("applicationName", "test-application-name")
+            .hasFieldOrPropertyWithValue("connectTimeout", Duration.ofMillis(1000))
             .hasFieldOrPropertyWithValue("database", "test-database")
             .hasFieldOrPropertyWithValue("host", "test-host")
             .hasFieldOrPropertyWithValue("password", "test-password")
diff --git a/src/test/java/io/r2dbc/postgresql/client/ReactorNettyClientTest.java b/src/test/java/io/r2dbc/postgresql/client/ReactorNettyClientTest.java
@@ -16,6 +16,7 @@
 
 package io.r2dbc.postgresql.client;
 
+import io.netty.channel.ConnectTimeoutException;
 import io.r2dbc.postgresql.PostgresqlConnectionConfiguration;
 import io.r2dbc.postgresql.PostgresqlConnectionFactory;
 import io.r2dbc.postgresql.PostgresqlServerErrorException;
@@ -35,6 +36,7 @@
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.stream.IntStream;
 
@@ -148,6 +150,24 @@ void parallelExchange() {
             .verifyComplete();
     }
 
+    @Test
+    void timeoutTest() {
+        PostgresqlConnectionFactory postgresqlConnectionFactory = new PostgresqlConnectionFactory(PostgresqlConnectionConfiguration.builder()
+            .host("example.com")
+            .port(81)
+            .username("test")
+            .password("test")
+            .database(SERVER.getDatabase())
+            .applicationName(ReactorNettyClientTest.class.getName())
+            .connectTimeout(Duration.ofMillis(200))
+            .build());
+
+        postgresqlConnectionFactory.create()
+            .as(StepVerifier::create)
+            .expectError(ConnectTimeoutException.class)
+            .verify(Duration.ofMillis(500));
+    }
+
     @Nested
     @TestInstance(TestInstance.Lifecycle.PER_CLASS)
     final class ScramTest {

Original file line number	Diff line number	Diff line change
`@@ -47,7 +47,7 @@ public PostgresqlConnectionFactory(PostgresqlConnectionConfiguration configurati`
`47`	`47`	`this(Mono.defer(() -> {`
`48`	`48`	`Assert.requireNonNull(configuration, "configuration must not be null");`
`49`	`49`
`50`		`- return ReactorNettyClient.connect(configuration.getHost(), configuration.getPort()).cast(Client.class);`
	`50`	`+ return ReactorNettyClient.connect(configuration.getHost(), configuration.getPort(), configuration.getConnectTimeout()).cast(Client.class);`
`51`	`51`	`}), configuration);`
`52`	`52`	`}`
`53`	`53`