PYTHON-5212 - Do not hold Topology lock while resetting pool

Author: NoahStapp

pymongo/asynchronous/pool.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import collections
 import contextlib
 import logging
@@ -860,8 +861,13 @@ async def _reset(
         # PoolClosedEvent but that reset() SHOULD close sockets *after*
         # publishing the PoolClearedEvent.
         if close:
-            for conn in sockets:
-                await conn.close_conn(ConnectionClosedReason.POOL_CLOSED)
+            if not _IS_SYNC:
+                await asyncio.gather(
+                    *[conn._close_conn(ConnectionClosedReason.POOL_CLOSED) for conn in sockets]
+                )
+            else:
+                for conn in sockets:
+                    await conn.close_conn(ConnectionClosedReason.POOL_CLOSED)
             if self.enabled_for_cmap:
                 assert listeners is not None
                 listeners.publish_pool_closed(self.address)
@@ -891,8 +897,13 @@ async def _reset(
                         serverPort=self.address[1],
                         serviceId=service_id,
                     )
-            for conn in sockets:
-                await conn.close_conn(ConnectionClosedReason.STALE)
+            if not _IS_SYNC:
+                await asyncio.gather(
+                    *[conn._close_conn(ConnectionClosedReason.STALE) for conn in sockets]
+                )
+            else:
+                for conn in sockets:
+                    await conn.close_conn(ConnectionClosedReason.STALE)
 
     async def update_is_writable(self, is_writable: Optional[bool]) -> None:
         """Updates the is_writable attribute on all sockets currently in the
@@ -938,8 +949,13 @@ async def remove_stale_sockets(self, reference_generation: int) -> None:
                     and self.conns[-1].idle_time_seconds() > self.opts.max_idle_time_seconds
                 ):
                     close_conns.append(self.conns.pop())
-            for conn in close_conns:
-                await conn.close_conn(ConnectionClosedReason.IDLE)
+            if not _IS_SYNC:
+                await asyncio.gather(
+                    *[conn._close_conn(ConnectionClosedReason.IDLE) for conn in close_conns]
+                )
+            else:
+                for conn in close_conns:
+                    await conn.close_conn(ConnectionClosedReason.IDLE)
 
         while True:
             async with self.size_cond:

pymongo/asynchronous/topology.py

@@ -529,12 +529,6 @@ async def _process_change(
             if not _IS_SYNC:
                 self._monitor_tasks.append(self._srv_monitor)
 
-        # Clear the pool from a failed heartbeat.
-        if reset_pool:
-            server = self._servers.get(server_description.address)
-            if server:
-                await server.pool.reset(interrupt_connections=interrupt_connections)
-
         # Wake anything waiting in select_servers().
         self._condition.notify_all()
 
@@ -557,6 +551,11 @@ async def on_change(
             # that didn't include this server.
             if self._opened and self._description.has_server(server_description.address):
                 await self._process_change(server_description, reset_pool, interrupt_connections)
+        # Clear the pool from a failed heartbeat, done outside the lock to avoid blocking on connection close.
+        if self._opened and self._description.has_server(server_description.address) and reset_pool:
+            server = self._servers.get(server_description.address)
+            if server:
+                await server.pool.reset(interrupt_connections=interrupt_connections)
 
     async def _process_srv_update(self, seedlist: list[tuple[str, Any]]) -> None:
         """Process a new seedlist on an opened topology.

pymongo/synchronous/pool.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import collections
 import contextlib
 import logging
@@ -858,8 +859,13 @@ def _reset(
         # PoolClosedEvent but that reset() SHOULD close sockets *after*
         # publishing the PoolClearedEvent.
         if close:
-            for conn in sockets:
-                conn.close_conn(ConnectionClosedReason.POOL_CLOSED)
+            if not _IS_SYNC:
+                asyncio.gather(
+                    *[conn._close_conn(ConnectionClosedReason.POOL_CLOSED) for conn in sockets]
+                )
+            else:
+                for conn in sockets:
+                    conn.close_conn(ConnectionClosedReason.POOL_CLOSED)
             if self.enabled_for_cmap:
                 assert listeners is not None
                 listeners.publish_pool_closed(self.address)
@@ -889,8 +895,13 @@ def _reset(
                         serverPort=self.address[1],
                         serviceId=service_id,
                     )
-            for conn in sockets:
-                conn.close_conn(ConnectionClosedReason.STALE)
+            if not _IS_SYNC:
+                asyncio.gather(
+                    *[conn._close_conn(ConnectionClosedReason.STALE) for conn in sockets]
+                )
+            else:
+                for conn in sockets:
+                    conn.close_conn(ConnectionClosedReason.STALE)
 
     def update_is_writable(self, is_writable: Optional[bool]) -> None:
         """Updates the is_writable attribute on all sockets currently in the
@@ -934,8 +945,13 @@ def remove_stale_sockets(self, reference_generation: int) -> None:
                     and self.conns[-1].idle_time_seconds() > self.opts.max_idle_time_seconds
                 ):
                     close_conns.append(self.conns.pop())
-            for conn in close_conns:
-                conn.close_conn(ConnectionClosedReason.IDLE)
+            if not _IS_SYNC:
+                asyncio.gather(
+                    *[conn._close_conn(ConnectionClosedReason.IDLE) for conn in close_conns]
+                )
+            else:
+                for conn in close_conns:
+                    conn.close_conn(ConnectionClosedReason.IDLE)
 
         while True:
             with self.size_cond:

pymongo/synchronous/topology.py

@@ -529,12 +529,6 @@ def _process_change(
             if not _IS_SYNC:
                 self._monitor_tasks.append(self._srv_monitor)
 
-        # Clear the pool from a failed heartbeat.
-        if reset_pool:
-            server = self._servers.get(server_description.address)
-            if server:
-                server.pool.reset(interrupt_connections=interrupt_connections)
-
         # Wake anything waiting in select_servers().
         self._condition.notify_all()
 
@@ -557,6 +551,11 @@ def on_change(
             # that didn't include this server.
             if self._opened and self._description.has_server(server_description.address):
                 self._process_change(server_description, reset_pool, interrupt_connections)
+        # Clear the pool from a failed heartbeat, done outside the lock to avoid blocking on connection close.
+        if self._opened and self._description.has_server(server_description.address) and reset_pool:
+            server = self._servers.get(server_description.address)
+            if server:
+                server.pool.reset(interrupt_connections=interrupt_connections)
 
     def _process_srv_update(self, seedlist: list[tuple[str, Any]]) -> None:
         """Process a new seedlist on an opened topology.