Add a Request.wait_for_disconnection() method (#4200) (#8504)

Co-authored-by: Andrew Svetlov <[email protected]>
Co-authored-by: Gustavo J. A. M. Carneiro <[email protected]>
Co-authored-by: J. Nick Koston <[email protected]>

CHANGES/2492.feature

@@ -0,0 +1 @@
+Add a Request.wait_for_disconnection() method, as means of allowing request handlers to be notified of premature client disconnections.

aiohttp/web_protocol.py

@@ -609,6 +609,7 @@ async def finish_response(
         can get exception information. Returns True if the client disconnects
         prematurely.
         """
+        request._finish()
         if self._request_parser is not None:
             self._request_parser.set_upgraded(False)
             self._upgrade = False

aiohttp/web_request.py

@@ -19,6 +19,7 @@
     MutableMapping,
     Optional,
     Pattern,
+    Set,
     Tuple,
     Union,
     cast,
@@ -49,6 +50,7 @@
     reify,
     sentinel,
     set_exception,
+    set_result,
 )
 from .http_parser import RawRequestMessage
 from .http_writer import HttpVersion
@@ -144,6 +146,7 @@ class BaseRequest(MutableMapping[str, Any], HeadersMixin):
             "_loop",
             "_transport_sslcontext",
             "_transport_peername",
+            "_disconnection_waiters",
         ]
     )
 
@@ -191,6 +194,7 @@ def __init__(
         self._task = task
         self._client_max_size = client_max_size
         self._loop = loop
+        self._disconnection_waiters: Set[asyncio.Future[None]] = set()
 
         transport = self._protocol.transport
         assert transport is not None
@@ -818,6 +822,21 @@ async def _prepare_hook(self, response: StreamResponse) -> None:
 
     def _cancel(self, exc: BaseException) -> None:
         set_exception(self._payload, exc)
+        for fut in self._disconnection_waiters:
+            set_result(fut, None)
+
+    def _finish(self) -> None:
+        for fut in self._disconnection_waiters:
+            fut.cancel()
+
+    async def wait_for_disconnection(self) -> None:
+        loop = asyncio.get_event_loop()
+        fut = loop.create_future()  # type: asyncio.Future[None]
+        self._disconnection_waiters.add(fut)
+        try:
+            await fut
+        finally:
+            self._disconnection_waiters.remove(fut)
 
 
 class Request(BaseRequest):

docs/web_reference.rst

@@ -510,6 +510,18 @@ and :ref:`aiohttp-web-signals` handlers.
           required work will be processed by :mod:`aiohttp.web`
           internal machinery.
 
+   .. method:: wait_for_disconnection()
+
+      Returns when the connection that sent this request closes
+
+      If there is no client disconnection during request handling, this
+      coroutine gets cancelled automatically at the end of this request being
+      handled.
+
+      This can be used in handlers as a means of receiving a notification of
+      premature client disconnection.
+
+      .. versionadded:: 3.10
 
 .. class:: Request