Skip to content

Commit 8027999

Browse files
mbeijenVictorPrinsKludex
authored
perf: Rewrite _assign_requests_to_connections as a single-pass loop (#974)
Co-authored-by: Victor Prins <[email protected]> Co-authored-by: Marcelo Trylesinski <[email protected]>
1 parent de30f39 commit 8027999

4 files changed

Lines changed: 182 additions & 82 deletions

File tree

‎src/httpcore2/httpcore2/_async/connection_pool.py‎

Lines changed: 54 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -261,61 +261,74 @@ def _assign_requests_to_connections(self) -> list[AsyncConnectionInterface]:
261261
Any closing connections are returned, allowing the I/O for closing
262262
those connections to be handled separately.
263263
"""
264-
closing_connections = []
264+
closing_connections: list[AsyncConnectionInterface] = []
265+
retained_connections: list[AsyncConnectionInterface] = []
265266

266-
# First we handle cleaning up any connections that are closed,
267-
# have expired their keep-alive, or surplus idle connections.
268-
for connection in list(self._connections):
267+
# First we handle cleaning up any connections that are closed
268+
# or have expired their keep-alive, in a single pass.
269+
for connection in self._connections:
269270
if connection.is_closed():
270-
# log: "removing closed connection"
271-
self._connections.remove(connection)
271+
continue
272272
elif connection.has_expired():
273-
# log: "closing expired connection"
274-
self._connections.remove(connection)
275-
closing_connections.append(connection)
276-
elif (
277-
connection.is_idle()
278-
and sum(connection.is_idle() for connection in self._connections) > self._max_keepalive_connections
279-
):
280-
# log: "closing idle connection"
281-
self._connections.remove(connection)
282273
closing_connections.append(connection)
274+
else:
275+
retained_connections.append(connection)
276+
277+
# Then we close any surplus idle connections, to enforce the
278+
# max_keepalive_connections setting.
279+
idle_surplus = (
280+
sum(connection.is_idle() for connection in retained_connections) - self._max_keepalive_connections
281+
)
282+
if idle_surplus > 0:
283+
kept: list[AsyncConnectionInterface] = []
284+
for connection in retained_connections:
285+
if idle_surplus > 0 and connection.is_idle():
286+
closing_connections.append(connection)
287+
idle_surplus -= 1
288+
else:
289+
kept.append(connection)
290+
retained_connections = kept
291+
292+
self._connections = retained_connections
293+
294+
# Snapshot the set of reusable connections once, rather than rebuilding
295+
# it per queued request — this is what brings the loop from O(N*M) to
296+
# O(N+M) in the common case.
297+
available_connections = [connection for connection in self._connections if connection.is_available()]
298+
new_connection_budget = self._max_connections - len(self._connections)
283299

284300
# Assign queued requests to connections.
285-
queued_requests = [request for request in self._requests if request.is_queued()]
286-
for pool_request in queued_requests:
301+
for pool_request in self._requests:
302+
if not pool_request.is_queued():
303+
continue
287304
origin = pool_request.request.url.origin
288-
available_connections = [
289-
connection
290-
for connection in self._connections
291-
if connection.can_handle_request(origin) and connection.is_available()
292-
]
293-
idle_connections = [connection for connection in self._connections if connection.is_idle()]
294305

295306
# There are three cases for how we may be able to handle the request:
296307
#
297308
# 1. There is an existing connection that can handle the request.
298309
# 2. We can create a new connection to handle the request.
299310
# 3. We can close an idle connection and then create a new connection
300311
# to handle the request.
301-
if available_connections:
302-
# log: "reusing existing connection"
303-
connection = available_connections[0]
304-
pool_request.assign_to_connection(connection)
305-
elif len(self._connections) < self._max_connections:
306-
# log: "creating new connection"
307-
connection = self.create_connection(origin)
308-
self._connections.append(connection)
309-
pool_request.assign_to_connection(connection)
310-
elif idle_connections:
311-
# log: "closing idle connection"
312-
connection = idle_connections[0]
313-
self._connections.remove(connection)
314-
closing_connections.append(connection)
315-
# log: "creating new connection"
316-
connection = self.create_connection(origin)
317-
self._connections.append(connection)
318-
pool_request.assign_to_connection(connection)
312+
for connection in available_connections:
313+
if connection.can_handle_request(origin):
314+
pool_request.assign_to_connection(connection)
315+
break
316+
else:
317+
if new_connection_budget > 0:
318+
connection = self.create_connection(origin)
319+
self._connections.append(connection)
320+
pool_request.assign_to_connection(connection)
321+
new_connection_budget -= 1
322+
continue
323+
for idx, connection in enumerate(available_connections):
324+
if connection.is_idle():
325+
del available_connections[idx]
326+
self._connections.remove(connection)
327+
closing_connections.append(connection)
328+
connection = self.create_connection(origin)
329+
self._connections.append(connection)
330+
pool_request.assign_to_connection(connection)
331+
break
319332

320333
return closing_connections
321334

‎src/httpcore2/httpcore2/_sync/connection_pool.py‎

Lines changed: 54 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -261,61 +261,74 @@ def _assign_requests_to_connections(self) -> list[ConnectionInterface]:
261261
Any closing connections are returned, allowing the I/O for closing
262262
those connections to be handled separately.
263263
"""
264-
closing_connections = []
264+
closing_connections: list[ConnectionInterface] = []
265+
retained_connections: list[ConnectionInterface] = []
265266

266-
# First we handle cleaning up any connections that are closed,
267-
# have expired their keep-alive, or surplus idle connections.
268-
for connection in list(self._connections):
267+
# First we handle cleaning up any connections that are closed
268+
# or have expired their keep-alive, in a single pass.
269+
for connection in self._connections:
269270
if connection.is_closed():
270-
# log: "removing closed connection"
271-
self._connections.remove(connection)
271+
continue
272272
elif connection.has_expired():
273-
# log: "closing expired connection"
274-
self._connections.remove(connection)
275-
closing_connections.append(connection)
276-
elif (
277-
connection.is_idle()
278-
and sum(connection.is_idle() for connection in self._connections) > self._max_keepalive_connections
279-
):
280-
# log: "closing idle connection"
281-
self._connections.remove(connection)
282273
closing_connections.append(connection)
274+
else:
275+
retained_connections.append(connection)
276+
277+
# Then we close any surplus idle connections, to enforce the
278+
# max_keepalive_connections setting.
279+
idle_surplus = (
280+
sum(connection.is_idle() for connection in retained_connections) - self._max_keepalive_connections
281+
)
282+
if idle_surplus > 0:
283+
kept: list[ConnectionInterface] = []
284+
for connection in retained_connections:
285+
if idle_surplus > 0 and connection.is_idle():
286+
closing_connections.append(connection)
287+
idle_surplus -= 1
288+
else:
289+
kept.append(connection)
290+
retained_connections = kept
291+
292+
self._connections = retained_connections
293+
294+
# Snapshot the set of reusable connections once, rather than rebuilding
295+
# it per queued request — this is what brings the loop from O(N*M) to
296+
# O(N+M) in the common case.
297+
available_connections = [connection for connection in self._connections if connection.is_available()]
298+
new_connection_budget = self._max_connections - len(self._connections)
283299

284300
# Assign queued requests to connections.
285-
queued_requests = [request for request in self._requests if request.is_queued()]
286-
for pool_request in queued_requests:
301+
for pool_request in self._requests:
302+
if not pool_request.is_queued():
303+
continue
287304
origin = pool_request.request.url.origin
288-
available_connections = [
289-
connection
290-
for connection in self._connections
291-
if connection.can_handle_request(origin) and connection.is_available()
292-
]
293-
idle_connections = [connection for connection in self._connections if connection.is_idle()]
294305

295306
# There are three cases for how we may be able to handle the request:
296307
#
297308
# 1. There is an existing connection that can handle the request.
298309
# 2. We can create a new connection to handle the request.
299310
# 3. We can close an idle connection and then create a new connection
300311
# to handle the request.
301-
if available_connections:
302-
# log: "reusing existing connection"
303-
connection = available_connections[0]
304-
pool_request.assign_to_connection(connection)
305-
elif len(self._connections) < self._max_connections:
306-
# log: "creating new connection"
307-
connection = self.create_connection(origin)
308-
self._connections.append(connection)
309-
pool_request.assign_to_connection(connection)
310-
elif idle_connections:
311-
# log: "closing idle connection"
312-
connection = idle_connections[0]
313-
self._connections.remove(connection)
314-
closing_connections.append(connection)
315-
# log: "creating new connection"
316-
connection = self.create_connection(origin)
317-
self._connections.append(connection)
318-
pool_request.assign_to_connection(connection)
312+
for connection in available_connections:
313+
if connection.can_handle_request(origin):
314+
pool_request.assign_to_connection(connection)
315+
break
316+
else:
317+
if new_connection_budget > 0:
318+
connection = self.create_connection(origin)
319+
self._connections.append(connection)
320+
pool_request.assign_to_connection(connection)
321+
new_connection_budget -= 1
322+
continue
323+
for idx, connection in enumerate(available_connections):
324+
if connection.is_idle():
325+
del available_connections[idx]
326+
self._connections.remove(connection)
327+
closing_connections.append(connection)
328+
connection = self.create_connection(origin)
329+
self._connections.append(connection)
330+
pool_request.assign_to_connection(connection)
331+
break
319332

320333
return closing_connections
321334

‎tests/httpcore2/_async/test_connection_pool.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -461,6 +461,43 @@ async def test_connection_pool_with_no_keepalive_connections_allowed() -> None:
461461
assert info == []
462462

463463

464+
@pytest.mark.anyio
465+
async def test_connection_pool_closes_idle_connection_for_different_origin() -> None:
466+
"""
467+
When the pool is at 'max_connections' and an incoming request is for an
468+
origin with no reusable connection, an IDLE connection to a different
469+
origin is closed to make room for the new connection.
470+
"""
471+
network_backend = httpcore2.AsyncMockBackend(
472+
[
473+
b"HTTP/1.1 200 OK\r\n",
474+
b"Content-Type: plain/text\r\n",
475+
b"Content-Length: 13\r\n",
476+
b"\r\n",
477+
b"Hello, world!",
478+
b"HTTP/1.1 200 OK\r\n",
479+
b"Content-Type: plain/text\r\n",
480+
b"Content-Length: 13\r\n",
481+
b"\r\n",
482+
b"Hello, world!",
483+
]
484+
)
485+
486+
async with httpcore2.AsyncConnectionPool(network_backend=network_backend, max_connections=1) as pool:
487+
# An initial request to a.com leaves a single IDLE connection in the pool.
488+
response = await pool.request("GET", "https://a.com/")
489+
assert response.status == 200
490+
info = [repr(c) for c in pool.connections]
491+
assert info == ["<AsyncHTTPConnection ['https://a.com:443', HTTP/1.1, IDLE, Request Count: 1]>"]
492+
493+
# A request to b.com cannot reuse the a.com connection and the pool is full,
494+
# so the IDLE a.com connection is closed and replaced with a b.com connection.
495+
response = await pool.request("GET", "https://b.com/")
496+
assert response.status == 200
497+
info = [repr(c) for c in pool.connections]
498+
assert info == ["<AsyncHTTPConnection ['https://b.com:443', HTTP/1.1, IDLE, Request Count: 1]>"]
499+
500+
464501
@pytest.mark.trio
465502
async def test_connection_pool_concurrency() -> None:
466503
"""

‎tests/httpcore2/_sync/test_connection_pool.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -462,6 +462,43 @@ def test_connection_pool_with_no_keepalive_connections_allowed() -> None:
462462

463463

464464

465+
def test_connection_pool_closes_idle_connection_for_different_origin() -> None:
466+
"""
467+
When the pool is at 'max_connections' and an incoming request is for an
468+
origin with no reusable connection, an IDLE connection to a different
469+
origin is closed to make room for the new connection.
470+
"""
471+
network_backend = httpcore2.MockBackend(
472+
[
473+
b"HTTP/1.1 200 OK\r\n",
474+
b"Content-Type: plain/text\r\n",
475+
b"Content-Length: 13\r\n",
476+
b"\r\n",
477+
b"Hello, world!",
478+
b"HTTP/1.1 200 OK\r\n",
479+
b"Content-Type: plain/text\r\n",
480+
b"Content-Length: 13\r\n",
481+
b"\r\n",
482+
b"Hello, world!",
483+
]
484+
)
485+
486+
with httpcore2.ConnectionPool(network_backend=network_backend, max_connections=1) as pool:
487+
# An initial request to a.com leaves a single IDLE connection in the pool.
488+
response = pool.request("GET", "https://a.com/")
489+
assert response.status == 200
490+
info = [repr(c) for c in pool.connections]
491+
assert info == ["<HTTPConnection ['https://a.com:443', HTTP/1.1, IDLE, Request Count: 1]>"]
492+
493+
# A request to b.com cannot reuse the a.com connection and the pool is full,
494+
# so the IDLE a.com connection is closed and replaced with a b.com connection.
495+
response = pool.request("GET", "https://b.com/")
496+
assert response.status == 200
497+
info = [repr(c) for c in pool.connections]
498+
assert info == ["<HTTPConnection ['https://b.com:443', HTTP/1.1, IDLE, Request Count: 1]>"]
499+
500+
501+
465502
def test_connection_pool_concurrency() -> None:
466503
"""
467504
HTTP/1.1 requests made in concurrency must not ever exceed the maximum number

0 commit comments

Comments
 (0)