Catch asyncio.CancelledError and raise GeneratorExit in aiter functions

[gjmhmm8]

Hello, I am using httpx to make asynchronous requests in a generator function. However, I encountered a problem when I tried to cancel the generator using generator.close() or asyncio.gather(*tasks, return_exceptions=True). The generator did not exit gracefully and raised an asyncio.CancelledError instead of a return. This caused some unwanted side effects and made it difficult to handle the cancellation properly.

I looked into the source code of httpx and found that the problem was in the aiter_raw and aiter_bytes functions in the httpx/_client.py module. These functions use async for to iterate over the response stream, but they do not catch the asyncio.CancelledError that may be raised when the stream is cancelled. According to the [documentation] of asyncio.CancelledError, this exception should be caught and either re-raised or suppressed. In this case, I think it would make sense to re-raise it as a GeneratorExit, which is the expected exception for generator termination.

I suggest adding a try-except block around the async for loop in the aiter_raw and aiter_bytes functions, like this:

async def aiter_raw(
    self, chunk_size: typing.Optional[int] = None
) -> typing.AsyncIterator[bytes]:
    """
    A byte-iterator over the raw response content.
    """
    if self.is_stream_consumed:
        raise StreamConsumed()
    if self.is_closed:
        raise StreamClosed()
    if not isinstance(self.stream, AsyncByteStream):
        raise RuntimeError("Attempted to call an async iterator on an sync stream.")

    try:
        self.is_stream_consumed = True
        self._num_bytes_downloaded = 0
        chunker = ByteChunker(chunk_size=chunk_size)

        with request_context(request=self._request):
            async for raw_stream_bytes in self.stream:
                self._num_bytes_downloaded += len(raw_stream_bytes)
                for chunk in chunker.decode(raw_stream_bytes):
                    await asyncio.sleep(0)
                    yield chunk
        for chunk in chunker.flush():
            await asyncio.sleep(0)
            yield chunk
    except asyncio.CancelledError:
        return
    finally:
        await self.aclose()

[tomchristie]

I'm going to start this off as a discussion rather than an issue.
Because I'm sometimes a draconian bore like that. Sorry not sorry.

Okay, could you help me out by providing a minimal reproducible example?

Bits of this that I think we'd need to understand more clearly...

• Can we deal with this just once as low down in the stack as possible, or do we need to handling it through all our async iteration cases?
• How does trio handle this case? We support either asyncio or trio, so how does resolving this for both cases look?

[DeoLeung]

I got a similar issue

SSLWantReadError: The operation did not complete (read) (_ssl.c:2580)
  File "anyio/streams/tls.py", line 131, in _call_sslobject_method
    result = func(*args)
  File "ssl.py", line 921, in read
    v = self._sslobj.read(len)
CancelledError: null
  File "xxx/api.py", line 203, in generate_messages
    async for i in _res.stream_message(result):
  File "xxx/openai.py", line 26, in stream_message
    async for i in self.aiter_lines():
  File "httpx/_models.py", line 963, in aiter_lines
    async for text in self.aiter_text():
  File "httpx/_models.py", line 950, in aiter_text
    async for byte_content in self.aiter_bytes():
  File "httpx/_models.py", line 929, in aiter_bytes
    async for raw_bytes in self.aiter_raw():
  File "httpx/_models.py", line 987, in aiter_raw
    async for raw_stream_bytes in self.stream:
  File "httpx/_client.py", line 149, in __aiter__
    async for chunk in self._stream:
  File "httpx/_transports/default.py", line 254, in __aiter__
    async for part in self._httpcore_stream:
  File "httpcore/_async/connection_pool.py", line 361, in __aiter__
    async for part in self._stream:
  File "httpcore/_async/http2.py", line 582, in __aiter__
    raise exc
  File "httpcore/_async/http2.py", line 572, in __aiter__
    async for chunk in self._connection._receive_response_body(
  File "httpcore/_async/http2.py", line 313, in _receive_response_body
    event = await self._receive_stream_event(request, stream_id)
  File "httpcore/_async/http2.py", line 333, in _receive_stream_event
    await self._receive_events(request, stream_id)
  File "httpcore/_async/http2.py", line 361, in _receive_events
    events = await self._read_incoming_data(request)
  File "httpcore/_async/http2.py", line 438, in _read_incoming_data
    data = await self._network_stream.read(self.READ_NUM_BYTES, timeout)
  File "httpcore/_backends/anyio.py", line 34, in read
    return await self._stream.receive(max_bytes=max_bytes)
  File "anyio/streams/tls.py", line 196, in receive
    data = await self._call_sslobject_method(self._ssl_object.read, max_bytes)
  File "anyio/streams/tls.py", line 138, in _call_sslobject_method
    data = await self.transport_stream.receive()
  File "anyio/streams/tls.py", line 138, in _call_sslobject_method
    data = await self.transport_stream.receive()
  File "anyio/_backends/_asyncio.py", line 1203, in receive
    await self._protocol.read_event.wait()
  File "asyncio/locks.py", line 213, in wait
    await fut

raise a proper exception except CancelledError will be better I think