AsyncConnectionPool gets stuck upon task group cancellation

[juneoh]

PoolTimeout is raised with 0 active requests, but 5 active connections. I.e.,

<AsyncConnectionPool [Requests: 0 active, 0 queued | Connections: 5 active, 0 idle]>

Case 1: Active connection's AsyncHTTPConnection.info() returns CONNECTING

1. Initial pool state: <AsyncConnectionPool [Requests: 5 active, 251 queued | Connections: 5 active, 0 idle]> where all the requests belong to the same anyio task group.
2. The task group is cancelled.
3. For an active request holding an active connection, trio.Cancelled is caught from AsyncConnectionPool.handle_async_request()
4. AsyncConnectionPool._assign_requests_to_connections() creates a new connection in CONNECTING state, and assigns it to a pending request.
5. However, when the pending request's time comes, they never get to initialize the connection because connection = await pool_request.wait_for_connection(timeout=timeout) also raises trio.Cancelled.
6. The request with the new CONNECTING connection is removed, but the connection remains.
7. The connection is no longer able to be assigned to requests, because it is not AsyncHTTPConnection.is_available() (assuming HTTP/1)

Case 2: Active connection's AsyncHTTP11Connection.info() returns [..] NEW [..]

1. Similar to case 1, but this time trio.Cancelled arises from async with self._state_lock within AsyncHTTP11Connection.handle_async_request().
2. The exception bubbles up to AsyncConnectionPool.handle_async_request().
3. AsyncConnectionPool._assign_requests_to_connections() does not remove the connection, because it is NEW.
4. The connection is no longer able to be assigned to requests, because it is not IDLE as mandated by AsyncHTTP11Connection.is_available().

Traceback (the line numbers may differ due to my additional logging):

Traceback (most recent call last):
  File "python3.11/site-packages/httpcore/_async/connection_pool.py", line 201, in handle_async_request
    response = await connection.handle_async_request(
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python3.11/site-packages/httpcore/_async/connection.py", line 105, in handle_async_request
    return await self._connection.handle_async_request(request)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python3.11/site-packages/httpcore/_async/http11.py", line 79, in handle_async_request
    async with self._state_lock:
  File "python3.11/site-packages/httpcore/_synchronization.py", line 74, in __aenter__
    await self._trio_lock.acquire()
  File "python3.11/site-packages/trio/_sync.py", line 585, in acquire
    await trio.lowlevel.checkpoint_if_cancelled()
  File "python3.11/site-packages/trio/_core/_run.py", line 2817, in checkpoint_if_cancelled
    await _core.checkpoint()
  File "python3.11/site-packages/trio/_core/_run.py", line 2795, in checkpoint
    await _core.wait_task_rescheduled(lambda _: _core.Abort.SUCCEEDED)
  File "python3.11/site-packages/trio/_core/_traps.py", line 179, in wait_task_rescheduled
    return (await _async_yield(WaitTaskRescheduled(abort_func))).unwrap()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "python3.11/site-packages/outcome/_impl.py", line 213, in unwrap
    raise captured_error
  File "python3.11/site-packages/trio/_core/_run.py", line 1454, in raise_cancel
    raise Cancelled._create()
trio.Cancelled: Cancelled

[juneoh]

Partial reproduction for case 1:

import functools
import unittest.mock

import httpcore
import trio


async def main():
    pool = httpcore.AsyncConnectionPool(max_connections=1)
    request = functools.partial(
        pool.request,
        "GET",
        "https://www.example.com/",
        extensions={"timeout": {"pool": 10}},
    )

    async def side_effect(*args, **kwargs):
        print(pool)
        await trio.sleep(1)
        raise trio.Cancelled._create()

    with unittest.mock.patch(
        "httpcore._async.connection_pool.AsyncPoolRequest.wait_for_connection",
        side_effect,
    ):
        async with trio.open_nursery() as nursery:
            nursery.start_soon(request)
            nursery.start_soon(request)

    print(pool)
    print(pool._connections[0].info())

    await request()


trio.run(main)

[juneoh]

Seems like wait_for_connection() is more likely to raise Cancelled if tasks are created recursively within a task group, e.g. for DFS.

import unittest.mock

import anyio
import httpcore


async def main():
    pool = httpcore.AsyncConnectionPool(max_connections=1)
    task_group = anyio.create_task_group()

    async def request(depth=0):
        print(pool)
        if depth < 2:
            task_group.start_soon(request, depth + 1)
            if depth == 1:
                raise RuntimeError
        await pool.request(
            "GET",
            "https://www.example.com/",
            extensions={"timeout": {"pool": 10}},
        )

    async def side_effect(*args, **kwargs):
        await anyio.sleep(2)

    with unittest.mock.patch(
        "httpcore._async.connection_pool.AsyncHTTPConnection._connect",
        side_effect,
    ):
        try:
            async with task_group:
                task_group.start_soon(request)
                await anyio.sleep(0.01)
        except Exception:
            pass

    print(pool)
    print(pool._connections[0].info())

    await pool.request(
        "GET",
        "https://www.example.com/",
        extensions={"timeout": {"pool": 10}},
    )


anyio.run(main)

[juneoh]

Reproduction for case 2:

import functools
import unittest.mock

import httpcore
import trio


class MockContextManager:
    async def __aenter__(self):
        await trio.sleep(1)
        raise trio.Cancelled._create()

    async def __aexit__(self):
        pass


class MockAsyncHTTP11Connection(httpcore.AsyncHTTP11Connection):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._state_lock = MockContextManager()


async def main():
    pool = httpcore.AsyncConnectionPool(max_connections=1)
    request = functools.partial(
        pool.request,
        "GET",
        "https://www.example.com/",
        extensions={"timeout": {"pool": 10}},
    )

    with unittest.mock.patch(
        "httpcore._async.connection.AsyncHTTP11Connection",
        MockAsyncHTTP11Connection,
    ):
        async with trio.open_nursery() as nursery:
            nursery.start_soon(request)
            nursery.start_soon(request)

    print(pool)
    print(pool._connections[0].info())

    await request()


trio.run(main)

Please forgive the mocks, since it's very difficult to time task cancellations.