PYTHON-1272 Fix deadlock when garbage collecting pinned cursors and sessions (#642)

ShaneHarvey · ShaneHarvey · commit 1b2a599706d6 · 2021-06-23T10:31:29.000-07:00
It's not safe to return the pinned connection to the pool from within Cursor.del because the Pool's lock may be held by a python thread while the cyclic garbage collector runs. Instead we send the cursor cleanup request to the client's background thread. The thread will send killCursors on the pinned socket and then return the socket to the pool. Also fixed a similar bug when garbage collecting a pinned session. (cherry picked from commit 6bc5e08)
diff --git a/pymongo/client_session.py b/pymongo/client_session.py
@@ -288,14 +288,15 @@ class _TxnState(object):
 
 class _Transaction(object):
     """Internal class to hold transaction information in a ClientSession."""
-    def __init__(self, opts):
+    def __init__(self, opts, client):
         self.opts = opts
         self.state = _TxnState.NONE
         self.sharded = False
         self.pinned_address = None
         self.sock_mgr = None
         self.recovery_token = None
         self.attempt = 0
+        self.client = client
 
     def active(self):
         return self.state in (_TxnState.STARTING, _TxnState.IN_PROGRESS)
@@ -329,6 +330,13 @@ def reset(self):
         self.recovery_token = None
         self.attempt = 0
 
+    def __del__(self):
+        if self.sock_mgr:
+            # Reuse the cursor closing machinery to return the socket to the
+            # pool soon.
+            self.client._close_cursor_soon(0, None, self.sock_mgr)
+            self.sock_mgr = None
+
 
 def _reraise_with_unknown_commit(exc):
     """Re-raise an exception with the UnknownTransactionCommitResult label."""
@@ -382,7 +390,7 @@ def __init__(self, client, server_session, options, authset, implicit):
         self._operation_time = None
         # Is this an implicitly created session?
         self._implicit = implicit
-        self._transaction = _Transaction(None)
+        self._transaction = _Transaction(None, client)
 
     def end_session(self):
         """Finish this session. If a transaction has started, abort it.
diff --git a/pymongo/command_cursor.py b/pymongo/command_cursor.py
@@ -65,29 +65,31 @@ def __init__(self, collection, cursor_info, address, retrieved=0,
             raise TypeError("max_await_time_ms must be an integer or None")
 
     def __del__(self):
-        if self.__id and not self.__killed:
-            self.__die()
+        self.__die()
 
     def __die(self, synchronous=False):
         """Closes this cursor.
         """
         already_killed = self.__killed
         self.__killed = True
         if self.__id and not already_killed:
+            cursor_id = self.__id
             address = _CursorAddress(
                 self.__address, self.__collection.full_name)
-            if synchronous:
-                self.__collection.database.client._close_cursor_now(
-                    self.__id, address, session=self.__session,
-                    sock_mgr=self.__sock_mgr)
-            else:
-                # The cursor will be closed later in a different session.
-                self.__collection.database.client._close_cursor(
-                    self.__id, address)
-        if self.__sock_mgr:
-            self.__sock_mgr.close()
-            self.__sock_mgr = None
-        self.__end_session(synchronous)
+        else:
+            # Skip killCursors.
+            cursor_id = 0
+            address = None
+        self.__collection.database.client._cleanup_cursor(
+            synchronous,
+            cursor_id,
+            address,
+            self.__sock_mgr,
+            self.__session,
+            self.__explicit_session)
+        if not self.__explicit_session:
+            self.__session = None
+        self.__sock_mgr = None
 
     def __end_session(self, synchronous):
         if self.__session and not self.__explicit_session:
@@ -186,7 +188,7 @@ def __send_message(self, operation):
             self.__id = response.data.cursor_id
 
         if self.__id == 0:
-            self.__die(True)
+            self.close()
         self.__data = deque(documents)
 
     def _unpack_response(self, response, cursor_id, codec_options,
diff --git a/pymongo/cursor.py b/pymongo/cursor.py
@@ -37,7 +37,6 @@
                              _RawBatchGetMore,
                              _Query,
                              _RawBatchQuery)
-from pymongo.monitoring import ConnectionClosedReason
 from pymongo.response import PinnedResponse
 
 # These errors mean that the server has already killed the cursor so there is
@@ -109,28 +108,23 @@ class CursorType(object):
     """
 
 
-# This has to be an old style class due to
-# http://bugs.jython.org/issue1057
-class _SocketManager:
+class _SocketManager(object):
     """Used with exhaust cursors to ensure the socket is returned.
     """
     def __init__(self, sock, more_to_come):
         self.sock = sock
         self.more_to_come = more_to_come
-        self.__closed = False
+        self.closed = False
         self.lock = threading.Lock()
 
-    def __del__(self):
-        self.close()
-
     def update_exhaust(self, more_to_come):
         self.more_to_come = more_to_come
 
     def close(self):
         """Return this instance's socket to the connection pool.
         """
-        if not self.__closed:
-            self.__closed = True
+        if not self.closed:
+            self.closed = True
             self.sock.unpin()
             self.sock = None
 
@@ -159,6 +153,7 @@ def __init__(self, collection, filter=None, projection=None, skip=0,
         """
         # Initialize all attributes used in __del__ before possibly raising
         # an error to avoid attribute errors during garbage collection.
+        self.__collection = collection
         self.__id = None
         self.__exhaust = False
         self.__sock_mgr = None
@@ -211,7 +206,6 @@ def __init__(self, collection, filter=None, projection=None, skip=0,
                 projection = {"_id": 1}
             projection = helpers._fields_list_to_dict(projection, "projection")
 
-        self.__collection = collection
         self.__spec = spec
         self.__projection = projection
         self.__skip = skip
@@ -297,6 +291,7 @@ def rewind(self):
         be sent to the server, even if the resultant data has already been
         retrieved by this cursor.
         """
+        self.close()
         self.__data = deque()
         self.__id = None
         self.__address = None
@@ -353,29 +348,23 @@ def __die(self, synchronous=False):
 
         self.__killed = True
         if self.__id and not already_killed:
-            if self.__exhaust and self.__sock_mgr:
-                # If this is an exhaust cursor and we haven't completely
-                # exhausted the result set we *must* close the socket
-                # to stop the server from sending more data.
-                self.__sock_mgr.sock.close_socket(
-                    ConnectionClosedReason.ERROR)
-            else:
-                address = _CursorAddress(
-                    self.__address, self.__collection.full_name)
-                if synchronous:
-                    self.__collection.database.client._close_cursor_now(
-                        self.__id, address, session=self.__session,
-                        sock_mgr=self.__sock_mgr)
-                else:
-                    # The cursor will be closed later in a different session.
-                    self.__collection.database.client._close_cursor(
-                        self.__id, address)
-        if self.__sock_mgr:
-            self.__sock_mgr.close()
-            self.__sock_mgr = None
-        if self.__session and not self.__explicit_session:
-            self.__session._end_session(lock=synchronous)
+            cursor_id = self.__id
+            address = _CursorAddress(
+                self.__address, self.__collection.full_name)
+        else:
+            # Skip killCursors.
+            cursor_id = 0
+            address = None
+        self.__collection.database.client._cleanup_cursor(
+            synchronous,
+            cursor_id,
+            address,
+            self.__sock_mgr,
+            self.__session,
+            self.__explicit_session)
+        if not self.__explicit_session:
             self.__session = None
+        self.__sock_mgr = None
 
     def close(self):
         """Explicitly close / kill this cursor.
@@ -1098,10 +1087,10 @@ def __send_message(self, operation):
         if self.__id == 0:
             # Don't wait for garbage collection to call __del__, return the
             # socket and the session to the pool now.
-            self.__die()
+            self.close()
 
         if self.__limit and self.__id and self.__limit <= self.__retrieved:
-            self.__die()
+            self.close()
 
     def _unpack_response(self, response, cursor_id, codec_options,
                          user_fields=None, legacy_response=False):
diff --git a/pymongo/mongo_client.py b/pymongo/mongo_client.py
@@ -1648,10 +1648,47 @@ def close_cursor(self, cursor_id, address=None):
         if not isinstance(cursor_id, integer_types):
             raise TypeError("cursor_id must be an instance of (int, long)")
 
-        self._close_cursor(cursor_id, address)
+        self._close_cursor_soon(cursor_id, address)
 
-    def _close_cursor(self, cursor_id, address):
-        """Send a kill cursors message with the given id.
+    def _cleanup_cursor(self, locks_allowed, cursor_id, address, sock_mgr,
+                        session, explicit_session):
+        """Cleanup a cursor from cursor.close() or __del__.
+
+        This method handles cleanup for Cursors/CommandCursors including any
+        pinned connection or implicit session attached at the time the cursor
+        was closed or garbage collected.
+
+        :Parameters:
+          - `locks_allowed`: True if we are allowed to acquire locks.
+          - `cursor_id`: The cursor id which may be 0.
+          - `address`: The _CursorAddress.
+          - `sock_mgr`: The _SocketManager for the pinned connection or None.
+          - `session`: The cursor's session.
+          - `explicit_session`: True if the session was passed explicitly.
+        """
+        if locks_allowed:
+            if cursor_id:
+                if sock_mgr and sock_mgr.more_to_come:
+                    # If this is an exhaust cursor and we haven't completely
+                    # exhausted the result set we *must* close the socket
+                    # to stop the server from sending more data.
+                    sock_mgr.sock.close_socket(
+                        ConnectionClosedReason.ERROR)
+                else:
+                    self._close_cursor_now(
+                        cursor_id, address, session=session,
+                        sock_mgr=sock_mgr)
+            if sock_mgr:
+                sock_mgr.close()
+        else:
+            # The cursor will be closed later in a different session.
+            if cursor_id or sock_mgr:
+                self._close_cursor_soon(cursor_id, address, sock_mgr)
+        if session and not explicit_session:
+            session._end_session(lock=locks_allowed)
+
+    def _close_cursor_soon(self, cursor_id, address, sock_mgr=None):
+        """Request that a cursor and/or connection be cleaned up soon
 
         What closing the cursor actually means depends on this client's
         cursor manager. If there is none, the cursor is closed asynchronously
@@ -1660,7 +1697,7 @@ def _close_cursor(self, cursor_id, address):
         if self.__cursor_manager is not None:
             self.__cursor_manager.close(cursor_id, address)
         else:
-            self.__kill_cursors_queue.append((address, [cursor_id]))
+            self.__kill_cursors_queue.append((address, cursor_id, sock_mgr))
 
     def _close_cursor_now(self, cursor_id, address=None, session=None,
                           sock_mgr=None):
@@ -1687,7 +1724,7 @@ def _close_cursor_now(self, cursor_id, address=None, session=None,
                         [cursor_id], address, self._get_topology(), session)
             except PyMongoError:
                 # Make another attempt to kill the cursor later.
-                self.__kill_cursors_queue.append((address, [cursor_id]))
+                self._close_cursor_soon(cursor_id, address)
 
     def kill_cursors(self, cursor_ids, address=None):
         """DEPRECATED - Send a kill cursors message soon with the given ids.
@@ -1718,7 +1755,8 @@ def kill_cursors(self, cursor_ids, address=None):
             raise TypeError("cursor_ids must be a list")
 
         # "Atomic", needs no lock.
-        self.__kill_cursors_queue.append((address, cursor_ids))
+        for cursor_id in cursor_ids:
+            self.__kill_cursors_queue.append((address, cursor_id, None))
 
     def _kill_cursors(self, cursor_ids, address, topology, session):
         """Send a kill cursors message with the given ids."""
@@ -1783,15 +1821,26 @@ def _kill_cursor_impl(self, cursor_ids, address, session, sock_info):
     def _process_kill_cursors(self):
         """Process any pending kill cursors requests."""
         address_to_cursor_ids = defaultdict(list)
+        pinned_cursors = []
 
         # Other threads or the GC may append to the queue concurrently.
         while True:
             try:
-                address, cursor_ids = self.__kill_cursors_queue.pop()
+                address, cursor_id, sock_mgr = self.__kill_cursors_queue.pop()
             except IndexError:
                 break
 
-            address_to_cursor_ids[address].extend(cursor_ids)
+            if sock_mgr:
+                pinned_cursors.append((address, cursor_id, sock_mgr))
+            else:
+                address_to_cursor_ids[address].append(cursor_id)
+
+        for address, cursor_id, sock_mgr in pinned_cursors:
+            try:
+                self._cleanup_cursor(True, cursor_id, address, sock_mgr,
+                                     None, False)
+            except Exception:
+                helpers._handle_exception()
 
         # Don't re-open topology if it's closed and there's no pending cursors.
         if address_to_cursor_ids:
diff --git a/test/__init__.py b/test/__init__.py
@@ -185,8 +185,8 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     def __del__(self):
         if self._enabled:
-            print(
-                '\nERROR: client_knobs still enabled! HEARTBEAT_FREQUENCY=%s, '
+            msg = (
+                'ERROR: client_knobs still enabled! HEARTBEAT_FREQUENCY=%s, '
                 'MIN_HEARTBEAT_INTERVAL=%s, KILL_CURSOR_FREQUENCY=%s, '
                 'EVENTS_QUEUE_FREQUENCY=%s, stack:\n%s' % (
                     common.HEARTBEAT_FREQUENCY,
@@ -195,6 +195,7 @@ def __del__(self):
                     common.EVENTS_QUEUE_FREQUENCY,
                     self._stack))
             self.disable()
+            raise Exception(msg)
 
 
 def _all_users(db):
diff --git a/test/test_gridfs.py b/test/test_gridfs.py
@@ -33,6 +33,7 @@
 from pymongo.read_preferences import ReadPreference
 from gridfs.errors import CorruptGridFile, FileExists, NoFile
 from test.test_replica_set_client import TestReplicaSetClientBase
+from gridfs.grid_file import GridOutCursor
 from test import (client_context,
                   unittest,
                   IntegrationTest)
@@ -429,6 +430,14 @@ def test_gridfs_find(self):
         cursor.close()
         self.assertRaises(TypeError, self.fs.find, {}, {"_id": True})
 
+    def test_delete_not_initialized(self):
+        # Creating a cursor with invalid arguments will not run __init__
+        # but will still call __del__.
+        cursor = GridOutCursor.__new__(GridOutCursor)  # Skip calling __init__
+        with self.assertRaises(TypeError):
+            cursor.__init__(self.db.fs.files, {}, {"_id": True})
+        cursor.__del__()  # no error
+
     def test_gridfs_find_one(self):
         self.assertEqual(None, self.fs.find_one())
 
diff --git a/test/test_load_balancer.py b/test/test_load_balancer.py
diff --git a/test/test_retryable_reads.py b/test/test_retryable_reads.py
