PYTHON-4494 - AsyncMongoClient._cleanup_cursor needs to be synchronous (mongodb#1680)

Author: NoahStapp

pymongo/asynchronous/client_session.py

@@ -531,6 +531,12 @@ async def _end_session(self, lock: bool) -> None:
                 self._client._return_server_session(self._server_session)
                 self._server_session = None
 
+    def _end_implicit_session(self) -> None:
+        # Implicit sessions can't be part of transactions or pinned connections
+        if self._server_session is not None:
+            self._client._return_server_session(self._server_session)
+            self._server_session = None
+
     def _check_ended(self) -> None:
         if self._server_session is None:
             raise InvalidOperation("Cannot use ended session")

pymongo/asynchronous/command_cursor.py

@@ -81,8 +81,8 @@ def __init__(
         self._explicit_session = explicit_session
         self._killed = self._id == 0
         self._comment = comment
-        if _IS_SYNC and self._killed:
-            self._end_session(True)  # type: ignore[unused-coroutine]
+        if self._killed:
+            self._end_session()
 
         if "ns" in cursor_info:  # noqa: SIM401
             self._ns = cursor_info["ns"]
@@ -95,8 +95,7 @@ def __init__(
             raise TypeError("max_await_time_ms must be an integer or None")
 
     def __del__(self) -> None:
-        if _IS_SYNC:
-            self._die(False)  # type: ignore[unused-coroutine]
+        self._die_no_lock()
 
     def batch_size(self, batch_size: int) -> AsyncCommandCursor[_DocumentType]:
         """Limits the number of documents returned in one batch. Each batch
@@ -198,8 +197,7 @@ def session(self) -> Optional[ClientSession]:
             return self._session
         return None
 
-    async def _die(self, synchronous: bool = False) -> None:
-        """Closes this cursor."""
+    def _prepare_to_die(self) -> tuple[int, Optional[_CursorAddress]]:
         already_killed = self._killed
         self._killed = True
         if self._id and not already_killed:
@@ -210,8 +208,22 @@ async def _die(self, synchronous: bool = False) -> None:
             # Skip killCursors.
             cursor_id = 0
             address = None
-        await self._collection.database.client._cleanup_cursor(
-            synchronous,
+        return cursor_id, address
+
+    def _die_no_lock(self) -> None:
+        """Closes this cursor without acquiring a lock."""
+        cursor_id, address = self._prepare_to_die()
+        self._collection.database.client._cleanup_cursor_no_lock(
+            cursor_id, address, self._sock_mgr, self._session, self._explicit_session
+        )
+        if not self._explicit_session:
+            self._session = None
+        self._sock_mgr = None
+
+    async def _die_lock(self) -> None:
+        """Closes this cursor."""
+        cursor_id, address = self._prepare_to_die()
+        await self._collection.database.client._cleanup_cursor_lock(
             cursor_id,
             address,
             self._sock_mgr,
@@ -222,14 +234,14 @@ async def _die(self, synchronous: bool = False) -> None:
             self._session = None
         self._sock_mgr = None
 
-    async def _end_session(self, synchronous: bool) -> None:
+    def _end_session(self) -> None:
         if self._session and not self._explicit_session:
-            await self._session._end_session(lock=synchronous)
+            self._session._end_implicit_session()
             self._session = None
 
     async def close(self) -> None:
         """Explicitly close / kill this cursor."""
-        await self._die(True)
+        await self._die_lock()
 
     async def _send_message(self, operation: _GetMore) -> None:
         """Send a getmore message and handle the response."""
@@ -243,7 +255,7 @@ async def _send_message(self, operation: _GetMore) -> None:
                 # Don't send killCursors because the cursor is already closed.
                 self._killed = True
             if exc.timeout:
-                await self._die(False)
+                self._die_no_lock()
             else:
                 # Return the session and pinned connection, if necessary.
                 await self.close()
@@ -305,7 +317,7 @@ async def _refresh(self) -> int:
                 )
             )
         else:  # Cursor id is zero nothing else to return
-            await self._die(True)
+            await self._die_lock()
 
         return len(self._data)
 

pymongo/asynchronous/cursor.py

@@ -259,8 +259,7 @@ def retrieved(self) -> int:
         return self._retrieved
 
     def __del__(self) -> None:
-        if _IS_SYNC:
-            self._die()  # type: ignore[unused-coroutine]
+        self._die_no_lock()
 
     def clone(self) -> AsyncCursor[_DocumentType]:
         """Get a clone of this cursor.
@@ -996,14 +995,7 @@ def _deepcopy(
                 y[key] = value
         return y
 
-    async def _die(self, synchronous: bool = False) -> None:
-        """Closes this cursor."""
-        try:
-            already_killed = self._killed
-        except AttributeError:
-            # ___init__ did not run to completion (or at all).
-            return
-
+    def _prepare_to_die(self, already_killed: bool) -> tuple[int, Optional[_CursorAddress]]:
         self._killed = True
         if self._id and not already_killed:
             cursor_id = self._id
@@ -1013,8 +1005,34 @@ async def _die(self, synchronous: bool = False) -> None:
             # Skip killCursors.
             cursor_id = 0
             address = None
-        await self._collection.database.client._cleanup_cursor(
-            synchronous,
+        return cursor_id, address
+
+    def _die_no_lock(self) -> None:
+        """Closes this cursor without acquiring a lock."""
+        try:
+            already_killed = self._killed
+        except AttributeError:
+            # ___init__ did not run to completion (or at all).
+            return
+
+        cursor_id, address = self._prepare_to_die(already_killed)
+        self._collection.database.client._cleanup_cursor_no_lock(
+            cursor_id, address, self._sock_mgr, self._session, self._explicit_session
+        )
+        if not self._explicit_session:
+            self._session = None
+        self._sock_mgr = None
+
+    async def _die_lock(self) -> None:
+        """Closes this cursor."""
+        try:
+            already_killed = self._killed
+        except AttributeError:
+            # ___init__ did not run to completion (or at all).
+            return
+
+        cursor_id, address = self._prepare_to_die(already_killed)
+        await self._collection.database.client._cleanup_cursor_lock(
             cursor_id,
             address,
             self._sock_mgr,
@@ -1027,7 +1045,7 @@ async def _die(self, synchronous: bool = False) -> None:
 
     async def close(self) -> None:
         """Explicitly close / kill this cursor."""
-        await self._die(True)
+        await self._die_lock()
 
     async def distinct(self, key: str) -> list:
         """Get a list of distinct values for `key` among all documents
@@ -1080,7 +1098,7 @@ async def _send_message(self, operation: Union[_Query, _GetMore]) -> None:
                 # Don't send killCursors because the cursor is already closed.
                 self._killed = True
             if exc.timeout:
-                await self._die(False)
+                self._die_no_lock()
             else:
                 await self.close()
             # If this is a tailable cursor the error is likely

pymongo/asynchronous/mongo_client.py

@@ -1857,48 +1857,63 @@ async def _retryable_write(
         async with self._tmp_session(session) as s:
             return await self._retry_with_session(retryable, func, s, bulk, operation, operation_id)
 
-    async def _cleanup_cursor(
+    def _cleanup_cursor_no_lock(
         self,
-        locks_allowed: bool,
         cursor_id: int,
         address: Optional[_CursorAddress],
         conn_mgr: _ConnectionManager,
         session: Optional[ClientSession],
         explicit_session: bool,
     ) -> None:
-        """Cleanup a cursor from cursor.close() or __del__.
+        """Cleanup a cursor from __del__ without locking.
+
+        This method handles cleanup for Cursors/CommandCursors including any
+        pinned connection attached at the time the cursor
+        was garbage collected.
+
+        :param cursor_id: The cursor id which may be 0.
+        :param address: The _CursorAddress.
+        :param conn_mgr: The _ConnectionManager for the pinned connection or None.
+        """
+        # The cursor will be closed later in a different session.
+        if cursor_id or conn_mgr:
+            self._close_cursor_soon(cursor_id, address, conn_mgr)
+        if session and not explicit_session:
+            session._end_implicit_session()
+
+    async def _cleanup_cursor_lock(
+        self,
+        cursor_id: int,
+        address: Optional[_CursorAddress],
+        conn_mgr: _ConnectionManager,
+        session: Optional[ClientSession],
+        explicit_session: bool,
+    ) -> None:
+        """Cleanup a cursor from cursor.close() using a lock.
 
         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.
 
-        :param locks_allowed: True if we are allowed to acquire locks.
         :param cursor_id: The cursor id which may be 0.
         :param address: The _CursorAddress.
         :param conn_mgr: The _ConnectionManager for the pinned connection or None.
         :param session: The cursor's session.
         :param explicit_session: True if the session was passed explicitly.
         """
-        if locks_allowed:
-            if cursor_id:
-                if conn_mgr and conn_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.
-                    assert conn_mgr.conn is not None
-                    conn_mgr.conn.close_conn(ConnectionClosedReason.ERROR)
-                else:
-                    await self._close_cursor_now(
-                        cursor_id, address, session=session, conn_mgr=conn_mgr
-                    )
-            if conn_mgr:
-                await conn_mgr.close()
-        else:
-            # The cursor will be closed later in a different session.
-            if cursor_id or conn_mgr:
-                self._close_cursor_soon(cursor_id, address, conn_mgr)
+        if cursor_id:
+            if conn_mgr and conn_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.
+                assert conn_mgr.conn is not None
+                conn_mgr.conn.close_conn(ConnectionClosedReason.ERROR)
+            else:
+                await self._close_cursor_now(cursor_id, address, session=session, conn_mgr=conn_mgr)
+        if conn_mgr:
+            await conn_mgr.close()
         if session and not explicit_session:
-            await session._end_session(lock=locks_allowed)
+            session._end_implicit_session()
 
     async def _close_cursor_now(
         self,
@@ -1978,7 +1993,7 @@ async def _process_kill_cursors(self) -> None:
 
         for address, cursor_id, conn_mgr in pinned_cursors:
             try:
-                await self._cleanup_cursor(True, cursor_id, address, conn_mgr, None, False)
+                await self._cleanup_cursor_lock(cursor_id, address, conn_mgr, None, False)
             except Exception as exc:
                 if isinstance(exc, InvalidOperation) and self._topology._closed:
                     # Raise the exception when client is closed so that it

pymongo/synchronous/client_session.py

@@ -531,6 +531,12 @@ def _end_session(self, lock: bool) -> None:
                 self._client._return_server_session(self._server_session)
                 self._server_session = None
 
+    def _end_implicit_session(self) -> None:
+        # Implicit sessions can't be part of transactions or pinned connections
+        if self._server_session is not None:
+            self._client._return_server_session(self._server_session)
+            self._server_session = None
+
     def _check_ended(self) -> None:
         if self._server_session is None:
             raise InvalidOperation("Cannot use ended session")

pymongo/synchronous/command_cursor.py

@@ -81,8 +81,8 @@ def __init__(
         self._explicit_session = explicit_session
         self._killed = self._id == 0
         self._comment = comment
-        if _IS_SYNC and self._killed:
-            self._end_session(True)  # type: ignore[unused-coroutine]
+        if self._killed:
+            self._end_session()
 
         if "ns" in cursor_info:  # noqa: SIM401
             self._ns = cursor_info["ns"]
@@ -95,8 +95,7 @@ def __init__(
             raise TypeError("max_await_time_ms must be an integer or None")
 
     def __del__(self) -> None:
-        if _IS_SYNC:
-            self._die(False)  # type: ignore[unused-coroutine]
+        self._die_no_lock()
 
     def batch_size(self, batch_size: int) -> CommandCursor[_DocumentType]:
         """Limits the number of documents returned in one batch. Each batch
@@ -198,8 +197,7 @@ def session(self) -> Optional[ClientSession]:
             return self._session
         return None
 
-    def _die(self, synchronous: bool = False) -> None:
-        """Closes this cursor."""
+    def _prepare_to_die(self) -> tuple[int, Optional[_CursorAddress]]:
         already_killed = self._killed
         self._killed = True
         if self._id and not already_killed:
@@ -210,8 +208,22 @@ def _die(self, synchronous: bool = False) -> None:
             # Skip killCursors.
             cursor_id = 0
             address = None
-        self._collection.database.client._cleanup_cursor(
-            synchronous,
+        return cursor_id, address
+
+    def _die_no_lock(self) -> None:
+        """Closes this cursor without acquiring a lock."""
+        cursor_id, address = self._prepare_to_die()
+        self._collection.database.client._cleanup_cursor_no_lock(
+            cursor_id, address, self._sock_mgr, self._session, self._explicit_session
+        )
+        if not self._explicit_session:
+            self._session = None
+        self._sock_mgr = None
+
+    def _die_lock(self) -> None:
+        """Closes this cursor."""
+        cursor_id, address = self._prepare_to_die()
+        self._collection.database.client._cleanup_cursor_lock(
             cursor_id,
             address,
             self._sock_mgr,
@@ -222,14 +234,14 @@ def _die(self, synchronous: bool = False) -> None:
             self._session = None
         self._sock_mgr = None
 
-    def _end_session(self, synchronous: bool) -> None:
+    def _end_session(self) -> None:
         if self._session and not self._explicit_session:
-            self._session._end_session(lock=synchronous)
+            self._session._end_implicit_session()
             self._session = None
 
     def close(self) -> None:
         """Explicitly close / kill this cursor."""
-        self._die(True)
+        self._die_lock()
 
     def _send_message(self, operation: _GetMore) -> None:
         """Send a getmore message and handle the response."""
@@ -243,7 +255,7 @@ def _send_message(self, operation: _GetMore) -> None:
                 # Don't send killCursors because the cursor is already closed.
                 self._killed = True
             if exc.timeout:
-                self._die(False)
+                self._die_no_lock()
             else:
                 # Return the session and pinned connection, if necessary.
                 self.close()
@@ -305,7 +317,7 @@ def _refresh(self) -> int:
                 )
             )
         else:  # Cursor id is zero nothing else to return
-            self._die(True)
+            self._die_lock()
 
         return len(self._data)