From 5cfb0045955f8b5e9e8a1e0505fa2b5ed4f7bde2 Mon Sep 17 00:00:00 2001
From: Sean Quah <8349537+squahtx@users.noreply.github.com>
Date: Tue, 10 May 2022 14:06:08 +0100
Subject: [PATCH] Add ability to cancel disconnected requests to
 `SynapseRequest` (#12588)

Signed-off-by: Sean Quah <seanq@element.io>
---
 changelog.d/12588.misc |  1 +
 synapse/http/site.py   | 24 +++++++++++++++++++++++-
 2 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 changelog.d/12588.misc

diff --git a/changelog.d/12588.misc b/changelog.d/12588.misc
new file mode 100644
index 000000000000..f62d5c8e210c
--- /dev/null
+++ b/changelog.d/12588.misc
@@ -0,0 +1 @@
+Add ability to cancel disconnected requests to `SynapseRequest`.
diff --git a/synapse/http/site.py b/synapse/http/site.py
index 0b85a57d7787..f7f1c5704285 100644
--- a/synapse/http/site.py
+++ b/synapse/http/site.py
@@ -19,6 +19,7 @@
 import attr
 from zope.interface import implementer
 
+from twisted.internet.defer import Deferred
 from twisted.internet.interfaces import IAddress, IReactorTime
 from twisted.python.failure import Failure
 from twisted.web.http import HTTPChannel
@@ -91,6 +92,13 @@ def __init__(
         # we can't yet create the logcontext, as we don't know the method.
         self.logcontext: Optional[LoggingContext] = None
 
+        # The `Deferred` to cancel if the client disconnects early. Expected to be set
+        # by `Resource.render`.
+        self.render_deferred: Optional["Deferred[None]"] = None
+        # A boolean indicating whether `_render_deferred` should be cancelled if the
+        # client disconnects early. Expected to be set during `Resource.render`.
+        self.is_render_cancellable = False
+
         global _next_request_seq
         self.request_seq = _next_request_seq
         _next_request_seq += 1
@@ -357,7 +365,21 @@ def connectionLost(self, reason: Union[Failure, Exception]) -> None:
                     {"event": "client connection lost", "reason": str(reason.value)}
                 )
 
-            if not self._is_processing:
+            if self._is_processing:
+                if self.is_render_cancellable:
+                    if self.render_deferred is not None:
+                        # Throw a cancellation into the request processing, in the hope
+                        # that it will finish up sooner than it normally would.
+                        # The `self.processing()` context manager will call
+                        # `_finished_processing()` when done.
+                        with PreserveLoggingContext():
+                            self.render_deferred.cancel()
+                    else:
+                        logger.error(
+                            "Connection from client lost, but have no Deferred to "
+                            "cancel even though the request is marked as cancellable."
+                        )
+            else:
                 self._finished_processing()
 
     def _started_processing(self, servlet_name: str) -> None: