Don't close socket after operation failure PYTHON-395

A socket must be closed after a network error, but not after an error response to getLastError. Sockets were wrongly closed in this case if auto_start_request=False.

Author: A. Jesse Jiryu Davis

pymongo/connection.py

@@ -832,6 +832,9 @@ def _send_message(self, message, with_last_error=False):
 
             self.__pool.maybe_return_socket(sock_info)
             return rv
+        except OperationFailure:
+            self.__pool.maybe_return_socket(sock_info)
+            raise
         except (ConnectionFailure, socket.error), e:
             self.disconnect()
             raise AutoReconnect(str(e))

pymongo/replica_set_connection.py

@@ -1002,6 +1002,9 @@ def _send_message(self, msg, safe=False, _connection_to_use=None):
                 rv = self.__check_response_to_last_error(response)
             member.pool.maybe_return_socket(sock_info)
             return rv
+        except OperationFailure:
+            member.pool.maybe_return_socket(sock_info)
+            raise
         except(ConnectionFailure, socket.error), why:
             member.pool.discard_socket(sock_info)
             if _connection_to_use in (None, -1):

test/test_connection.py

@@ -623,5 +623,42 @@ def interrupter():
             db.foo.find().next()
         )
 
+    def test_operation_failure_without_request(self):
+        # Ensure Connection doesn't close socket after it gets an error
+        # response to getLastError. PYTHON-395.
+        c = get_connection(auto_start_request=False)
+        pool = c._Connection__pool
+        self.assertEqual(1, len(pool.sockets))
+        old_sock_info = iter(pool.sockets).next()
+        c.pymongo_test.test.drop()
+        c.pymongo_test.test.insert({'_id': 'foo'}, safe=True)
+        self.assertRaises(
+            OperationFailure,
+            c.pymongo_test.test.insert, {'_id': 'foo'}, safe=True)
+
+        self.assertEqual(1, len(pool.sockets))
+        new_sock_info = iter(pool.sockets).next()
+        self.assertEqual(old_sock_info, new_sock_info)
+
+    def test_operation_failure_with_request(self):
+        # Ensure Connection doesn't close socket after it gets an error
+        # response to getLastError. PYTHON-395.
+        c = get_connection(auto_start_request=True)
+        pool = c._Connection__pool
+
+        # Connection has reserved a socket for this thread
+        self.assertTrue(isinstance(pool._get_request_state(), SocketInfo))
+
+        old_sock_info = pool._get_request_state()
+        c.pymongo_test.test.drop()
+        c.pymongo_test.test.insert({'_id': 'foo'}, safe=True)
+        self.assertRaises(
+            OperationFailure,
+            c.pymongo_test.test.insert, {'_id': 'foo'}, safe=True)
+
+        # OperationFailure doesn't affect the request socket
+        self.assertEqual(old_sock_info, pool._get_request_state())
+
+
 if __name__ == "__main__":
     unittest.main()

test/test_replica_set_connection.py

@@ -36,6 +36,7 @@
 from pymongo.replica_set_connection import ReplicaSetConnection
 from pymongo.replica_set_connection import _partition_node
 from pymongo.database import Database
+from pymongo.pool import SocketInfo
 from pymongo.errors import (AutoReconnect,
                             ConfigurationError,
                             ConnectionFailure,
@@ -648,6 +649,46 @@ def sigalarm(num, frame):
             if old_signal_handler:
                 signal.signal(signal.SIGALRM, old_signal_handler)
 
+    def test_operation_failure_without_request(self):
+        # Ensure ReplicaSetConnection doesn't close socket after it gets an
+        # error response to getLastError. PYTHON-395.
+        c = self._get_connection(auto_start_request=False)
+        pool = c._ReplicaSetConnection__members[self.primary].pool
+        self.assertEqual(1, len(pool.sockets))
+        old_sock_info = iter(pool.sockets).next()
+        c.pymongo_test.test.drop()
+        c.pymongo_test.test.insert({'_id': 'foo'}, safe=True)
+        self.assertRaises(
+            OperationFailure,
+            c.pymongo_test.test.insert, {'_id': 'foo'}, safe=True)
+
+        self.assertEqual(1, len(pool.sockets))
+        new_sock_info = iter(pool.sockets).next()
+
+        self.assertEqual(old_sock_info, new_sock_info)
+        c.close()
+
+    def test_operation_failure_with_request(self):
+        # Ensure ReplicaSetConnection doesn't close socket after it gets an
+        # error response to getLastError. PYTHON-395.
+        c = self._get_connection(auto_start_request=True)
+        c.pymongo_test.test.find_one()
+        pool = c._ReplicaSetConnection__members[self.primary].pool
+
+        # Connection has reserved a socket for this thread
+        self.assertTrue(isinstance(pool._get_request_state(), SocketInfo))
+
+        old_sock_info = pool._get_request_state()
+        c.pymongo_test.test.drop()
+        c.pymongo_test.test.insert({'_id': 'foo'}, safe=True)
+        self.assertRaises(
+            OperationFailure,
+            c.pymongo_test.test.insert, {'_id': 'foo'}, safe=True)
+
+        # OperationFailure doesn't affect the request socket
+        self.assertEqual(old_sock_info, pool._get_request_state())
+        c.close()
+
     def test_auto_start_request(self):
         for bad_horrible_value in (None, 5, 'hi!'):
             self.assertRaises(