add 'network_timeout' option to Connection instances to specify timeout for socket operations

Mike Dirolf · Mike Dirolf · commit d7289d0c836e · 2009-10-13T16:20:25.000-04:00
diff --git a/pymongo/connection.py b/pymongo/connection.py
@@ -50,7 +50,7 @@ class Connection(object): # TODO support auth for pooling
 
     def __init__(self, host=None, port=None, pool_size=None,
                  auto_start_request=None, timeout=None, slave_okay=False,
-                 _connect=True):
+                 network_timeout=None, _connect=True):
         """Open a new connection to a Mongo instance at host:port.
 
         The resultant connection object has connection-pooling built in. It
@@ -82,6 +82,8 @@ def __init__(self, host=None, port=None, pool_size=None,
           - `timeout` (optional): max time to wait when attempting to acquire a
             connection from the connection pool before raising an exception -
             can be set to -1 to wait indefinitely
+          - `network_timeout` (optional): timeout (in seconds) to use for socket
+            operations - default is no timeout
         """
         if host is None:
             host = self.HOST
@@ -128,6 +130,8 @@ def __init__(self, host=None, port=None, pool_size=None,
         self.__sockets = [None for _ in range(self.__pool_size)]
         self.__currently_resetting = False
 
+        self.__network_timeout = network_timeout
+
         # cache of existing indexes used by ensure_index ops
         self.__index_cache = {}
 
@@ -296,7 +300,7 @@ def __find_master(self):
                     sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                     sock.settimeout(_CONNECT_TIMEOUT)
                     sock.connect((host, port))
-                    sock.settimeout(None)
+                    sock.settimeout(self.__network_timeout)
                     master = self.__master(sock)
                     if master is True:
                         self.__host = host
@@ -350,7 +354,7 @@ def __connect(self, socket_number):
             sock = self.__sockets[socket_number]
             sock.settimeout(_CONNECT_TIMEOUT)
             sock.connect((self.host(), self.port()))
-            sock.settimeout(None)
+            sock.settimeout(self.__network_timeout)
             _logger.debug("connected")
             return
         except socket.error:
@@ -493,8 +497,8 @@ def receive(length):
             while len(message) < length:
                 try:
                     chunk = sock.recv(length - len(message))
-                except socket.error:
-                    raise ConnectionFailure("connection closed")
+                except socket.error, e:
+                    raise ConnectionFailure(e)
                 if chunk == "":
                     raise ConnectionFailure("connection closed")
                 message += chunk
diff --git a/test/test_connection.py b/test/test_connection.py
@@ -128,5 +128,30 @@ def iterate():
 
         self.assertRaises(TypeError, iterate)
 
+    def test_socket_timeout(self):
+        no_timeout = Connection(self.host, self.port)
+        timeout = Connection(self.host, self.port, network_timeout=0.25)
+
+        no_timeout.pymongo_test.drop_collection("test")
+
+        no_timeout.pymongo_test.test.save({"x": 1})
+
+        where_func = """function (doc) {
+  var d = new Date().getTime() + 500;
+  var x = new Date().getTime();
+  while (x < d) {
+    x = new Date().getTime();
+  }
+  return true;
+}"""
+
+        def get_x(db):
+            return db.test.find().where(where_func).next()["x"]
+
+        self.assertEqual(1, get_x(no_timeout.pymongo_test))
+        self.assertRaises(ConnectionFailure, get_x, timeout.pymongo_test)
+        self.assertEqual(1, no_timeout.pymongo_test.test.find().next()["x"])
+
+
 if __name__ == "__main__":
     unittest.main()
