Fix cancellations being swallowed

CHANGES/9030.bugfix.rst

@@ -0,0 +1 @@
+Fixed (on Python 3.11+) some edge cases where a task cancellation may get incorrectly suppressed -- by :user:`Dreamsorcerer`.

aiohttp/client_reqrep.py

@@ -561,11 +561,8 @@
         """Support coroutines that yields bytes objects."""
         # 100 response
         if self._continue is not None:
-            try:
-                await writer.drain()
-                await self._continue
-            except asyncio.CancelledError:
-                return
+            await writer.drain()
+            await self._continue
 
         protocol = conn.protocol
         assert protocol is not None
@@ -594,6 +591,7 @@
         except asyncio.CancelledError:
             # Body hasn't been fully sent, so connection can't be reused.
             conn.close()
+            raise
         except Exception as underlying_exc:
             set_exception(
                 protocol,
@@ -696,8 +694,15 @@
 
     async def close(self) -> None:
         if self._writer is not None:
-            with contextlib.suppress(asyncio.CancelledError):
+            try:
                 await self._writer
+            except asyncio.CancelledError:
+                if (
+                    sys.version_info >= (3, 11)
+                    and (task := asyncio.current_task())
+                    and task.cancelling()
+                ):
+                    raise
 
     def terminate(self) -> None:
         if self._writer is not None:
@@ -1040,7 +1045,15 @@
 
     async def _wait_released(self) -> None:
         if self._writer is not None:
-            await self._writer
+            try:
+                await self._writer
+            except asyncio.CancelledError:
+                if (
+                    sys.version_info >= (3, 11)
+                    and (task := asyncio.current_task())
+                    and task.cancelling()
+                ):
+                    raise
         self._release_connection()
 
     def _cleanup_writer(self) -> None:
@@ -1057,7 +1070,15 @@
 
     async def wait_for_close(self) -> None:
         if self._writer is not None:
-            await self._writer
+            try:
+                await self._writer
+            except asyncio.CancelledError:
+                if (
+                    sys.version_info >= (3, 11)
+                    and (task := asyncio.current_task())
+                    and task.cancelling()
+                ):
+                    raise
         self.release()
 
     async def read(self) -> bytes:

aiohttp/web_protocol.py

@@ -285,17 +285,31 @@ async def shutdown(self, timeout: Optional[float] = 15.0) -> None:
 
         # Wait for graceful handler completion
         if self._handler_waiter is not None:
-            with suppress(asyncio.CancelledError, asyncio.TimeoutError):
+            try:
                 async with ceil_timeout(timeout):
                     await self._handler_waiter
+            except (asyncio.CancelledError, asyncio.TimeoutError):
+                if (
+                    sys.version_info >= (3, 11)
+                    and (task := asyncio.current_task())
+                    and task.cancelling()
+                ):
+                    raise
         # Then cancel handler and wait
-        with suppress(asyncio.CancelledError, asyncio.TimeoutError):
+        try:
             async with ceil_timeout(timeout):
                 if self._current_request is not None:
                     self._current_request._cancel(asyncio.CancelledError())
 
                 if self._task_handler is not None and not self._task_handler.done():
                     await self._task_handler
+        except (asyncio.CancelledError, asyncio.TimeoutError):
+            if (
+                sys.version_info >= (3, 11)
+                and (task := asyncio.current_task())
+                and task.cancelling()
+            ):
+                raise
 
         # force-close non-idle handler
         if self._task_handler is not None:
@@ -523,8 +537,6 @@ async def start(self) -> None:
                     # wait for next request
                     self._waiter = loop.create_future()
                     await self._waiter
-                except asyncio.CancelledError:
-                    break
                 finally:
                     self._waiter = None
 
@@ -551,7 +563,7 @@ async def start(self) -> None:
                     task = loop.create_task(coro)
                 try:
                     resp, reset = await task
-                except (asyncio.CancelledError, ConnectionError):
+                except ConnectionError:
                     self.log_debug("Ignored premature client disconnection")
                     break
 
@@ -577,23 +589,29 @@ async def start(self) -> None:
                         now = loop.time()
                         end_t = now + lingering_time
 
-                        with suppress(asyncio.TimeoutError, asyncio.CancelledError):
+                        try:
                             while not payload.is_eof() and now < end_t:
                                 async with ceil_timeout(end_t - now):
                                     # read and ignore
                                     await payload.readany()
                                 now = loop.time()
+                        except (asyncio.CancelledError, asyncio.TimeoutError):
+                            if (
+                                sys.version_info >= (3, 11)
+                                and (t := asyncio.current_task())
+                                and t.cancelling()
+                            ):
+                                raise
 
                     # if payload still uncompleted
                     if not payload.is_eof() and not self._force_close:
                         self.log_debug("Uncompleted request.")
                         self.close()
 
                 set_exception(payload, PayloadAccessError())
-
             except asyncio.CancelledError:
-                self.log_debug("Ignored premature client disconnection ")
-                break
+                self.log_debug("Ignored premature client disconnection")
+                raise
             except Exception as exc:
                 self.log_exception("Unhandled exception", exc_info=exc)
                 self.force_close()

tests/test_client_request.py

@@ -2,6 +2,7 @@
 import hashlib
 import io
 import pathlib
+import sys
 import zlib
 from http.cookies import BaseCookie, Morsel, SimpleCookie
 from typing import Any, AsyncIterator, Callable, Dict, Iterator, List, Protocol
@@ -1234,6 +1235,22 @@
     assert isinstance(exc, aiohttp.ClientOSError)
 
 
+@pytest.mark.skipif(sys.version_info < (3, 11), reason="Needs Task.cancelling()")
+async def test_cancel_close(loop: asyncio.AbstractEventLoop, conn: mock.Mock) -> None:
+    req = ClientRequest("get", URL("http://python.org"), loop=loop)
+    req._writer = asyncio.Future()  # type: ignore[assignment]
+
+    t = asyncio.create_task(req.close())
+
+    # Start waiting on _writer
+    await asyncio.sleep(0)
+
+    t.cancel()
+    # Cancellation should not be suppressed.
+    with pytest.raises(asyncio.CancelledError):
+        await t
+
+
 async def test_terminate(loop: asyncio.AbstractEventLoop, conn: mock.Mock) -> None:
     req = ClientRequest("get", URL("http://python.org"), loop=loop)
 

tests/test_web_functional.py

@@ -4,6 +4,7 @@
 import json
 import pathlib
 import socket
+import sys
 import zlib
 from typing import Any, NoReturn, Optional
 from unittest import mock
@@ -188,6 +189,45 @@
     resp.release()
 
 
+@pytest.mark.skipif(sys.version_info < (3, 11), reason="Needs Task.cancelling()")
+async def test_cancel_shutdown(aiohttp_client: Any) -> None:
+    async def handler(request):
+        t = asyncio.create_task(request.protocol.shutdown())
+        # Ensure it's started waiting
+        await asyncio.sleep(0)
+
+        t.cancel()
+        # Cancellation should not be suppressed
+        with pytest.raises(asyncio.CancelledError):
+            await t
+
+        # Repeat for second waiter in shutdown()
+        with mock.patch.object(request.protocol, "_handler_waiter", None):
+            with mock.patch.object(
+                request.protocol,
+                "_current_request",
+                autospec=True,
+                spec_set=True,
+            ):
+                t = asyncio.create_task(request.protocol.shutdown())
+                await asyncio.sleep(0)
+
+                t.cancel()
+                with pytest.raises(asyncio.CancelledError):
+                    await t
+
+        return web.Response(body=b"OK")
+
+    app = web.Application()
+    app.router.add_get("/", handler)
+    client = await aiohttp_client(app)
+
+    async with client.get("/") as resp:
+        assert resp.status == 200
+        txt = await resp.text()
+        assert txt == "OK"
+
+
 async def test_post_form(aiohttp_client: Any) -> None:
     async def handler(request):
         data = await request.post()