From a3e1a2bbc27c67e752a1142484360d6a058a70df Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 23 Jul 2026 09:18:30 +0200 Subject: [PATCH 1/6] Assign each released connection to a single queued request Assignment does not change an HTTP/1.1 connection's state, so a newly idle connection was handed to every queued request in one pass and re-picked by later passes until the winner sent on it. Every loser woke up, failed with ConnectionNotAvailable, re-entered the queue, and triggered another full O(n) assignment pass - quadratic churn at high queue depth. Drop a connection from the availability snapshot once assigned, and exclude idle connections already reserved by an in-flight request when building the snapshot. 1000 concurrent requests against a local server drop from 5.1s to 0.9s with the default pool, and from 100s to 1.1s with max_connections=1. HTTP/2 multiplexing is unaffected: an active h2 connection is not idle, so it stays available to additional streams. --- .../httpcore2/_async/connection_pool.py | 17 +++++++- .../httpcore2/_sync/connection_pool.py | 17 +++++++- .../httpcore2/_async/test_connection_pool.py | 43 +++++++++++++++++++ tests/httpcore2/_sync/test_connection_pool.py | 43 +++++++++++++++++++ 4 files changed, 116 insertions(+), 4 deletions(-) diff --git a/src/httpcore2/httpcore2/_async/connection_pool.py b/src/httpcore2/httpcore2/_async/connection_pool.py index 7011176e..20b72cb5 100644 --- a/src/httpcore2/httpcore2/_async/connection_pool.py +++ b/src/httpcore2/httpcore2/_async/connection_pool.py @@ -303,7 +303,16 @@ 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`. + available_connections = [ + connection + for connection in self._connections + if connection.is_available() and not (connection.is_idle() and connection in request_connections) + ] new_connection_budget = self._max_connections - len(self._connections) # Assign queued requests to connections. @@ -318,9 +327,13 @@ 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) + if connection.is_idle(): + # An HTTP/1.1 connection (or an idle HTTP/2 one) can + # only take this single request until it is released. + del available_connections[idx] break else: if new_connection_budget > 0: diff --git a/src/httpcore2/httpcore2/_sync/connection_pool.py b/src/httpcore2/httpcore2/_sync/connection_pool.py index 287d9fcf..f09b7146 100644 --- a/src/httpcore2/httpcore2/_sync/connection_pool.py +++ b/src/httpcore2/httpcore2/_sync/connection_pool.py @@ -303,7 +303,16 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]: # 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`. + available_connections = [ + connection + for connection in self._connections + if connection.is_available() and not (connection.is_idle() and connection in request_connections) + ] new_connection_budget = self._max_connections - len(self._connections) # Assign queued requests to connections. @@ -318,9 +327,13 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]: # 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) + if connection.is_idle(): + # An HTTP/1.1 connection (or an idle HTTP/2 one) can + # only take this single request until it is released. + del available_connections[idx] break else: if new_connection_budget > 0: diff --git a/tests/httpcore2/_async/test_connection_pool.py b/tests/httpcore2/_async/test_connection_pool.py index 832c9c1f..2e38892a 100644 --- a/tests/httpcore2/_async/test_connection_pool.py +++ b/tests/httpcore2/_async/test_connection_pool.py @@ -789,3 +789,46 @@ async def trace(name: str, kwargs: dict[str, typing.Any]) -> None: "http11.response_closed.started", "http11.response_closed.complete", ] + + +@pytest.mark.trio +async def test_connection_pool_assigns_released_connection_to_one_queued_request() -> None: + """ + A released connection must be handed to exactly one queued request. + + Assigning it to every queued request wakes them all, only for all but one + to fail with `ConnectionNotAvailable` and re-enter the queue, degrading + quadratically with queue depth. + """ + + class CountingPool(httpcore2.AsyncConnectionPool): + assign_passes = 0 + + def _assign_requests_to_connections(self) -> list[httpcore2.AsyncConnectionInterface]: + CountingPool.assign_passes += 1 + return super()._assign_requests_to_connections() + + network_backend = httpcore2.AsyncMockBackend( + [ + b"HTTP/1.1 200 OK\r\n", + b"Content-Type: plain/text\r\n", + b"Content-Length: 13\r\n", + b"\r\n", + b"Hello, world!", + ] + * 10 + ) + + async def fetch(pool: httpcore2.AsyncConnectionPool) -> None: + async with pool.stream("GET", "https://example.com/") as response: + await response.aread() + assert response.status == 200 + + async with CountingPool(max_connections=1, network_backend=network_backend) as pool: + async with concurrency.open_nursery() as nursery: + for _ in range(10): + nursery.start_soon(fetch, pool) + + # Exactly two passes per request: one when it is queued, one when it + # releases its connection. + assert CountingPool.assign_passes == 2 * 10 diff --git a/tests/httpcore2/_sync/test_connection_pool.py b/tests/httpcore2/_sync/test_connection_pool.py index 137d4fd8..022341b9 100644 --- a/tests/httpcore2/_sync/test_connection_pool.py +++ b/tests/httpcore2/_sync/test_connection_pool.py @@ -789,3 +789,46 @@ def trace(name: str, kwargs: dict[str, typing.Any]) -> None: "http11.response_closed.started", "http11.response_closed.complete", ] + + + +def test_connection_pool_assigns_released_connection_to_one_queued_request() -> None: + """ + A released connection must be handed to exactly one queued request. + + Assigning it to every queued request wakes them all, only for all but one + to fail with `ConnectionNotAvailable` and re-enter the queue, degrading + quadratically with queue depth. + """ + + class CountingPool(httpcore2.ConnectionPool): + assign_passes = 0 + + def _assign_requests_to_connections(self) -> list[httpcore2.ConnectionInterface]: + CountingPool.assign_passes += 1 + return super()._assign_requests_to_connections() + + network_backend = httpcore2.MockBackend( + [ + b"HTTP/1.1 200 OK\r\n", + b"Content-Type: plain/text\r\n", + b"Content-Length: 13\r\n", + b"\r\n", + b"Hello, world!", + ] + * 10 + ) + + def fetch(pool: httpcore2.ConnectionPool) -> None: + with pool.stream("GET", "https://example.com/") as response: + response.read() + assert response.status == 200 + + with CountingPool(max_connections=1, network_backend=network_backend) as pool: + with concurrency.open_nursery() as nursery: + for _ in range(10): + nursery.start_soon(fetch, pool) + + # Exactly two passes per request: one when it is queued, one when it + # releases its connection. + assert CountingPool.assign_passes == 2 * 10 From b9f33ba5f886a2e947f6c9d93efacdd7aa3def93 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 23 Jul 2026 09:42:07 +0200 Subject: [PATCH 2/6] Exempt multiplexing connections from single-assignment reservation An idle HTTP/2 connection can serve further requests while reserved, so treating it like HTTP/1.1 could leave a queued burst waiting for the next pool event instead of multiplexing. Add can_multiplex() to the connection interface (False by default, True for established HTTP/2) and only apply the reserved-idle exclusion to connections that cannot multiplex. --- src/httpcore2/httpcore2/_async/connection.py | 3 ++ .../httpcore2/_async/connection_pool.py | 12 +++-- src/httpcore2/httpcore2/_async/http2.py | 3 ++ src/httpcore2/httpcore2/_async/interfaces.py | 10 ++++ src/httpcore2/httpcore2/_sync/connection.py | 3 ++ .../httpcore2/_sync/connection_pool.py | 12 +++-- src/httpcore2/httpcore2/_sync/http2.py | 3 ++ src/httpcore2/httpcore2/_sync/interfaces.py | 10 ++++ .../httpcore2/_async/test_connection_pool.py | 52 +++++++++++++++++++ tests/httpcore2/_sync/test_connection_pool.py | 52 +++++++++++++++++++ 10 files changed, 150 insertions(+), 10 deletions(-) diff --git a/src/httpcore2/httpcore2/_async/connection.py b/src/httpcore2/httpcore2/_async/connection.py index ff3509bd..e3fb8596 100644 --- a/src/httpcore2/httpcore2/_async/connection.py +++ b/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 diff --git a/src/httpcore2/httpcore2/_async/connection_pool.py b/src/httpcore2/httpcore2/_async/connection_pool.py index 20b72cb5..e254fd41 100644 --- a/src/httpcore2/httpcore2/_async/connection_pool.py +++ b/src/httpcore2/httpcore2/_async/connection_pool.py @@ -307,11 +307,13 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]: # 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`. + # 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) + 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) @@ -330,9 +332,9 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]: for idx, connection in enumerate(available_connections): if connection.can_handle_request(origin): pool_request.assign_to_connection(connection) - if connection.is_idle(): - # An HTTP/1.1 connection (or an idle HTTP/2 one) can - # only take this single request until it is released. + 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: diff --git a/src/httpcore2/httpcore2/_async/http2.py b/src/httpcore2/httpcore2/_async/http2.py index 6f2b4b35..7d2dab5d 100644 --- a/src/httpcore2/httpcore2/_async/http2.py +++ b/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 diff --git a/src/httpcore2/httpcore2/_async/interfaces.py b/src/httpcore2/httpcore2/_async/interfaces.py index f394d843..008d295d 100644 --- a/src/httpcore2/httpcore2/_async/interfaces.py +++ b/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. diff --git a/src/httpcore2/httpcore2/_sync/connection.py b/src/httpcore2/httpcore2/_sync/connection.py index 6c213cb4..5634ce99 100644 --- a/src/httpcore2/httpcore2/_sync/connection.py +++ b/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 diff --git a/src/httpcore2/httpcore2/_sync/connection_pool.py b/src/httpcore2/httpcore2/_sync/connection_pool.py index f09b7146..20912e5c 100644 --- a/src/httpcore2/httpcore2/_sync/connection_pool.py +++ b/src/httpcore2/httpcore2/_sync/connection_pool.py @@ -307,11 +307,13 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]: # 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`. + # 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) + 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) @@ -330,9 +332,9 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]: for idx, connection in enumerate(available_connections): if connection.can_handle_request(origin): pool_request.assign_to_connection(connection) - if connection.is_idle(): - # An HTTP/1.1 connection (or an idle HTTP/2 one) can - # only take this single request until it is released. + 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: diff --git a/src/httpcore2/httpcore2/_sync/http2.py b/src/httpcore2/httpcore2/_sync/http2.py index b0c42c71..992b42ce 100644 --- a/src/httpcore2/httpcore2/_sync/http2.py +++ b/src/httpcore2/httpcore2/_sync/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 diff --git a/src/httpcore2/httpcore2/_sync/interfaces.py b/src/httpcore2/httpcore2/_sync/interfaces.py index bbe7c7e6..0e66bd1e 100644 --- a/src/httpcore2/httpcore2/_sync/interfaces.py +++ b/src/httpcore2/httpcore2/_sync/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. diff --git a/tests/httpcore2/_async/test_connection_pool.py b/tests/httpcore2/_async/test_connection_pool.py index 2e38892a..11b9a5d4 100644 --- a/tests/httpcore2/_async/test_connection_pool.py +++ b/tests/httpcore2/_async/test_connection_pool.py @@ -832,3 +832,55 @@ async def fetch(pool: httpcore2.AsyncConnectionPool) -> None: # Exactly two passes per request: one when it is queued, one when it # releases its connection. assert CountingPool.assign_passes == 2 * 10 + + +@pytest.mark.trio +async def test_connection_pool_multiplexes_idle_http2_connection_within_a_pass() -> None: + """ + A burst of requests arriving while a warmed HTTP/2 connection is idle + must be assigned to it immediately, not serialized behind the first + request's reservation. + """ + + class QueueObservingPool(httpcore2.AsyncConnectionPool): + max_queued_after_pass = 0 + + def _assign_requests_to_connections(self) -> list[httpcore2.AsyncConnectionInterface]: + closing = super()._assign_requests_to_connections() + queued = sum(request.is_queued() for request in self._requests) + QueueObservingPool.max_queued_after_pass = max(QueueObservingPool.max_queued_after_pass, queued) + return closing + + def response_frames(stream_id: int) -> list[bytes]: + return [ + hyperframe.frame.HeadersFrame( + stream_id=stream_id, + data=hpack.Encoder().encode([(b":status", b"200")]), + flags=["END_HEADERS"], + ).serialize(), + hyperframe.frame.DataFrame(stream_id=stream_id, data=b"Hello, world!", flags=["END_STREAM"]).serialize(), + ] + + network_backend = httpcore2.AsyncMockBackend( + buffer=[ + hyperframe.frame.SettingsFrame().serialize(), + *response_frames(1), + *response_frames(3), + *response_frames(5), + *response_frames(7), + ], + http2=True, + ) + + async def fetch(pool: httpcore2.AsyncConnectionPool) -> None: + response = await pool.request("GET", "https://example.com/") + assert response.status == 200 + + async with QueueObservingPool(network_backend=network_backend, max_connections=1, http2=True) as pool: + # Warm the connection; it returns to the pool IDLE. + await fetch(pool) + async with concurrency.open_nursery() as nursery: + for _ in range(3): + nursery.start_soon(fetch, pool) + + assert QueueObservingPool.max_queued_after_pass == 0 diff --git a/tests/httpcore2/_sync/test_connection_pool.py b/tests/httpcore2/_sync/test_connection_pool.py index 022341b9..693b7371 100644 --- a/tests/httpcore2/_sync/test_connection_pool.py +++ b/tests/httpcore2/_sync/test_connection_pool.py @@ -832,3 +832,55 @@ def fetch(pool: httpcore2.ConnectionPool) -> None: # Exactly two passes per request: one when it is queued, one when it # releases its connection. assert CountingPool.assign_passes == 2 * 10 + + + +def test_connection_pool_multiplexes_idle_http2_connection_within_a_pass() -> None: + """ + A burst of requests arriving while a warmed HTTP/2 connection is idle + must be assigned to it immediately, not serialized behind the first + request's reservation. + """ + + class QueueObservingPool(httpcore2.ConnectionPool): + max_queued_after_pass = 0 + + def _assign_requests_to_connections(self) -> list[httpcore2.ConnectionInterface]: + closing = super()._assign_requests_to_connections() + queued = sum(request.is_queued() for request in self._requests) + QueueObservingPool.max_queued_after_pass = max(QueueObservingPool.max_queued_after_pass, queued) + return closing + + def response_frames(stream_id: int) -> list[bytes]: + return [ + hyperframe.frame.HeadersFrame( + stream_id=stream_id, + data=hpack.Encoder().encode([(b":status", b"200")]), + flags=["END_HEADERS"], + ).serialize(), + hyperframe.frame.DataFrame(stream_id=stream_id, data=b"Hello, world!", flags=["END_STREAM"]).serialize(), + ] + + network_backend = httpcore2.MockBackend( + buffer=[ + hyperframe.frame.SettingsFrame().serialize(), + *response_frames(1), + *response_frames(3), + *response_frames(5), + *response_frames(7), + ], + http2=True, + ) + + def fetch(pool: httpcore2.ConnectionPool) -> None: + response = pool.request("GET", "https://example.com/") + assert response.status == 200 + + with QueueObservingPool(network_backend=network_backend, max_connections=1, http2=True) as pool: + # Warm the connection; it returns to the pool IDLE. + fetch(pool) + with concurrency.open_nursery() as nursery: + for _ in range(3): + nursery.start_soon(fetch, pool) + + assert QueueObservingPool.max_queued_after_pass == 0 From 0ed00a0b5130a6ca1a24aac2f8ce7d89d9c27546 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 23 Jul 2026 10:43:42 +0200 Subject: [PATCH 3/6] Stop scanning the request queue once no connection can be assigned Each assignment pass walked every in-flight request even when the pool was saturated, and re-probed reserved idle connections for expiry with an is_readable socket check on every interleaved pass. Break out of the assignment loop once no connection is available and no new one may be created, and skip expiry checks and surplus-keepalive eviction for connections reserved by an assigned request - they were health-checked at assignment time, and evicting them would hand the winning request a closed connection. 1000 unbounded concurrent requests against a local server now complete in 0.41s versus 0.51s sequential, compared to 0.93s before this change and 5.1s before #1075. --- .../httpcore2/_async/connection_pool.py | 29 +++++++++++++------ .../httpcore2/_sync/connection_pool.py | 29 +++++++++++++------ 2 files changed, 40 insertions(+), 18 deletions(-) diff --git a/src/httpcore2/httpcore2/_async/connection_pool.py b/src/httpcore2/httpcore2/_async/connection_pool.py index e254fd41..2c5e3980 100644 --- a/src/httpcore2/httpcore2/_async/connection_pool.py +++ b/src/httpcore2/httpcore2/_async/connection_pool.py @@ -264,34 +264,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). + # 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 = {r.connection for r in self._requests} # 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: @@ -317,8 +323,13 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]: ] 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 diff --git a/src/httpcore2/httpcore2/_sync/connection_pool.py b/src/httpcore2/httpcore2/_sync/connection_pool.py index 20912e5c..a6051bac 100644 --- a/src/httpcore2/httpcore2/_sync/connection_pool.py +++ b/src/httpcore2/httpcore2/_sync/connection_pool.py @@ -264,34 +264,40 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]: closing_connections: list[ConnectionInterface] = [] retained_connections: list[ConnectionInterface] = [] - # Connections currently referenced by an active request (including - # connections that are in the process of being established). + # 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 = {r.connection for r in self._requests} # 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[ConnectionInterface] = [] 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: @@ -317,8 +323,13 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]: ] 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 From d6582d8dd92eb66e374bbb4561c9340ef4768921 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Thu, 23 Jul 2026 11:03:47 +0200 Subject: [PATCH 4/6] Maintain request-connection reservations incrementally --- .../httpcore2/_async/connection_pool.py | 32 ++++++++++++++++--- .../httpcore2/_sync/connection_pool.py | 32 ++++++++++++++++--- 2 files changed, 54 insertions(+), 10 deletions(-) diff --git a/src/httpcore2/httpcore2/_async/connection_pool.py b/src/httpcore2/httpcore2/_async/connection_pool.py index 2c5e3980..dbf9441f 100644 --- a/src/httpcore2/httpcore2/_async/connection_pool.py +++ b/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,7 +232,9 @@ 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 @@ -235,6 +242,7 @@ async def handle_async_request(self, request: Request) -> Response: 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 @@ -267,7 +288,7 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]: # 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 = {r.connection for r in self._requests} + 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. Reserved @@ -342,7 +363,7 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]: # to handle the request. 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. @@ -352,7 +373,7 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]: 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): @@ -362,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 @@ -434,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() diff --git a/src/httpcore2/httpcore2/_sync/connection_pool.py b/src/httpcore2/httpcore2/_sync/connection_pool.py index a6051bac..7f08ee8f 100644 --- a/src/httpcore2/httpcore2/_sync/connection_pool.py +++ b/src/httpcore2/httpcore2/_sync/connection_pool.py @@ -112,6 +112,11 @@ def __init__( self._connections: list[ConnectionInterface] = [] self._requests: list[PoolRequest] = [] + # 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[ConnectionInterface, 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,7 +232,9 @@ def handle_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 @@ -235,6 +242,7 @@ def handle_request(self, request: Request) -> Response: 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 @@ def handle_request(self, request: Request) -> Response: extensions=response.extensions, ) + def _reserve_connection(self, pool_request: PoolRequest, connection: ConnectionInterface) -> 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: PoolRequest) -> 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[ConnectionInterface]: """ Manage the state of the connection pool, assigning incoming @@ -267,7 +288,7 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]: # 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 = {r.connection for r in self._requests} + 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. Reserved @@ -342,7 +363,7 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]: # to handle the request. 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. @@ -352,7 +373,7 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]: 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): @@ -362,7 +383,7 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]: 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 @@ -434,6 +455,7 @@ def close(self) -> None: self._stream.close() 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() From 58b1b656e976871a7b7935aa272bac04dd7ebc75 Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Sun, 2 Aug 2026 20:47:56 +0100 Subject: [PATCH 5/6] Fix sync connection reservation races --- .../httpcore2/_async/connection_pool.py | 20 ++-- .../httpcore2/_sync/connection_pool.py | 20 ++-- .../httpcore2/_async/test_connection_pool.py | 101 ++++++++++++++++++ tests/httpcore2/_sync/test_connection_pool.py | 101 ++++++++++++++++++ 4 files changed, 222 insertions(+), 20 deletions(-) diff --git a/src/httpcore2/httpcore2/_async/connection_pool.py b/src/httpcore2/httpcore2/_async/connection_pool.py index dbf9441f..a6682ba6 100644 --- a/src/httpcore2/httpcore2/_async/connection_pool.py +++ b/src/httpcore2/httpcore2/_async/connection_pool.py @@ -331,16 +331,16 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]: # it per queued request — this is what brings the loop from O(N*M) to # O(N+M) in the common case. # - # 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. + # An established non-multiplexing connection already assigned to an + # in-flight request is reserved. Its state may transition from IDLE to + # ACTIVE after `is_available()` returns, so use `is_connected()` here + # rather than checking its mutable idle state. Multiplexing connections + # and not-yet-connected HTTP/2 candidates remain available. 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()) + and not (connection.is_connected() and connection in request_connections and not connection.can_multiplex()) ] new_connection_budget = self._max_connections - len(self._connections) @@ -363,11 +363,11 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]: # to handle the request. for idx, connection in enumerate(available_connections): if connection.can_handle_request(origin): - 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. + if connection.is_connected() and not connection.can_multiplex(): + # Remove an established HTTP/1.1 connection before + # waking the request, which may transition it to ACTIVE. del available_connections[idx] + self._reserve_connection(pool_request, connection) break else: if new_connection_budget > 0: diff --git a/src/httpcore2/httpcore2/_sync/connection_pool.py b/src/httpcore2/httpcore2/_sync/connection_pool.py index 7f08ee8f..62ac6dfa 100644 --- a/src/httpcore2/httpcore2/_sync/connection_pool.py +++ b/src/httpcore2/httpcore2/_sync/connection_pool.py @@ -331,16 +331,16 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]: # it per queued request — this is what brings the loop from O(N*M) to # O(N+M) in the common case. # - # 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. + # An established non-multiplexing connection already assigned to an + # in-flight request is reserved. Its state may transition from IDLE to + # ACTIVE after `is_available()` returns, so use `is_connected()` here + # rather than checking its mutable idle state. Multiplexing connections + # and not-yet-connected HTTP/2 candidates remain available. 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()) + and not (connection.is_connected() and connection in request_connections and not connection.can_multiplex()) ] new_connection_budget = self._max_connections - len(self._connections) @@ -363,11 +363,11 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]: # to handle the request. for idx, connection in enumerate(available_connections): if connection.can_handle_request(origin): - 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. + if connection.is_connected() and not connection.can_multiplex(): + # Remove an established HTTP/1.1 connection before + # waking the request, which may transition it to ACTIVE. del available_connections[idx] + self._reserve_connection(pool_request, connection) break else: if new_connection_budget > 0: diff --git a/tests/httpcore2/_async/test_connection_pool.py b/tests/httpcore2/_async/test_connection_pool.py index 11b9a5d4..0273990d 100644 --- a/tests/httpcore2/_async/test_connection_pool.py +++ b/tests/httpcore2/_async/test_connection_pool.py @@ -884,3 +884,104 @@ async def fetch(pool: httpcore2.AsyncConnectionPool) -> None: nursery.start_soon(fetch, pool) assert QueueObservingPool.max_queued_after_pass == 0 + + +class RacingConnection(httpcore2.AsyncConnectionInterface): + def __init__(self, *, activate_during_availability_check: bool = False) -> None: + self.active = False + self.closed = False + self.activate_during_availability_check = activate_during_availability_check + + async def handle_async_request(self, request: httpcore2.Request) -> httpcore2.Response: + raise NotImplementedError + + async def aclose(self) -> None: + self.closed = True + + def info(self) -> str: + return "racing connection" + + def can_handle_request(self, origin: httpcore2.Origin) -> bool: + return True + + def is_connected(self) -> bool: + return not self.closed + + def is_available(self) -> bool: + available = not self.active and not self.closed + if self.activate_during_availability_check: + self.active = True + return available + + def has_expired(self) -> bool: + return False + + def is_idle(self) -> bool: + return not self.active and not self.closed + + def is_closed(self) -> bool: + return self.closed + + +class RacingPoolRequest: + def __init__( + self, + *, + connection: RacingConnection | None = None, + activate_on_assignment: bool = False, + ) -> None: + self.request = httpcore2.Request("GET", "https://example.com/") + self.connection = connection + self.activate_on_assignment = activate_on_assignment + + def assign_to_connection(self, connection: RacingConnection) -> None: + self.connection = connection + if self.activate_on_assignment: + connection.active = True + + def is_queued(self) -> bool: + return self.connection is None + + +@pytest.mark.anyio +async def test_connection_pool_removes_connection_before_waking_request() -> None: + """ + A sync request may start using a connection as soon as its event is set. + Remove a non-multiplexing connection from the candidate snapshot before + assigning it, so that state transition cannot make it eligible again. + """ + connection = RacingConnection() + first = RacingPoolRequest(activate_on_assignment=True) + second = RacingPoolRequest() + pool = httpcore2.AsyncConnectionPool(max_connections=1) + pool_state = typing.cast(typing.Any, pool) + pool_state._connections = [connection] + pool_state._requests = [first, second] + + pool._assign_requests_to_connections() + + assert first.connection is connection + assert second.connection is None + await pool.aclose() + + +@pytest.mark.anyio +async def test_connection_pool_excludes_reserved_connection_during_state_transition() -> None: + """ + A connection may transition from IDLE to ACTIVE between `is_available()` + and the reservation check. Established reserved HTTP/1.1 connections must + remain excluded regardless of that mutable state. + """ + connection = RacingConnection(activate_during_availability_check=True) + active = RacingPoolRequest(connection=connection) + queued = RacingPoolRequest() + pool = httpcore2.AsyncConnectionPool(max_connections=1) + pool_state = typing.cast(typing.Any, pool) + pool_state._connections = [connection] + pool_state._requests = [active, queued] + pool_state._request_connections = {connection: 1} + + pool._assign_requests_to_connections() + + assert queued.connection is None + await pool.aclose() diff --git a/tests/httpcore2/_sync/test_connection_pool.py b/tests/httpcore2/_sync/test_connection_pool.py index 693b7371..d9cef938 100644 --- a/tests/httpcore2/_sync/test_connection_pool.py +++ b/tests/httpcore2/_sync/test_connection_pool.py @@ -884,3 +884,104 @@ def fetch(pool: httpcore2.ConnectionPool) -> None: nursery.start_soon(fetch, pool) assert QueueObservingPool.max_queued_after_pass == 0 + + +class RacingConnection(httpcore2.ConnectionInterface): + def __init__(self, *, activate_during_availability_check: bool = False) -> None: + self.active = False + self.closed = False + self.activate_during_availability_check = activate_during_availability_check + + def handle_request(self, request: httpcore2.Request) -> httpcore2.Response: + raise NotImplementedError + + def close(self) -> None: + self.closed = True + + def info(self) -> str: + return "racing connection" + + def can_handle_request(self, origin: httpcore2.Origin) -> bool: + return True + + def is_connected(self) -> bool: + return not self.closed + + def is_available(self) -> bool: + available = not self.active and not self.closed + if self.activate_during_availability_check: + self.active = True + return available + + def has_expired(self) -> bool: + return False + + def is_idle(self) -> bool: + return not self.active and not self.closed + + def is_closed(self) -> bool: + return self.closed + + +class RacingPoolRequest: + def __init__( + self, + *, + connection: RacingConnection | None = None, + activate_on_assignment: bool = False, + ) -> None: + self.request = httpcore2.Request("GET", "https://example.com/") + self.connection = connection + self.activate_on_assignment = activate_on_assignment + + def assign_to_connection(self, connection: RacingConnection) -> None: + self.connection = connection + if self.activate_on_assignment: + connection.active = True + + def is_queued(self) -> bool: + return self.connection is None + + + +def test_connection_pool_removes_connection_before_waking_request() -> None: + """ + A sync request may start using a connection as soon as its event is set. + Remove a non-multiplexing connection from the candidate snapshot before + assigning it, so that state transition cannot make it eligible again. + """ + connection = RacingConnection() + first = RacingPoolRequest(activate_on_assignment=True) + second = RacingPoolRequest() + pool = httpcore2.ConnectionPool(max_connections=1) + pool_state = typing.cast(typing.Any, pool) + pool_state._connections = [connection] + pool_state._requests = [first, second] + + pool._assign_requests_to_connections() + + assert first.connection is connection + assert second.connection is None + pool.close() + + + +def test_connection_pool_excludes_reserved_connection_during_state_transition() -> None: + """ + A connection may transition from IDLE to ACTIVE between `is_available()` + and the reservation check. Established reserved HTTP/1.1 connections must + remain excluded regardless of that mutable state. + """ + connection = RacingConnection(activate_during_availability_check=True) + active = RacingPoolRequest(connection=connection) + queued = RacingPoolRequest() + pool = httpcore2.ConnectionPool(max_connections=1) + pool_state = typing.cast(typing.Any, pool) + pool_state._connections = [connection] + pool_state._requests = [active, queued] + pool_state._request_connections = {connection: 1} + + pool._assign_requests_to_connections() + + assert queued.connection is None + pool.close() From 13f976f581c60ec82b4bb19bb708ef86bddca12e Mon Sep 17 00:00:00 2001 From: Marcelo Trylesinski Date: Sun, 2 Aug 2026 20:57:42 +0100 Subject: [PATCH 6/6] Preserve HTTP/2 multiplexing through proxies --- src/httpcore2/httpcore2/_async/http_proxy.py | 6 ++ src/httpcore2/httpcore2/_async/socks_proxy.py | 3 + src/httpcore2/httpcore2/_sync/http_proxy.py | 6 ++ src/httpcore2/httpcore2/_sync/socks_proxy.py | 3 + tests/httpcore2/_async/test_http_proxy.py | 1 + tests/httpcore2/_async/test_socks_proxy.py | 55 +++++++++++++++++++ tests/httpcore2/_sync/test_http_proxy.py | 1 + tests/httpcore2/_sync/test_socks_proxy.py | 55 +++++++++++++++++++ 8 files changed, 130 insertions(+) diff --git a/src/httpcore2/httpcore2/_async/http_proxy.py b/src/httpcore2/httpcore2/_async/http_proxy.py index 18d42ecf..98d4bfbe 100644 --- a/src/httpcore2/httpcore2/_async/http_proxy.py +++ b/src/httpcore2/httpcore2/_async/http_proxy.py @@ -216,6 +216,9 @@ def has_expired(self) -> bool: def is_idle(self) -> bool: return self._connection.is_idle() + def can_multiplex(self) -> bool: + return self._connection.can_multiplex() + def is_closed(self) -> bool: return self._connection.is_closed() @@ -345,6 +348,9 @@ def has_expired(self) -> bool: def is_idle(self) -> bool: return self._connection.is_idle() + def can_multiplex(self) -> bool: + return self._connection.can_multiplex() + def is_closed(self) -> bool: return self._connection.is_closed() diff --git a/src/httpcore2/httpcore2/_async/socks_proxy.py b/src/httpcore2/httpcore2/_async/socks_proxy.py index e42c8865..3f1a6c9a 100644 --- a/src/httpcore2/httpcore2/_async/socks_proxy.py +++ b/src/httpcore2/httpcore2/_async/socks_proxy.py @@ -311,6 +311,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: # pragma: no cover return self._connect_failed diff --git a/src/httpcore2/httpcore2/_sync/http_proxy.py b/src/httpcore2/httpcore2/_sync/http_proxy.py index 42e41421..33dd1f58 100644 --- a/src/httpcore2/httpcore2/_sync/http_proxy.py +++ b/src/httpcore2/httpcore2/_sync/http_proxy.py @@ -216,6 +216,9 @@ def has_expired(self) -> bool: def is_idle(self) -> bool: return self._connection.is_idle() + def can_multiplex(self) -> bool: + return self._connection.can_multiplex() + def is_closed(self) -> bool: return self._connection.is_closed() @@ -345,6 +348,9 @@ def has_expired(self) -> bool: def is_idle(self) -> bool: return self._connection.is_idle() + def can_multiplex(self) -> bool: + return self._connection.can_multiplex() + def is_closed(self) -> bool: return self._connection.is_closed() diff --git a/src/httpcore2/httpcore2/_sync/socks_proxy.py b/src/httpcore2/httpcore2/_sync/socks_proxy.py index 7547b59f..2d143aca 100644 --- a/src/httpcore2/httpcore2/_sync/socks_proxy.py +++ b/src/httpcore2/httpcore2/_sync/socks_proxy.py @@ -311,6 +311,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: # pragma: no cover return self._connect_failed diff --git a/tests/httpcore2/_async/test_http_proxy.py b/tests/httpcore2/_async/test_http_proxy.py index 046e3f01..2aaed3a1 100644 --- a/tests/httpcore2/_async/test_http_proxy.py +++ b/tests/httpcore2/_async/test_http_proxy.py @@ -174,6 +174,7 @@ async def test_proxy_tunneling_http2() -> None: assert info == [""] assert proxy.connections[0].is_idle() assert proxy.connections[0].is_available() + assert proxy.connections[0].can_multiplex() assert not proxy.connections[0].is_closed() # A connection on a tunneled proxy can only handle HTTPS requests to the same origin. diff --git a/tests/httpcore2/_async/test_socks_proxy.py b/tests/httpcore2/_async/test_socks_proxy.py index 3612868f..299d11b1 100644 --- a/tests/httpcore2/_async/test_socks_proxy.py +++ b/tests/httpcore2/_async/test_socks_proxy.py @@ -1,3 +1,8 @@ +import ssl +import typing + +import hpack +import hyperframe.frame import pytest import httpcore2 @@ -49,6 +54,56 @@ async def test_socks5_request() -> None: assert not proxy.connections[0].can_handle_request(httpcore2.Origin(b"https", b"other.com", 443)) +class HTTP2SocksStream(httpcore2.AsyncMockStream): + async def start_tls( + self, + ssl_context: ssl.SSLContext, + server_hostname: str | None = None, + timeout: float | None = None, + ) -> httpcore2.AsyncNetworkStream: + self._http2 = True + return self + + +class HTTP2SocksBackend(httpcore2.AsyncMockBackend): + async def connect_tcp( + self, + host: str, + port: int, + timeout: float | None = None, + local_address: str | None = None, + socket_options: typing.Iterable[httpcore2.SOCKET_OPTION] | None = None, + ) -> httpcore2.AsyncNetworkStream: + return HTTP2SocksStream(list(self._buffer)) + + +@pytest.mark.anyio +async def test_socks5_request_http2_can_multiplex() -> None: + network_backend = HTTP2SocksBackend( + [ + b"\x05\x00", + b"\x05\x00\x00\x01\xff\x00\x00\x01\x00\x50", + hyperframe.frame.SettingsFrame().serialize(), + hyperframe.frame.HeadersFrame( + stream_id=1, + data=hpack.Encoder().encode([(b":status", b"200")]), + flags=["END_HEADERS"], + ).serialize(), + hyperframe.frame.DataFrame(stream_id=1, data=b"Hello, world!", flags=["END_STREAM"]).serialize(), + ] + ) + + async with httpcore2.AsyncConnectionPool( + proxy=httpcore2.Proxy("socks5://localhost:8080/"), + network_backend=network_backend, + http2=True, + ) as proxy: + response = await proxy.request("GET", "https://example.com/") + + assert response.status == 200 + assert proxy.connections[0].can_multiplex() + + @pytest.mark.anyio async def test_authenticated_socks5_request() -> None: """ diff --git a/tests/httpcore2/_sync/test_http_proxy.py b/tests/httpcore2/_sync/test_http_proxy.py index 5e470f0e..db31f455 100644 --- a/tests/httpcore2/_sync/test_http_proxy.py +++ b/tests/httpcore2/_sync/test_http_proxy.py @@ -174,6 +174,7 @@ def test_proxy_tunneling_http2() -> None: assert info == [""] assert proxy.connections[0].is_idle() assert proxy.connections[0].is_available() + assert proxy.connections[0].can_multiplex() assert not proxy.connections[0].is_closed() # A connection on a tunneled proxy can only handle HTTPS requests to the same origin. diff --git a/tests/httpcore2/_sync/test_socks_proxy.py b/tests/httpcore2/_sync/test_socks_proxy.py index aa48c9a7..71139bbc 100644 --- a/tests/httpcore2/_sync/test_socks_proxy.py +++ b/tests/httpcore2/_sync/test_socks_proxy.py @@ -1,3 +1,8 @@ +import ssl +import typing + +import hpack +import hyperframe.frame import pytest import httpcore2 @@ -49,6 +54,56 @@ def test_socks5_request() -> None: assert not proxy.connections[0].can_handle_request(httpcore2.Origin(b"https", b"other.com", 443)) +class HTTP2SocksStream(httpcore2.MockStream): + def start_tls( + self, + ssl_context: ssl.SSLContext, + server_hostname: str | None = None, + timeout: float | None = None, + ) -> httpcore2.NetworkStream: + self._http2 = True + return self + + +class HTTP2SocksBackend(httpcore2.MockBackend): + def connect_tcp( + self, + host: str, + port: int, + timeout: float | None = None, + local_address: str | None = None, + socket_options: typing.Iterable[httpcore2.SOCKET_OPTION] | None = None, + ) -> httpcore2.NetworkStream: + return HTTP2SocksStream(list(self._buffer)) + + + +def test_socks5_request_http2_can_multiplex() -> None: + network_backend = HTTP2SocksBackend( + [ + b"\x05\x00", + b"\x05\x00\x00\x01\xff\x00\x00\x01\x00\x50", + hyperframe.frame.SettingsFrame().serialize(), + hyperframe.frame.HeadersFrame( + stream_id=1, + data=hpack.Encoder().encode([(b":status", b"200")]), + flags=["END_HEADERS"], + ).serialize(), + hyperframe.frame.DataFrame(stream_id=1, data=b"Hello, world!", flags=["END_STREAM"]).serialize(), + ] + ) + + with httpcore2.ConnectionPool( + proxy=httpcore2.Proxy("socks5://localhost:8080/"), + network_backend=network_backend, + http2=True, + ) as proxy: + response = proxy.request("GET", "https://example.com/") + + assert response.status == 200 + assert proxy.connections[0].can_multiplex() + + def test_authenticated_socks5_request() -> None: """