Merge pull request #294 from lutovich/1.1-config-socket-connect-timeout

Make socket connection timeout configurable

Author: Zhen Li

driver/src/main/java/org/neo4j/driver/internal/ConnectionSettings.java

@@ -29,7 +29,7 @@
  */
 public class ConnectionSettings
 {
-    public static final String DEFAULT_USER_AGENT = format( "neo4j-java/%s", driverVersion() );
+    private static final String DEFAULT_USER_AGENT = format( "neo4j-java/%s", driverVersion() );
 
     /**
      * Extracts the driver version from the driver jar MANIFEST.MF file.
@@ -52,16 +52,18 @@ private static String driverVersion()
 
     private final AuthToken authToken;
     private final String userAgent;
+    private final int timeoutMillis;
 
-    public ConnectionSettings( AuthToken authToken, String userAgent )
+    public ConnectionSettings( AuthToken authToken, String userAgent, int timeoutMillis )
     {
         this.authToken = authToken;
         this.userAgent = userAgent;
+        this.timeoutMillis = timeoutMillis;
     }
 
-    public ConnectionSettings( AuthToken authToken )
+    public ConnectionSettings( AuthToken authToken, int timeoutMillis )
     {
-        this( authToken, DEFAULT_USER_AGENT );
+        this( authToken, DEFAULT_USER_AGENT, timeoutMillis );
     }
 
     public AuthToken authToken()
@@ -74,4 +76,8 @@ public String userAgent()
         return userAgent;
     }
 
+    public int timeoutMillis()
+    {
+        return timeoutMillis;
+    }
 }

driver/src/main/java/org/neo4j/driver/internal/DriverFactory.java

@@ -120,7 +120,7 @@ ConnectionPool createConnectionPool( AuthToken authToken, SecurityPlan securityP
     {
         authToken = authToken == null ? AuthTokens.none() : authToken;
 
-        ConnectionSettings connectionSettings = new ConnectionSettings( authToken );
+        ConnectionSettings connectionSettings = new ConnectionSettings( authToken, config.connectionTimeoutMillis() );
         PoolSettings poolSettings = new PoolSettings( config.maxIdleConnectionPoolSize() );
         Connector connector = new SocketConnector( connectionSettings, securityPlan, config.logging() );
 

driver/src/main/java/org/neo4j/driver/internal/net/ChannelFactory.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright (c) 2002-2016 "Neo Technology,"
+ * Network Engine for Objects in Lund AB [http://neotechnology.com]
+ *
+ * This file is part of Neo4j.
+ *
+ * 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
+ *
+ *     http://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 org.neo4j.driver.internal.net;
+
+import java.io.IOException;
+import java.net.Socket;
+import java.net.SocketTimeoutException;
+import java.net.StandardSocketOptions;
+import java.nio.channels.ByteChannel;
+import java.nio.channels.SocketChannel;
+import java.security.GeneralSecurityException;
+
+import org.neo4j.driver.internal.security.SecurityPlan;
+import org.neo4j.driver.internal.security.TLSSocketChannel;
+import org.neo4j.driver.v1.Logger;
+
+class ChannelFactory
+{
+    static ByteChannel create( BoltServerAddress address, SecurityPlan securityPlan, int timeoutMillis, Logger log )
+            throws IOException, GeneralSecurityException
+    {
+        SocketChannel soChannel = SocketChannel.open();
+        soChannel.setOption( StandardSocketOptions.SO_REUSEADDR, true );
+        soChannel.setOption( StandardSocketOptions.SO_KEEPALIVE, true );
+        connect( soChannel, address, timeoutMillis );
+
+        ByteChannel channel;
+
+        if ( securityPlan.requiresEncryption() )
+        {
+            channel = TLSSocketChannel.create( address, securityPlan, soChannel, log );
+        }
+        else
+        {
+            channel = soChannel;
+        }
+
+        if ( log.isTraceEnabled() )
+        {
+            channel = new LoggingByteChannel( channel, log );
+        }
+
+        return channel;
+    }
+
+    private static void connect( SocketChannel soChannel, BoltServerAddress address, int timeoutMillis )
+            throws IOException
+    {
+        Socket socket = soChannel.socket();
+        try
+        {
+            socket.connect( address.toSocketAddress(), timeoutMillis );
+        }
+        catch ( SocketTimeoutException e )
+        {
+            throw new IOException( "Timeout " + timeoutMillis + "ms expired", e );
+        }
+    }
+}

driver/src/main/java/org/neo4j/driver/internal/net/SocketClient.java

@@ -20,17 +20,14 @@
 
 import java.io.IOException;
 import java.net.ConnectException;
-import java.net.StandardSocketOptions;
 import java.nio.ByteBuffer;
 import java.nio.channels.ByteChannel;
-import java.nio.channels.SocketChannel;
 import java.security.GeneralSecurityException;
 import java.util.Queue;
 
 import org.neo4j.driver.internal.messaging.Message;
 import org.neo4j.driver.internal.messaging.MessageFormat;
 import org.neo4j.driver.internal.security.SecurityPlan;
-import org.neo4j.driver.internal.security.TLSSocketChannel;
 import org.neo4j.driver.internal.util.BytePrinter;
 import org.neo4j.driver.v1.Logger;
 import org.neo4j.driver.v1.exceptions.ClientException;
@@ -49,6 +46,7 @@ public class SocketClient
 
     private final BoltServerAddress address;
     private final SecurityPlan securityPlan;
+    private final int timeoutMillis;
     private final Logger logger;
 
     private SocketProtocol protocol;
@@ -57,10 +55,11 @@ public class SocketClient
 
     private ByteChannel channel;
 
-    public SocketClient( BoltServerAddress address, SecurityPlan securityPlan, Logger logger )
+    public SocketClient( BoltServerAddress address, SecurityPlan securityPlan, int timeoutMillis, Logger logger )
     {
         this.address = address;
         this.securityPlan = securityPlan;
+        this.timeoutMillis = timeoutMillis;
         this.logger = logger;
         this.channel = null;
     }
@@ -121,7 +120,7 @@ public void start()
         try
         {
             logger.debug( "~~ [CONNECT] %s", address );
-            setChannel( ChannelFactory.create( address, securityPlan, logger ) );
+            setChannel( ChannelFactory.create( address, securityPlan, timeoutMillis, logger ) );
             protocol = negotiateProtocol();
             reader = protocol.reader();
             writer = protocol.writer();
@@ -280,36 +279,6 @@ public String toString()
         return "SocketClient[protocolVersion=" + version + "]";
     }
 
-    private static class ChannelFactory
-    {
-        public static ByteChannel create( BoltServerAddress address, SecurityPlan securityPlan, Logger logger )
-                throws IOException, GeneralSecurityException
-        {
-            SocketChannel soChannel = SocketChannel.open();
-            soChannel.setOption( StandardSocketOptions.SO_REUSEADDR, true );
-            soChannel.setOption( StandardSocketOptions.SO_KEEPALIVE, true );
-            soChannel.connect( address.toSocketAddress() );
-
-            ByteChannel channel;
-
-            if (securityPlan.requiresEncryption())
-            {
-                channel = TLSSocketChannel.create( address, securityPlan, soChannel, logger );
-            }
-            else
-            {
-                channel = soChannel;
-            }
-
-            if ( logger.isTraceEnabled() )
-            {
-                channel = new LoggingByteChannel( channel, logger );
-            }
-
-            return channel;
-        }
-    }
-
     public BoltServerAddress address()
     {
         return address;

driver/src/main/java/org/neo4j/driver/internal/net/SocketConnection.java

@@ -58,10 +58,10 @@ public class SocketConnection implements Connection
 
     private final Logger logger;
 
-    public SocketConnection( BoltServerAddress address, SecurityPlan securityPlan, Logging logging )
+    public SocketConnection( BoltServerAddress address, SecurityPlan securityPlan, int timeoutMillis, Logging logging )
     {
         this.logger = logging.getLog( format( "conn-%s", UUID.randomUUID().toString() ) );
-        this.socket = new SocketClient( address, securityPlan, logger );
+        this.socket = new SocketClient( address, securityPlan, timeoutMillis, logger );
         this.responseHandler = createResponseHandler( logger );
         this.socket.start();
     }

driver/src/main/java/org/neo4j/driver/internal/net/SocketConnector.java

@@ -47,7 +47,8 @@ public SocketConnector( ConnectionSettings connectionSettings, SecurityPlan secu
     @Override
     public final Connection connect( BoltServerAddress address )
     {
-        Connection connection = createConnection( address, securityPlan, logging );
+        Connection connection = createConnection( address, securityPlan, connectionSettings.timeoutMillis(),
+                logging );
 
         // Because SocketConnection is not thread safe, wrap it in this guard
         // to ensure concurrent access leads causes application errors
@@ -71,9 +72,10 @@ public final Connection connect( BoltServerAddress address )
      * <p>
      * <b>This method is package-private only for testing</b>
      */
-    Connection createConnection( BoltServerAddress address, SecurityPlan securityPlan, Logging logging )
+    Connection createConnection( BoltServerAddress address, SecurityPlan securityPlan, int timeoutMillis,
+            Logging logging )
     {
-        return new SocketConnection( address, securityPlan, logging );
+        return new SocketConnection( address, securityPlan, timeoutMillis, logging );
     }
 
     private static Map<String,Value> tokenAsMap( AuthToken token )

driver/src/main/java/org/neo4j/driver/v1/Config.java

@@ -61,6 +61,7 @@ public class Config
 
     private final int routingFailureLimit;
     private final long routingRetryDelayMillis;
+    private final int connectionTimeoutMillis;
 
     private Config( ConfigBuilder builder)
     {
@@ -73,6 +74,7 @@ private Config( ConfigBuilder builder)
         this.trustStrategy = builder.trustStrategy;
         this.routingFailureLimit = builder.routingFailureLimit;
         this.routingRetryDelayMillis = builder.routingRetryDelayMillis;
+        this.connectionTimeoutMillis = builder.connectionTimeoutMillis;
     }
 
     /**
@@ -127,6 +129,14 @@ public long idleTimeBeforeConnectionTest()
         return -1;
     }
 
+    /**
+     * @return the configured connection timeout value in milliseconds.
+     */
+    public int connectionTimeoutMillis()
+    {
+        return connectionTimeoutMillis;
+    }
+
     /**
      * @return the level of encryption required for all connections.
      */
@@ -176,7 +186,8 @@ public static class ConfigBuilder
         private EncryptionLevel encryptionLevel = EncryptionLevel.REQUIRED;
         private TrustStrategy trustStrategy = trustAllCertificates();
         private int routingFailureLimit = 1;
-        private long routingRetryDelayMillis = 5_000;
+        private long routingRetryDelayMillis = TimeUnit.SECONDS.toMillis( 5 );
+        private int connectionTimeoutMillis = (int) TimeUnit.SECONDS.toMillis( 5 );
 
         private ConfigBuilder() {}
 
@@ -367,6 +378,42 @@ public ConfigBuilder withRoutingRetryDelay( long delay, TimeUnit unit )
             return this;
         }
 
+        /**
+         * Specify socket connection timeout.
+         * <p>
+         * A timeout of zero is treated as an infinite timeout and will be bound by the timeout configured on the
+         * operating system level. The connection will block until established or an error occurs.
+         * <p>
+         * Timeout value should be greater or equal to zero and represent a valid {@code int} value when converted to
+         * {@link TimeUnit#MILLISECONDS milliseconds}.
+         * <p>
+         * The default value of this parameter is {@code 5 SECONDS}.
+         *
+         * @param value the timeout duration
+         * @param unit the unit in which duration is given
+         * @return this builder
+         * @throws IllegalArgumentException when given value is negative or does not fit in {@code int} when
+         * converted to milliseconds.
+         */
+        public ConfigBuilder withConnectionTimeout( long value, TimeUnit unit )
+        {
+            long connectionTimeoutMillis = unit.toMillis( value );
+            if ( connectionTimeoutMillis < 0 )
+            {
+                throw new IllegalArgumentException( String.format(
+                        "The connection timeout may not be smaller than 0, but was %d %s.", value, unit ) );
+            }
+            int connectionTimeoutMillisInt = (int) connectionTimeoutMillis;
+            if ( connectionTimeoutMillisInt != connectionTimeoutMillis )
+            {
+                throw new IllegalArgumentException( String.format(
+                        "The connection timeout must represent int value when converted to milliseconds %d.",
+                        connectionTimeoutMillis ) );
+            }
+            this.connectionTimeoutMillis = connectionTimeoutMillisInt;
+            return this;
+        }
+
         /**
          * Create a config instance from this builder.
          * @return a {@link Config} instance