Avoid HTTP/1 connection assignment churn

.github/workflows/test-suite.yml

@@ -3,9 +3,9 @@ name: Test Suite
 
 on:
   push:
-    branches: ["master"]
+    branches: ["master", "openpipe/linear-connection-pool-1.0.9"]
   pull_request:
-    branches: ["master"]
+    branches: ["master", "openpipe/linear-connection-pool-1.0.9"]
 
 jobs:
   tests:

httpcore/_async/connection.py

@@ -184,6 +184,12 @@ def is_available(self) -> bool:
             )
         return self._connection.is_available()
 
+    def _is_multiplexable(self) -> bool:
+        if self._connection is None:
+            # Preserve speculative HTTP/2 sharing until the protocol is known.
+            return self.is_available()
+        return self._connection._is_multiplexable()
+
     def has_expired(self) -> bool:
         if self._connection is None:
             return self._connect_failed

httpcore/_async/connection_pool.py

@@ -4,6 +4,7 @@
 import sys
 import types
 import typing
+from collections import OrderedDict
 
 from .._backends.auto import AutoBackend
 from .._backends.base import SOCKET_OPTION, AsyncNetworkBackend
@@ -119,6 +120,7 @@ def __init__(
         # and the set of connections that are servicing those requests.
         self._connections: list[AsyncConnectionInterface] = []
         self._requests: list[AsyncPoolRequest] = []
+        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,
@@ -241,14 +243,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: nocover
 
         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()
 
@@ -267,6 +272,24 @@ 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:
+            pool_request.connection = 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
@@ -279,27 +302,36 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]:
         """
         closing_connections: list[AsyncConnectionInterface] = []
         retained_connections: list[AsyncConnectionInterface] = []
+        assigned = self._request_connections
 
         # First we handle cleaning up any connections that are closed
         # or have expired their keep-alive, in a single pass.
         for connection in self._connections:
             if connection.is_closed():
                 continue
-            elif connection.has_expired():
+            elif connection not in assigned 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.
         idle_surplus = (
-            sum(connection.is_idle() for connection in retained_connections)
+            sum(
+                connection.is_idle()
+                for connection in retained_connections
+                if connection not in assigned
+            )
             - 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 not in assigned
+                    and connection.is_idle()
+                ):
                     closing_connections.append(connection)
                     idle_surplus -= 1
                 else:
@@ -310,15 +342,18 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]:
 
         # Snapshot the reusable connections once instead of rebuilding the
         # list for every queued request.
-        available_connections = [
-            connection
+        available_connections = OrderedDict(
+            (id(connection), connection)
             for connection in self._connections
             if connection.is_available()
-        ]
+            and (connection not in assigned or connection._is_multiplexable())
+        )
         new_connection_budget = self._max_connections - len(self._connections)
 
         # Assign queued requests to connections.
         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
@@ -329,29 +364,31 @@ 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 connection in available_connections.values():
                 if connection.can_handle_request(origin):
-                    pool_request.assign_to_connection(connection)
+                    if not connection._is_multiplexable():
+                        del available_connections[id(connection)]
+                    self._reserve_connection(pool_request, connection)
                     break
             else:
                 if new_connection_budget > 0:
                     connection = self.create_connection(origin)
                     self._connections.append(connection)
-                    if connection.is_available():
-                        available_connections.append(connection)
-                    pool_request.assign_to_connection(connection)
+                    if connection.is_available() and connection._is_multiplexable():
+                        available_connections[id(connection)] = connection
+                    self._reserve_connection(pool_request, connection)
                     new_connection_budget -= 1
                     continue
-                for idx, connection in enumerate(available_connections):
-                    if connection.is_idle():
-                        del available_connections[idx]
+                for connection in available_connections.values():
+                    if connection not in assigned and connection.is_idle():
+                        del available_connections[id(connection)]
                         self._connections.remove(connection)
                         closing_connections.append(connection)
                         connection = self.create_connection(origin)
                         self._connections.append(connection)
-                        if connection.is_available():
-                            available_connections.append(connection)
-                        pool_request.assign_to_connection(connection)
+                        if connection.is_available() and connection._is_multiplexable():
+                            available_connections[id(connection)] = connection
+                        self._reserve_connection(pool_request, connection)
                         break
 
         return closing_connections
@@ -425,13 +462,18 @@ async def __aiter__(self) -> typing.AsyncIterator[bytes]:
             raise exc from None
 
     async def aclose(self) -> None:
-        if not self._closed:
+        with self._pool._optional_thread_lock:
+            if self._closed:
+                return
             self._closed = True
+
+        try:
             with AsyncShieldCancellation():
                 if hasattr(self._stream, "aclose"):
                     await self._stream.aclose()
-
+        finally:
             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()
 

httpcore/_async/http2.py

@@ -519,6 +519,9 @@ def is_available(self) -> bool:
             )
         )
 
+    def _is_multiplexable(self) -> bool:
+        return True
+
     def has_expired(self) -> bool:
         now = time.monotonic()
         return self._expire_at is not None and now > self._expire_at

httpcore/_async/http_proxy.py

@@ -217,6 +217,9 @@ def info(self) -> str:
     def is_available(self) -> bool:
         return self._connection.is_available()
 
+    def _is_multiplexable(self) -> bool:
+        return self._connection._is_multiplexable()
+
     def has_expired(self) -> bool:
         return self._connection.has_expired()
 
@@ -354,6 +357,9 @@ def info(self) -> str:
     def is_available(self) -> bool:
         return self._connection.is_available()
 
+    def _is_multiplexable(self) -> bool:
+        return self._connection._is_multiplexable()
+
     def has_expired(self) -> bool:
         return self._connection.has_expired()
 

httpcore/_async/interfaces.py

@@ -112,6 +112,10 @@ def is_available(self) -> bool:
         """
         raise NotImplementedError()  # pragma: nocover
 
+    def _is_multiplexable(self) -> bool:
+        """Whether multiple pool requests may hold an assignment at once."""
+        return False
+
     def has_expired(self) -> bool:
         """
         Return `True` if the connection is in a state where it should be closed.

httpcore/_async/socks_proxy.py

@@ -317,6 +317,11 @@ def is_available(self) -> bool:
             )
         return self._connection.is_available()
 
+    def _is_multiplexable(self) -> bool:
+        if self._connection is None:
+            return self.is_available()
+        return self._connection._is_multiplexable()
+
     def has_expired(self) -> bool:
         if self._connection is None:  # pragma: nocover
             return self._connect_failed

httpcore/_sync/connection.py

@@ -184,6 +184,12 @@ def is_available(self) -> bool:
             )
         return self._connection.is_available()
 
+    def _is_multiplexable(self) -> bool:
+        if self._connection is None:
+            # Preserve speculative HTTP/2 sharing until the protocol is known.
+            return self.is_available()
+        return self._connection._is_multiplexable()
+
     def has_expired(self) -> bool:
         if self._connection is None:
             return self._connect_failed