Refactor connection waiters to be cancellation safe (#9671)

bdraco · web-flow · commit 37d9fe68387e · 2024-11-05T10:12:44.000-06:00
diff --git a/CHANGES/9670.bugfix.rst b/CHANGES/9670.bugfix.rst
@@ -0,0 +1 @@
+9671.bugfix.rst
diff --git a/CHANGES/9671.bugfix.rst b/CHANGES/9671.bugfix.rst
@@ -0,0 +1,3 @@
+Fixed a deadlock that could occur while attempting to get a new connection slot after a timeout -- by :user:`bdraco`.
+
+The connector was not cancellation-safe.
diff --git a/aiohttp/connector.py b/aiohttp/connector.py
@@ -6,7 +6,7 @@
 import sys
 import traceback
 import warnings
-from collections import defaultdict, deque
+from collections import OrderedDict, defaultdict
 from contextlib import suppress
 from http import HTTPStatus
 from itertools import chain, cycle, islice
@@ -252,9 +252,11 @@ def __init__(
         self._force_close = force_close
 
         # {host_key: FIFO list of waiters}
-        self._waiters: DefaultDict[ConnectionKey, deque[asyncio.Future[None]]] = (
-            defaultdict(deque)
-        )
+        # The FIFO is implemented with an OrderedDict with None keys because
+        # python does not have an ordered set.
+        self._waiters: DefaultDict[
+            ConnectionKey, OrderedDict[asyncio.Future[None], None]
+        ] = defaultdict(OrderedDict)
 
         self._loop = loop
         self._factory = functools.partial(ResponseHandler, loop=loop)
@@ -335,7 +337,7 @@ def _cleanup(self) -> None:
             # recreate it ever!
             self._cleanup_handle = None
 
-        now = self._loop.time()
+        now = monotonic()
         timeout = self._keepalive_timeout
 
         if self._conns:
@@ -366,14 +368,6 @@ def _cleanup(self) -> None:
                 timeout_ceil_threshold=self._timeout_ceil_threshold,
             )
 
-    def _drop_acquired_per_host(
-        self, key: "ConnectionKey", val: ResponseHandler
-    ) -> None:
-        if conns := self._acquired_per_host.get(key):
-            conns.remove(val)
-            if not conns:
-                del self._acquired_per_host[key]
-
     def _cleanup_closed(self) -> None:
         """Double confirmation for transport close.
 
@@ -446,6 +440,9 @@ def _close_immediately(self) -> List[Awaitable[object]]:
         finally:
             self._conns.clear()
             self._acquired.clear()
+            for keyed_waiters in self._waiters.values():
+                for keyed_waiter in keyed_waiters:
+                    keyed_waiter.cancel()
             self._waiters.clear()
             self._cleanup_handle = None
             self._cleanup_closed_transports.clear()
@@ -489,117 +486,109 @@ async def connect(
     ) -> Connection:
         """Get from pool or create new connection."""
         key = req.connection_key
-        available = self._available_connections(key)
-        wait_for_conn = available <= 0 or key in self._waiters
-        if not wait_for_conn and (proto := self._get(key)) is not None:
+        if (conn := await self._get(key, traces)) is not None:
             # If we do not have to wait and we can get a connection from the pool
             # we can avoid the timeout ceil logic and directly return the connection
-            return await self._reused_connection(key, proto, traces)
+            return conn
 
         async with ceil_timeout(timeout.connect, timeout.ceil_threshold):
-            # Wait if there are no available connections or if there are/were
-            # waiters (i.e. don't steal connection from a waiter about to wake up)
-            if wait_for_conn:
+            if self._available_connections(key) <= 0:
                 await self._wait_for_available_connection(key, traces)
-                if (proto := self._get(key)) is not None:
-                    return await self._reused_connection(key, proto, traces)
+                if (conn := await self._get(key, traces)) is not None:
+                    return conn
 
             placeholder = cast(
                 ResponseHandler, _TransportPlaceholder(self._placeholder_future)
             )
             self._acquired.add(placeholder)
             self._acquired_per_host[key].add(placeholder)
 
-            if traces:
-                for trace in traces:
-                    await trace.send_connection_create_start()
-
             try:
+                # Traces are done inside the try block to ensure that the
+                # that the placeholder is still cleaned up if an exception
+                # is raised.
+                if traces:
+                    for trace in traces:
+                        await trace.send_connection_create_start()
                 proto = await self._create_connection(req, traces, timeout)
-                if self._closed:
-                    proto.close()
-                    raise ClientConnectionError("Connector is closed.")
+                if traces:
+                    for trace in traces:
+                        await trace.send_connection_create_end()
             except BaseException:
-                if not self._closed:
-                    self._acquired.remove(placeholder)
-                    self._drop_acquired_per_host(key, placeholder)
-                    self._release_waiter()
+                self._release_acquired(key, placeholder)
                 raise
             else:
-                if not self._closed:
-                    self._acquired.remove(placeholder)
-                    self._drop_acquired_per_host(key, placeholder)
-
-            if traces:
-                for trace in traces:
-                    await trace.send_connection_create_end()
-
-            return self._acquired_connection(proto, key)
-
-    async def _reused_connection(
-        self, key: "ConnectionKey", proto: ResponseHandler, traces: List["Trace"]
-    ) -> Connection:
-        if traces:
-            # Acquire the connection to prevent race conditions with limits
-            placeholder = cast(
-                ResponseHandler, _TransportPlaceholder(self._placeholder_future)
-            )
-            self._acquired.add(placeholder)
-            self._acquired_per_host[key].add(placeholder)
-            for trace in traces:
-                await trace.send_connection_reuseconn()
-            self._acquired.remove(placeholder)
-            self._drop_acquired_per_host(key, placeholder)
-        return self._acquired_connection(proto, key)
+                if self._closed:
+                    proto.close()
+                    raise ClientConnectionError("Connector is closed.")
 
-    def _acquired_connection(
-        self, proto: ResponseHandler, key: "ConnectionKey"
-    ) -> Connection:
-        """Mark proto as acquired and wrap it in a Connection object."""
+        # The connection was successfully created, drop the placeholder
+        # and add the real connection to the acquired set. There should
+        # be no awaits after the proto is added to the acquired set
+        # to ensure that the connection is not left in the acquired set
+        # on cancellation.
+        acquired_per_host = self._acquired_per_host[key]
+        self._acquired.remove(placeholder)
+        acquired_per_host.remove(placeholder)
         self._acquired.add(proto)
-        self._acquired_per_host[key].add(proto)
+        acquired_per_host.add(proto)
         return Connection(self, key, proto, self._loop)
 
     async def _wait_for_available_connection(
         self, key: "ConnectionKey", traces: List["Trace"]
     ) -> None:
-        """Wait until there is an available connection."""
-        fut: asyncio.Future[None] = self._loop.create_future()
-
-        # This connection will now count towards the limit.
-        self._waiters[key].append(fut)
+        """Wait for an available connection slot."""
+        # We loop here because there is a race between
+        # the connection limit check and the connection
+        # being acquired. If the connection is acquired
+        # between the check and the await statement, we
+        # need to loop again to check if the connection
+        # slot is still available.
+        attempts = 0
+        while True:
+            fut: asyncio.Future[None] = self._loop.create_future()
+            keyed_waiters = self._waiters[key]
+            keyed_waiters[fut] = None
+            if attempts:
+                # If we have waited before, we need to move the waiter
+                # to the front of the queue as otherwise we might get
+                # starved and hit the timeout.
+                keyed_waiters.move_to_end(fut, last=False)
 
-        if traces:
-            for trace in traces:
-                await trace.send_connection_queued_start()
+            try:
+                # Traces happen in the try block to ensure that the
+                # the waiter is still cleaned up if an exception is raised.
+                if traces:
+                    for trace in traces:
+                        await trace.send_connection_queued_start()
+                await fut
+                if traces:
+                    for trace in traces:
+                        await trace.send_connection_queued_end()
+            finally:
+                # pop the waiter from the queue if its still
+                # there and not already removed by _release_waiter
+                keyed_waiters.pop(fut, None)
+                if not self._waiters.get(key, True):
+                    del self._waiters[key]
 
-        try:
-            await fut
-        except BaseException as e:
-            if key in self._waiters:
-                # remove a waiter even if it was cancelled, normally it's
-                #  removed when it's notified
-                with suppress(ValueError):
-                    # fut may no longer be in list
-                    self._waiters[key].remove(fut)
-
-            raise e
-        finally:
-            if key in self._waiters and not self._waiters[key]:
-                del self._waiters[key]
+            if self._available_connections(key) > 0:
+                break
+            attempts += 1
 
-        if traces:
-            for trace in traces:
-                await trace.send_connection_queued_end()
+    async def _get(
+        self, key: "ConnectionKey", traces: List["Trace"]
+    ) -> Optional[Connection]:
+        """Get next reusable connection for the key or None.
 
-    def _get(self, key: "ConnectionKey") -> Optional[ResponseHandler]:
-        """Get next reusable connection for the key or None."""
+        The connection will be marked as acquired.
+        """
         try:
             conns = self._conns[key]
         except KeyError:
             return None
 
-        t1 = self._loop.time()
+        t1 = monotonic()
         while conns:
             proto, t0 = conns.pop()
             # We will we reuse the connection if its connected and
@@ -608,7 +597,16 @@ def _get(self, key: "ConnectionKey") -> Optional[ResponseHandler]:
                 if not conns:
                     # The very last connection was reclaimed: drop the key
                     del self._conns[key]
-                return proto
+                self._acquired.add(proto)
+                self._acquired_per_host[key].add(proto)
+                if traces:
+                    for trace in traces:
+                        try:
+                            await trace.send_connection_reuseconn()
+                        except BaseException:
+                            self._release_acquired(key, proto)
+                            raise
+                return Connection(self, key, proto, self._loop)
 
             # Connection cannot be reused, close it
             transport = proto.transport
@@ -642,25 +640,23 @@ def _release_waiter(self) -> None:
 
             waiters = self._waiters[key]
             while waiters:
-                waiter = waiters.popleft()
+                waiter, _ = waiters.popitem(last=False)
                 if not waiter.done():
                     waiter.set_result(None)
                     return
 
     def _release_acquired(self, key: "ConnectionKey", proto: ResponseHandler) -> None:
+        """Release acquired connection."""
         if self._closed:
             # acquired connection is already released on connector closing
             return
 
-        try:
-            self._acquired.remove(proto)
-            self._drop_acquired_per_host(key, proto)
-        except KeyError:  # pragma: no cover
-            # this may be result of undetermenistic order of objects
-            # finalization due garbage collection.
-            pass
-        else:
-            self._release_waiter()
+        self._acquired.discard(proto)
+        if conns := self._acquired_per_host.get(key):
+            conns.discard(proto)
+            if not conns:
+                del self._acquired_per_host[key]
+        self._release_waiter()
 
     def _release(
         self,
@@ -692,7 +688,7 @@ def _release(
         conns = self._conns.get(key)
         if conns is None:
             conns = self._conns[key] = []
-        conns.append((protocol, self._loop.time()))
+        conns.append((protocol, monotonic()))
 
         if self._cleanup_handle is None:
             self._cleanup_handle = helpers.weakref_handle(
diff --git a/tests/test_connector.py b/tests/test_connector.py

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	+Fixed a deadlock that could occur while attempting to get a new connection slot after a timeout -- by :user:`bdraco`.
	`2`	`+`
	`3`	`+The connector was not cancellation-safe.`