Maintain connection reservations incrementally in the pool

src/httpcore2/httpcore2/_async/connection.py

@@ -177,6 +177,9 @@ def is_idle(self) -> bool:
             return self._connect_failed
         return self._connection.is_idle()
 
+    def can_multiplex(self) -> bool:
+        return self._connection is not None and self._connection.can_multiplex()
+
     def is_closed(self) -> bool:
         if self._connection is None:
             return self._connect_failed

src/httpcore2/httpcore2/_async/connection_pool.py

@@ -112,6 +112,11 @@ def __init__(
         self._connections: list[AsyncConnectionInterface] = []
         self._requests: list[AsyncPoolRequest] = []
 
+        # Reference counts of connections held by in-flight requests,
+        # maintained incrementally so assignment passes never rebuild them
+        # by scanning the full request list.
+        self._request_connections: dict[AsyncConnectionInterface, int] = {}
+
         # We only mutate the state of the connection pool within an 'optional_thread_lock'
         # context. This holds a threading lock unless we're running in async mode,
         # in which case it is a no-op.
@@ -227,14 +232,17 @@ async def handle_async_request(self, request: Request) -> Response:
                     # handle a request, but then become unavailable.
                     #
                     # In this case we clear the connection and try again.
-                    pool_request.clear_connection()
+                    with self._optional_thread_lock:
+                        self._release_request_connection(pool_request)
+                        pool_request.clear_connection()
                 else:
                     break  # pragma: no cover
 
         except BaseException as exc:
             with self._optional_thread_lock:
                 # For any exception or cancellation we remove the request from
                 # the queue, and then re-assign requests to connections.
+                self._release_request_connection(pool_request)
                 self._requests.remove(pool_request)
                 closing = self._assign_requests_to_connections()
 
@@ -251,6 +259,19 @@ async def handle_async_request(self, request: Request) -> Response:
             extensions=response.extensions,
         )
 
+    def _reserve_connection(self, pool_request: AsyncPoolRequest, connection: AsyncConnectionInterface) -> None:
+        pool_request.assign_to_connection(connection)
+        self._request_connections[connection] = self._request_connections.get(connection, 0) + 1
+
+    def _release_request_connection(self, pool_request: AsyncPoolRequest) -> None:
+        connection = pool_request.connection
+        if connection is not None:
+            count = self._request_connections[connection] - 1
+            if count:
+                self._request_connections[connection] = count
+            else:
+                del self._request_connections[connection]
+
     def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]:
         """
         Manage the state of the connection pool, assigning incoming
@@ -264,34 +285,40 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]:
         closing_connections: list[AsyncConnectionInterface] = []
         retained_connections: list[AsyncConnectionInterface] = []
 
-        # Connections currently referenced by an active request (including
-        # connections that are in the process of being established).
-        request_connections = {r.connection for r in self._requests}
+        # Connections currently referenced by an in-flight request, including
+        # connections that are in the process of being established and idle
+        # connections reserved by an assigned-but-not-yet-sent request.
+        request_connections = self._request_connections
 
         # First we handle cleaning up any connections that are closed
-        # or have expired their keep-alive, in a single pass.
+        # or have expired their keep-alive, in a single pass. Reserved
+        # connections skip the expiry check: they were checked when assigned,
+        # and `has_expired()` on an idle connection probes the socket.
         for connection in self._connections:
+            reserved = connection in request_connections
             if connection.is_closed():
                 continue
-            elif not (connection.is_connected() or connection in request_connections):
+            elif not (connection.is_connected() or reserved):
                 # Garbage: a NEW-state connection whose request was cancelled
                 # before the TCP handshake completed.  Drop it without closing
                 # (there is no socket to close yet).
                 continue
-            elif connection.has_expired():
+            elif not reserved and connection.has_expired():
                 closing_connections.append(connection)
             else:
                 retained_connections.append(connection)
 
         # Then we close any surplus idle connections, to enforce the
-        # max_keepalive_connections setting.
+        # max_keepalive_connections setting. Reserved connections are not
+        # surplus: a request is about to be sent on them.
         idle_surplus = (
-            sum(connection.is_idle() for connection in retained_connections) - self._max_keepalive_connections
+            sum(connection.is_idle() and connection not in request_connections for connection in retained_connections)
+            - self._max_keepalive_connections
         )
         if idle_surplus > 0:
             kept: list[AsyncConnectionInterface] = []
             for connection in retained_connections:
-                if idle_surplus > 0 and connection.is_idle():
+                if idle_surplus > 0 and connection.is_idle() and connection not in request_connections:
                     closing_connections.append(connection)
                     idle_surplus -= 1
                 else:
@@ -303,11 +330,27 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]:
         # Snapshot the set of reusable connections once, rather than rebuilding
         # it per queued request — this is what brings the loop from O(N*M) to
         # O(N+M) in the common case.
-        available_connections = [connection for connection in self._connections if connection.is_available()]
+        #
+        # An idle connection already assigned to an in-flight request is
+        # reserved: it stays IDLE until the winning task sends on it, so
+        # without this exclusion the next pass would assign it again and the
+        # loser would churn through `ConnectionNotAvailable`. Multiplexing
+        # connections are exempt: they can take further requests while idle.
+        available_connections = [
+            connection
+            for connection in self._connections
+            if connection.is_available()
+            and not (connection.is_idle() and connection in request_connections and not connection.can_multiplex())
+        ]
         new_connection_budget = self._max_connections - len(self._connections)
 
-        # Assign queued requests to connections.
+        # Assign queued requests to connections. Once no connection is
+        # available and no new connection may be created, no queued request
+        # can be assigned, so the scan stops early: this keeps a pass on a
+        # saturated pool O(connections) rather than O(in-flight requests).
         for pool_request in self._requests:
+            if not available_connections and new_connection_budget <= 0:
+                break
             if not pool_request.is_queued():
                 continue
             origin = pool_request.request.url.origin
@@ -318,15 +361,19 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]:
             # 2. We can create a new connection to handle the request.
             # 3. We can close an idle connection and then create a new connection
             #    to handle the request.
-            for connection in available_connections:
+            for idx, connection in enumerate(available_connections):
                 if connection.can_handle_request(origin):
-                    pool_request.assign_to_connection(connection)
+                    self._reserve_connection(pool_request, connection)
+                    if connection.is_idle() and not connection.can_multiplex():
+                        # An idle HTTP/1.1 connection can only take this
+                        # single request until it is released.
+                        del available_connections[idx]
                     break
             else:
                 if new_connection_budget > 0:
                     connection = self.create_connection(origin)
                     self._connections.append(connection)
-                    pool_request.assign_to_connection(connection)
+                    self._reserve_connection(pool_request, connection)
                     new_connection_budget -= 1
                     continue
                 for idx, connection in enumerate(available_connections):
@@ -336,7 +383,7 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]:
                         closing_connections.append(connection)
                         connection = self.create_connection(origin)
                         self._connections.append(connection)
-                        pool_request.assign_to_connection(connection)
+                        self._reserve_connection(pool_request, connection)
                         break
 
         return closing_connections
@@ -408,6 +455,7 @@ async def aclose(self) -> None:
                     await self._stream.aclose()
 
             with self._pool._optional_thread_lock:
+                self._pool._release_request_connection(self._pool_request)
                 self._pool._requests.remove(self._pool_request)
                 closing = self._pool._assign_requests_to_connections()
 

src/httpcore2/httpcore2/_async/http2.py

@@ -495,6 +495,9 @@ def has_expired(self) -> bool:
     def is_idle(self) -> bool:
         return self._state == HTTPConnectionState.IDLE
 
+    def can_multiplex(self) -> bool:
+        return True
+
     def is_closed(self) -> bool:
         return self._state == HTTPConnectionState.CLOSED
 

src/httpcore2/httpcore2/_async/interfaces.py

@@ -140,6 +140,16 @@ def is_idle(self) -> bool:
         """
         raise NotImplementedError()  # pragma: no cover
 
+    def can_multiplex(self) -> bool:
+        """
+        Return `True` if the connection can serve multiple requests
+        concurrently, such as an established HTTP/2 connection.
+
+        The default covers HTTP/1.1-style implementations, which serve a
+        single request at a time.
+        """
+        return False
+
     def is_closed(self) -> bool:
         """
         Return `True` if the connection has been closed.

src/httpcore2/httpcore2/_sync/connection.py

@@ -177,6 +177,9 @@ def is_idle(self) -> bool:
             return self._connect_failed
         return self._connection.is_idle()
 
+    def can_multiplex(self) -> bool:
+        return self._connection is not None and self._connection.can_multiplex()
+
     def is_closed(self) -> bool:
         if self._connection is None:
             return self._connect_failed