PYTHON-679 - Add socketKeepAlive option.

pymongo/common.py

@@ -300,6 +300,7 @@ VALIDATORS = {
     'authsource': validate_basestring,
     'gssapiservicename': validate_basestring,
     'uuidrepresentation': validate_uuid_representation,
+    'socketkeepalive': validate_boolean
 }
 
 

pymongo/connection.py

@@ -111,6 +111,9 @@ class Connection(MongoClient):
           - `waitQueueMultiple`: (integer or None) Multiplied by max_pool_size
             to give the number of threads allowed to wait for a socket at one
             time. Defaults to ``None`` (no waiters).
+          - `socketKeepAlive`: (boolean) Whether to send periodic keep-alive
+            packets on connected sockets. Defaults to ``False`` (do not send
+            keep-alive packets).
           - `auto_start_request`: If ``True`` (the default), each thread that
             accesses this Connection has a socket allocated to it for the
             thread's lifetime.  This ensures consistent reads, even if you read

pymongo/mongo_client.py

@@ -149,6 +149,9 @@ class MongoClient(common.BaseObject):
           - `waitQueueMultiple`: (integer or None) Multiplied by max_pool_size
             to give the number of threads allowed to wait for a socket at one
             time. Defaults to ``None`` (no waiters).
+          - `socketKeepAlive`: (boolean) Whether to send periodic keep-alive
+            packets on connected sockets. Defaults to ``False`` (do not send
+            keep-alive packets).
           - `auto_start_request`: If ``True``, each thread that accesses
             this :class:`MongoClient` has a socket allocated to it for the
             thread's lifetime.  This ensures consistent reads, even if you
@@ -303,6 +306,7 @@ class MongoClient(common.BaseObject):
         self.__conn_timeout = options.get('connecttimeoutms', 20.0)
         self.__wait_queue_timeout = options.get('waitqueuetimeoutms')
         self.__wait_queue_multiple = options.get('waitqueuemultiple')
+        self.__socket_keepalive = options.get('socketkeepalive', False)
 
         self.__use_ssl = options.get('ssl', None)
         self.__ssl_keyfile = options.get('ssl_keyfile', None)
@@ -489,7 +493,8 @@ class MongoClient(common.BaseObject):
             ssl_cert_reqs=self.__ssl_cert_reqs,
             ssl_ca_certs=self.__ssl_ca_certs,
             wait_queue_timeout=self.__wait_queue_timeout,
-            wait_queue_multiple=self.__wait_queue_multiple)
+            wait_queue_multiple=self.__wait_queue_multiple,
+            socket_keepalive=self.__socket_keepalive)
 
     def __check_auth(self, sock_info):
         """Authenticate using cached database credentials.

pymongo/mongo_replica_set_client.py

@@ -500,6 +500,9 @@ class MongoReplicaSetClient(common.BaseObject):
           - `waitQueueMultiple`: (integer or None) Multiplied by max_pool_size
             to give the number of threads allowed to wait for a socket at one
             time. Defaults to ``None`` (no waiters).
+          - `socketKeepAlive`: (boolean) Whether to send periodic keep-alive
+            packets on connected sockets. Defaults to ``False`` (do not send
+            keep-alive packets).
           - `auto_start_request`: If ``True``, each thread that accesses
             this :class:`MongoReplicaSetClient` has a socket allocated to it
             for the thread's lifetime, for each member of the set. For
@@ -656,6 +659,7 @@ class MongoReplicaSetClient(common.BaseObject):
         self.__conn_timeout = self.__opts.get('connecttimeoutms', 20.0)
         self.__wait_queue_timeout = self.__opts.get('waitqueuetimeoutms')
         self.__wait_queue_multiple = self.__opts.get('waitqueuemultiple')
+        self.__socket_keepalive = self.__opts.get('socketkeepalive', False)
         self.__use_ssl = self.__opts.get('ssl', None)
         self.__ssl_keyfile = self.__opts.get('ssl_keyfile', None)
         self.__ssl_certfile = self.__opts.get('ssl_certfile', None)
@@ -1046,6 +1050,7 @@ class MongoReplicaSetClient(common.BaseObject):
             self.__use_ssl,
             wait_queue_timeout=self.__wait_queue_timeout,
             wait_queue_multiple=self.__wait_queue_multiple,
+            socket_keepalive=self.__socket_keepalive,
             use_greenlets=self.__use_greenlets,
             ssl_keyfile=self.__ssl_keyfile,
             ssl_certfile=self.__ssl_certfile,

pymongo/pool.py

@@ -100,7 +100,8 @@ class Pool:
     def __init__(self, pair, max_size, net_timeout, conn_timeout, use_ssl,
                  use_greenlets, ssl_keyfile=None, ssl_certfile=None,
                  ssl_cert_reqs=None, ssl_ca_certs=None,
-                 wait_queue_timeout=None, wait_queue_multiple=None):
+                 wait_queue_timeout=None, wait_queue_multiple=None,
+                 socket_keepalive=False):
         """
         :Parameters:
           - `pair`: a (hostname, port) tuple
@@ -136,6 +137,9 @@ class Pool:
             free sockets.
           - `wait_queue_multiple`: (integer) Multiplied by max_pool_size to give
             the number of threads allowed to wait for a socket at one time.
+          - `socket_keepalive`: (boolean) Whether to send periodic keep-alive
+            packets on connected sockets. Defaults to ``False`` (do not send
+            keep-alive packets).
         """
         # Only check a socket's health with _closed() every once in a while.
         # Can override for testing: 0 to always check, None to never check.
@@ -154,6 +158,7 @@ class Pool:
         self.conn_timeout = conn_timeout
         self.wait_queue_timeout = wait_queue_timeout
         self.wait_queue_multiple = wait_queue_multiple
+        self.socket_keepalive = socket_keepalive
         self.use_ssl = use_ssl
         self.ssl_keyfile = ssl_keyfile
         self.ssl_certfile = ssl_certfile
@@ -240,6 +245,8 @@ class Pool:
             try:
                 sock = socket.socket(af, socktype, proto)
                 sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE,
+                                self.socket_keepalive)
                 sock.settimeout(self.conn_timeout)
                 sock.connect(sa)
                 return sock

pymongo/replica_set_connection.py

@@ -118,6 +118,9 @@ class ReplicaSetConnection(MongoReplicaSetClient):
           - `waitQueueMultiple`: (integer or None) Multiplied by max_pool_size
             to give the number of threads allowed to wait for a socket at one
             time. Defaults to ``None`` (no waiters).
+          - `socketKeepAlive`: (boolean) Whether to send periodic keep-alive
+            packets on connected sockets. Defaults to ``False`` (do not send
+            keep-alive packets).
           - `auto_start_request`: If ``True`` (the default), each thread that
             accesses this :class:`ReplicaSetConnection` has a socket allocated
             to it for the thread's lifetime, for each member of the set. For

test/test_client.py

@@ -70,6 +70,7 @@ class TestClient(unittest.TestCase, TestRequestMixin):
                              connectTimeoutMS=20000,
                              waitQueueTimeoutMS=None,
                              waitQueueMultiple=None,
+                             socketKeepAlive=False,
                              auto_start_request=False,
                              use_greenlets=False,
                              replicaSet=None,
@@ -86,6 +87,7 @@ class TestClient(unittest.TestCase, TestRequestMixin):
         self.assertEqual(20.0, client._MongoClient__conn_timeout)
         self.assertEqual(None, client._MongoClient__wait_queue_timeout)
         self.assertEqual(None, client._MongoClient__wait_queue_multiple)
+        self.assertFalse(client._MongoClient__socket_keepalive)
         self.assertFalse(client.auto_start_request)
         self.assertFalse(client.use_greenlets)
         self.assertEqual(None, client._MongoClient__repl)
@@ -532,6 +534,10 @@ class TestClient(unittest.TestCase, TestRequestMixin):
         self.assertEqual(pool.wait_queue_multiple, 2)
         self.assertEqual(pool._socket_semaphore.waiter_semaphore.counter, 6)
 
+    def test_socketKeepAlive(self):
+        client = MongoClient(host, port, socketKeepAlive=True)
+        self.assertTrue(get_pool(client).socket_keepalive)
+
     def test_tz_aware(self):
         self.assertRaises(ConfigurationError, MongoClient, tz_aware='foo')
 

test/test_replica_set_client.py

@@ -131,6 +131,7 @@ class TestReplicaSetClient(TestReplicaSetClientBase, TestRequestMixin):
                                        connectTimeoutMS=20000,
                                        waitQueueTimeoutMS=None,
                                        waitQueueMultiple=None,
+                                       socketKeepAlive=False,
                                        auto_start_request=False,
                                        use_greenlets=False,
                                        replicaSet='myreplset',  # Required
@@ -151,6 +152,7 @@ class TestReplicaSetClient(TestReplicaSetClientBase, TestRequestMixin):
                          client._MongoReplicaSetClient__wait_queue_timeout)
         self.assertEqual(None,
                          client._MongoReplicaSetClient__wait_queue_multiple)
+        self.assertFalse(client._MongoReplicaSetClient__socket_keepalive)
         self.assertFalse(client.auto_start_request)
         self.assertFalse(client.use_greenlets)
         self.assertEqual('myreplset',
@@ -709,6 +711,11 @@ class TestReplicaSetClient(TestReplicaSetClientBase, TestRequestMixin):
         self.assertEqual(pool.wait_queue_multiple, 2)
         self.assertEqual(pool._socket_semaphore.waiter_semaphore.counter, 6)
 
+    def test_socketKeepAlive(self):
+        client = self._get_client(socketKeepAlive=True)
+        pool = get_pool(client)
+        self.assertTrue(pool.socket_keepalive)
+
     def test_tz_aware(self):
         self.assertRaises(ConfigurationError, MongoReplicaSetClient,
                           tz_aware='foo', replicaSet=self.name)