added socket_connect_timeout and socket_keepalive options. fixed #353

3 files changed, 63 insertions, 15 deletions

diff --git a/CHANGES b/CHANGES
index 16c6245..d64f5e3 100644
--- a/CHANGES
+++ b/CHANGES
@@ -49,9 +49,19 @@
     * Fixed a bug where some encodings (like utf-16) were unusable on Python 3
       as command names and literals would get encoded.
     * Added an SSLConnection class that allows for secure connections through
-      stunnel or other means. Currently users must create a connection pool,
-      specifying connection_class=SSLConnection to use this feature. Thanks
-      https://github.com/oranagra.
+      stunnel or other means. Construct and SSL connection with the sll=True
+      option on client classes, using the rediss:// scheme from an URL, or
+      by passing the SSLConnection class to a connection pool's
+      connection_class argument. Thanks https://github.com/oranagra.
+    * Added a socket_connect_timeout option to control how long to wait while
+      establishing a TCP connection before timing out. This lets the client
+      fail fast when attempting to connect to a downed server while keeping
+      a more lenient timeout for all other socket operations.
+    * Added TCP Keep-alive support by passing use the socket_keepalive=True
+      option. Finer grain control can be achieved using the
+      socket_keepalive_options option which expects a dictionary with any of
+      the keys (socket.TCP_KEEPIDLE, socket.TCP_KEEPCNT, socket.TCP_KEEPINTVL)
+      and integers for values. Thanks Yossi Gottlieb.
 * 2.9.1
     * IPv6 support. Thanks https://github.com/amashinchi
 * 2.9.0
diff --git a/redis/client.py b/redis/client.py
index bb63435..c933221 100755
--- a/redis/client.py
+++ b/redis/client.py
@@ -388,6 +388,8 @@ class StrictRedis(object):
 
     def __init__(self, host='localhost', port=6379,
                  db=0, password=None, socket_timeout=None,
+                 socket_connect_timeout=None,
+                 socket_keepalive=None, socket_keepalive_options=None,
                  connection_pool=None, charset='utf-8',
                  errors='strict', decode_responses=False,
                  unix_socket_path=None,
@@ -409,9 +411,13 @@ class StrictRedis(object):
                     'connection_class': UnixDomainSocketConnection
                 })
             else:
+                # TCP specific options
                 kwargs.update({
                     'host': host,
-                    'port': port
+                    'port': port,
+                    'socket_connect_timeout': socket_connect_timeout,
+                    'socket_keepalive': socket_keepalive,
+                    'socket_keepalive_options': socket_keepalive_options,
                 })
 
                 if ssl:
diff --git a/redis/connection.py b/redis/connection.py
index c209750..f7297a0 100755
--- a/redis/connection.py
+++ b/redis/connection.py
@@ -354,7 +354,9 @@ class Connection(object):
     description_format = "Connection<host=%(host)s,port=%(port)s,db=%(db)s>"
 
     def __init__(self, host='localhost', port=6379, db=0, password=None,
-                 socket_timeout=None, encoding='utf-8',
+                 socket_timeout=None, socket_connect_timeout=None,
+                 socket_keepalive=False, socket_keepalive_options=None,
+                 encoding='utf-8',
                  encoding_errors='strict', decode_responses=False,
                  parser_class=DefaultParser, socket_read_size=65536):
         self.pid = os.getpid()
@@ -363,6 +365,9 @@ class Connection(object):
         self.db = db
         self.password = password
         self.socket_timeout = socket_timeout
+        self.socket_connect_timeout = socket_connect_timeout or socket_timeout
+        self.socket_keepalive = socket_keepalive
+        self.socket_keepalive_options = socket_keepalive_options or {}
         self.encoding = encoding
         self.encoding_errors = encoding_errors
         self.decode_responses = decode_responses
@@ -415,16 +420,43 @@ class Connection(object):
 
     def _connect(self):
         "Create a TCP socket connection"
-        # in 2.6+ try to use IPv6/4 compatibility, else just original code
-        if hasattr(socket, 'create_connection'):
-            sock = socket.create_connection((self.host, self.port),
-                                            self.socket_timeout)
-        else:
-            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-            sock.settimeout(self.socket_timeout)
-            sock.connect((self.host, self.port))
-        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-        return sock
+        # we want to mimic what socket.create_connection does to support
+        # ipv4/ipv6, but we want to set options prior to calling
+        # socket.connect()
+        err = None
+        for res in socket.getaddrinfo(self.host, self.port, 0,
+                                      socket.SOCK_STREAM):
+            family, socktype, proto, canonname, socket_address = res
+            sock = None
+            try:
+                sock = socket.socket(family, socktype, proto)
+                # TCP_NODELAY
+                sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+
+                # TCP_KEEPALIVE
+                if self.socket_keepalive:
+                    sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+                    for k, v in iteritems(self.socket_keepalive_options):
+                        sock.setsockopt(socket.SOL_TCP, k, v)
+
+                # set the socket_connect_timeout before we connect
+                sock.settimeout(self.socket_connect_timeout)
+
+                # connect
+                sock.connect(socket_address)
+
+                # set the socket_timeout now that we're connected
+                sock.settimeout(self.socket_timeout)
+                return sock
+
+            except socket.error as _:
+                err = _
+                if sock is not None:
+                    sock.close()
+
+        if err is not None:
+            raise err
+        raise socket.error("socket.getaddrinfo returned an empty list")
 
     def _error_message(self, exception):
         # args for socket.error can either be (errno, "message")