deal with connect timeout

tornadis/client.py

@@ -24,23 +24,27 @@ class Client(object):
     Attributes:
         host (string): the host name to connect to.
         port (int): the port to connect to.
+        connect_timeout (int): connect timeout (seconds)
         subscribed (boolean): is the client object subscribed to redis
             (with pubsub methods).
         __reply_queue (toro.Queue): toro queue to put redis replies.
         __reader: hiredis reader object.
         __connection: tornadis low level Connection object.
     """
 
-    def __init__(self, host='localhost', port=6379, ioloop=None):
+    def __init__(self, host='localhost', port=6379, connect_timeout=20,
+                 ioloop=None):
         """Constructor.
 
         Args:
             host (string): the host name to connect to.
             port (int): the port to connect to.
+            connect_timeout (int): connect timeout (seconds)
             ioloop (IOLoop): the tornado ioloop to use.
         """
         self.host = host
         self.port = port
+        self.connect_timeout = connect_timeout
         self.subscribed = False
         self.__ioloop = ioloop or tornado.ioloop.IOLoop.instance()
         self.__connection = None
@@ -66,6 +70,7 @@ def connect(self):
         self.__reply_queue = toro.Queue()
         self.__reader = hiredis.Reader()
         self.__connection = Connection(host=self.host, port=self.port,
+                                       connect_timeout=self.connect_timeout,
                                        ioloop=self.__ioloop)
         yield self.__connection.connect()
         self.__connection.register_read_until_close_callback(cb1, cb2)

tornadis/connection.py

@@ -16,19 +16,23 @@ class Connection(object):
     Attributes:
         host (string): the host name to connect to.
         port (int): the port to connect to.
+        connect_timeout (int): connect timeout (seconds)
         connected (boolean): is the connection object really connected.
     """
 
-    def __init__(self, host='localhost', port=6379, ioloop=None):
+    def __init__(self, host='localhost', port=6379, connect_timeout=20,
+                 ioloop=None):
         """Constructor.
 
         Args:
             host (string): the host name to connect to.
             port (int): the port to connect to.
+            connect_timeout (int): connection timeout (seconds)
             ioloop (IOLoop): the tornado ioloop to use.
         """
         self.host = host
         self.port = port
+        self.connect_timeout = connect_timeout
         self.connected = False
         self.__stream = None
         self.__ioloop = ioloop or tornado.ioloop.IOLoop.instance()
@@ -46,14 +50,20 @@ def connect(self):
         self.__socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.__stream = tornado.iostream.IOStream(self.__socket,
                                                   io_loop=self.__ioloop)
-        future = self.__stream.connect((self.host, self.port))
+        cb = self._connect_timeout_callback
+        handle = self.__ioloop.call_later(self.connect_timeout, cb)
         try:
-            yield future
+            yield self.__stream.connect((self.host, self.port))
         except:
+            self.__ioloop.remove_timeout(handle)
             raise ConnectionError("can't connect to %s:%i" % (self.host,
                                                               self.port))
+        self.__ioloop.remove_timeout(handle)
         self.connected = True
 
+    def _connect_timeout_callback(self):
+        self.disconnect()
+
     def disconnect(self):
         """Disconnects the object.
         """