From ba0fc0abe1fc99c22d25be8c47fc71e7f375246f Mon Sep 17 00:00:00 2001
From: "Gustavo J. A. M. Carneiro" <gjcarneiro@gmail.com>
Date: Fri, 25 Oct 2019 18:07:35 +0100
Subject: [PATCH] Add a Request.wait_for_disconnection() method (#4200)

as means of allowing request handlers to be notified of premature client
disconnections.

Co-Authored-By: Andrew Svetlov <andrew.svetlov@gmail.com>
(cherry picked from commit 7f777333a4ec0043ddf2e8d67146a626089773d9)
---
 CHANGES/2492.feature    |  1 +
 aiohttp/web_protocol.py |  1 +
 aiohttp/web_request.py  | 19 +++++++++++++++++++
 docs/web_reference.rst  | 12 ++++++++++++
 4 files changed, 33 insertions(+)
 create mode 100644 CHANGES/2492.feature

diff --git a/CHANGES/2492.feature b/CHANGES/2492.feature
new file mode 100644
index 00000000000..5c98dbbbcf2
--- /dev/null
+++ b/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.
diff --git a/aiohttp/web_protocol.py b/aiohttp/web_protocol.py
index f083b13eb0f..88df4b31d24 100644
--- a/aiohttp/web_protocol.py
+++ b/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
diff --git a/aiohttp/web_request.py b/aiohttp/web_request.py
index 7d1694584ea..944661e8f16 100644
--- a/aiohttp/web_request.py
+++ b/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()  # type: Set[asyncio.Future[None]]
 
         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):
diff --git a/docs/web_reference.rst b/docs/web_reference.rst
index f96fd59f56a..3ccd82d87d2 100644
--- a/docs/web_reference.rst
+++ b/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.
 
+   .. comethod:: 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:: 4.0
 
 .. class:: Request