Inner task finished by best effort

encode · Jun 11, 2024 · 333c20c · 333c20c
1 parent 2402e8e
commit 333c20c
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 18 deletions.
diff --git 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
@@ -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
@@ -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...