diff --git a/httpcore/_async/http2.py b/httpcore/_async/http2.py index 6f42c7d4..da47fbb2 100644 --- a/httpcore/_async/http2.py +++ b/httpcore/_async/http2.py @@ -160,13 +160,11 @@ async def handle_async_request(self, request: Request) -> Response: }, ) except BaseException as exc: - - async def close() -> None: - kwargs = {"stream_id": stream_id} - async with Trace("response_closed", logger, request, kwargs): - await self._response_closed(stream_id=stream_id) - - await async_cancel_shield(close) + kwargs = {"stream_id": stream_id} + async with Trace("response_closed", logger, request, kwargs): + await async_cancel_shield( + lambda: self._response_closed(stream_id=stream_id) + ) if isinstance(exc, h2.exceptions.ProtocolError): # One case where h2 can raise a protocol error is when a diff --git a/httpcore/_sync/http2.py b/httpcore/_sync/http2.py index 05575338..ea0b02b7 100644 --- a/httpcore/_sync/http2.py +++ b/httpcore/_sync/http2.py @@ -160,13 +160,11 @@ def handle_request(self, request: Request) -> Response: }, ) except BaseException as exc: - - def close() -> None: - kwargs = {"stream_id": stream_id} - with Trace("response_closed", logger, request, kwargs): - self._response_closed(stream_id=stream_id) - - sync_cancel_shield(close) + kwargs = {"stream_id": stream_id} + with Trace("response_closed", logger, request, kwargs): + sync_cancel_shield( + lambda: self._response_closed(stream_id=stream_id) + ) if isinstance(exc, h2.exceptions.ProtocolError): # One case where h2 can raise a protocol error is when a diff --git a/httpcore/_synchronization.py b/httpcore/_synchronization.py index 1a20aefd..ccf22faf 100644 --- a/httpcore/_synchronization.py +++ b/httpcore/_synchronization.py @@ -205,10 +205,18 @@ async def async_cancel_shield( await shielded() else: inner_task = asyncio.create_task(shielded()) - try: - await asyncio.shield(inner_task) - except asyncio.CancelledError: - return + retry = False + while True: + try: + await asyncio.shield(inner_task) + break + except asyncio.CancelledError: + if inner_task.done() or retry: + break + # We may get multiple cancellations. + # Retry once to get inner_task finished here by best effort. + retry = True + continue # Our thread-based synchronization primitives...