From 5ebe74faa6c3896624cc76b8582b37fa4b9fe509 Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@element.io>
Date: Thu, 24 Feb 2022 18:41:44 +0000
Subject: [PATCH 1/2] Add `stop_cancellation` utility function

---
 changelog.d/12106.misc           |  1 +
 synapse/util/async_helpers.py    | 19 ++++++++++++
 tests/util/test_async_helpers.py | 50 +++++++++++++++++++++++++++++++-
 3 files changed, 69 insertions(+), 1 deletion(-)
 create mode 100644 changelog.d/12106.misc

diff --git a/changelog.d/12106.misc b/changelog.d/12106.misc
new file mode 100644
index 000000000000..d918e9e3b16d
--- /dev/null
+++ b/changelog.d/12106.misc
@@ -0,0 +1 @@
+Add `stop_cancellation` utility function to stop `Deferred`s from being cancelled.
diff --git a/synapse/util/async_helpers.py b/synapse/util/async_helpers.py
index 3f7299aff7eb..244c5ad3d0bc 100644
--- a/synapse/util/async_helpers.py
+++ b/synapse/util/async_helpers.py
@@ -655,3 +655,22 @@ def maybe_awaitable(value: Union[Awaitable[R], R]) -> Awaitable[R]:
         return value
 
     return DoneAwaitable(value)
+
+
+def stop_cancellation(deferred: "defer.Deferred[T]") -> "defer.Deferred[T]":
+    """Prevent a `Deferred` from being cancelled by wrapping it in another `Deferred`.
+
+    Args:
+        deferred: The `Deferred` to protect against cancellation. Must not follow the
+            Synapse logcontext rules.
+
+    Returns:
+        A new `Deferred`, which will contain the result of the original `Deferred`,
+        but will not propagate cancellation through to the original. When cancelled,
+        the new `Deferred` will fail with a `CancelledError` and will not follow the
+        Synapse logcontext rules. `make_deferred_yieldable` should be used to wrap
+        the new `Deferred`.
+    """
+    new_deferred: defer.Deferred[T] = defer.Deferred()
+    deferred.addBoth(new_deferred.callback)
+    return new_deferred
diff --git a/tests/util/test_async_helpers.py b/tests/util/test_async_helpers.py
index ab89cab81256..a375690a6e5e 100644
--- a/tests/util/test_async_helpers.py
+++ b/tests/util/test_async_helpers.py
@@ -21,7 +21,11 @@
     PreserveLoggingContext,
     current_context,
 )
-from synapse.util.async_helpers import ObservableDeferred, timeout_deferred
+from synapse.util.async_helpers import (
+    ObservableDeferred,
+    stop_cancellation,
+    timeout_deferred,
+)
 
 from tests.unittest import TestCase
 
@@ -171,3 +175,47 @@ def errback(res, deferred_name):
             )
             self.failureResultOf(timing_out_d, defer.TimeoutError)
             self.assertIs(current_context(), context_one)
+
+
+class StopCancellationTests(TestCase):
+    """Tests for the `stop_cancellation` function."""
+
+    def test_succeed(self):
+        """Test that the new `Deferred` receives the result."""
+        deferred: "Deferred[str]" = Deferred()
+        wrapper_deferred = stop_cancellation(deferred)
+
+        # Success should propagate through.
+        deferred.callback("success")
+        self.assertTrue(wrapper_deferred.called)
+        self.assertEqual("success", self.successResultOf(wrapper_deferred))
+
+    def test_failure(self):
+        """Test that the new `Deferred` receives the `Failure`."""
+        deferred: "Deferred[str]" = Deferred()
+        wrapper_deferred = stop_cancellation(deferred)
+
+        # Failure should propagate through.
+        deferred.errback(ValueError("abc"))
+        self.assertTrue(wrapper_deferred.called)
+        self.failureResultOf(wrapper_deferred, ValueError)
+        self.assertIsNone(deferred.result, "`Failure` was not consumed")
+
+    def test_cancellation(self):
+        """Test that cancellation of the new `Deferred` leaves the original running."""
+        deferred: "Deferred[str]" = Deferred()
+        wrapper_deferred = stop_cancellation(deferred)
+
+        # Cancel the new `Deferred`.
+        wrapper_deferred.cancel()
+        self.assertTrue(wrapper_deferred.called)
+        self.failureResultOf(wrapper_deferred, CancelledError)
+        self.assertFalse(
+            deferred.called, "Original `Deferred` was unexpectedly cancelled."
+        )
+
+        # Now make the inner `Deferred` fail.
+        # The `Failure` must be consumed, otherwise unwanted tracebacks will be printed
+        # in logs.
+        deferred.errback(ValueError("abc"))
+        self.assertIsNone(deferred.result, "`Failure` was not consumed")

From 44c86dcfccde8ce87ca3ecebb07b5af2dd9030ab Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@element.io>
Date: Tue, 1 Mar 2022 12:26:14 +0000
Subject: [PATCH 2/2] Use chainDeferred instead of addBoth

---
 synapse/util/async_helpers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/synapse/util/async_helpers.py b/synapse/util/async_helpers.py
index 244c5ad3d0bc..d9bb3e6fb8c7 100644
--- a/synapse/util/async_helpers.py
+++ b/synapse/util/async_helpers.py
@@ -672,5 +672,5 @@ def stop_cancellation(deferred: "defer.Deferred[T]") -> "defer.Deferred[T]":
         the new `Deferred`.
     """
     new_deferred: defer.Deferred[T] = defer.Deferred()
-    deferred.addBoth(new_deferred.callback)
+    deferred.chainDeferred(new_deferred)
     return new_deferred