Add delay_cancellation utility function (#12180)

`delay_cancellation` behaves like `stop_cancellation`, except it
delays `CancelledError`s until the original `Deferred` resolves.
This is handy for unifying cleanup paths and ensuring that uncancelled
coroutines don't use finished logcontexts.

Signed-off-by: Sean Quah <seanq@element.io>
This commit is contained in:
Sean Quah 2022-03-14 17:52:15 +00:00 committed by GitHub
parent 54f674f7a9
commit 90b2327066
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 161 additions and 12 deletions

1
changelog.d/12180.misc Normal file
View file

@ -0,0 +1 @@
Add `delay_cancellation` utility function, which behaves like `stop_cancellation` but waits until the original `Deferred` resolves before raising a `CancelledError`.

View file

@ -686,12 +686,48 @@ def stop_cancellation(deferred: "defer.Deferred[T]") -> "defer.Deferred[T]":
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`.
A new `Deferred`, which will contain the result of the original `Deferred`.
The new `Deferred` will not propagate cancellation through to the original.
When cancelled, the new `Deferred` will fail with a `CancelledError`.
The new `Deferred` will not follow the Synapse logcontext rules and should be
wrapped with `make_deferred_yieldable`.
"""
new_deferred: defer.Deferred[T] = defer.Deferred()
new_deferred: "defer.Deferred[T]" = defer.Deferred()
deferred.chainDeferred(new_deferred)
return new_deferred
def delay_cancellation(deferred: "defer.Deferred[T]") -> "defer.Deferred[T]":
"""Delay cancellation of a `Deferred` until it resolves.
Has the same effect as `stop_cancellation`, but the returned `Deferred` will not
resolve with a `CancelledError` until the original `Deferred` resolves.
Args:
deferred: The `Deferred` to protect against cancellation. May optionally follow
the Synapse logcontext rules.
Returns:
A new `Deferred`, which will contain the result of the original `Deferred`.
The new `Deferred` will not propagate cancellation through to the original.
When cancelled, the new `Deferred` will wait until the original `Deferred`
resolves before failing with a `CancelledError`.
The new `Deferred` will follow the Synapse logcontext rules if `deferred`
follows the Synapse logcontext rules. Otherwise the new `Deferred` should be
wrapped with `make_deferred_yieldable`.
"""
def handle_cancel(new_deferred: "defer.Deferred[T]") -> None:
# before the new deferred is cancelled, we `pause` it to stop the cancellation
# propagating. we then `unpause` it once the wrapped deferred completes, to
# propagate the exception.
new_deferred.pause()
new_deferred.errback(Failure(CancelledError()))
deferred.addBoth(lambda _: new_deferred.unpause())
new_deferred: "defer.Deferred[T]" = defer.Deferred(handle_cancel)
deferred.chainDeferred(new_deferred)
return new_deferred

View file

@ -13,6 +13,8 @@
# limitations under the License.
import traceback
from parameterized import parameterized_class
from twisted.internet import defer
from twisted.internet.defer import CancelledError, Deferred, ensureDeferred
from twisted.internet.task import Clock
@ -23,10 +25,12 @@ from synapse.logging.context import (
LoggingContext,
PreserveLoggingContext,
current_context,
make_deferred_yieldable,
)
from synapse.util.async_helpers import (
ObservableDeferred,
concurrently_execute,
delay_cancellation,
stop_cancellation,
timeout_deferred,
)
@ -313,13 +317,27 @@ class ConcurrentlyExecuteTest(TestCase):
self.successResultOf(d2)
class StopCancellationTests(TestCase):
"""Tests for the `stop_cancellation` function."""
@parameterized_class(
("wrapper",),
[("stop_cancellation",), ("delay_cancellation",)],
)
class CancellationWrapperTests(TestCase):
"""Common tests for the `stop_cancellation` and `delay_cancellation` functions."""
wrapper: str
def wrap_deferred(self, deferred: "Deferred[str]") -> "Deferred[str]":
if self.wrapper == "stop_cancellation":
return stop_cancellation(deferred)
elif self.wrapper == "delay_cancellation":
return delay_cancellation(deferred)
else:
raise ValueError(f"Unsupported wrapper type: {self.wrapper}")
def test_succeed(self):
"""Test that the new `Deferred` receives the result."""
deferred: "Deferred[str]" = Deferred()
wrapper_deferred = stop_cancellation(deferred)
wrapper_deferred = self.wrap_deferred(deferred)
# Success should propagate through.
deferred.callback("success")
@ -329,7 +347,7 @@ class StopCancellationTests(TestCase):
def test_failure(self):
"""Test that the new `Deferred` receives the `Failure`."""
deferred: "Deferred[str]" = Deferred()
wrapper_deferred = stop_cancellation(deferred)
wrapper_deferred = self.wrap_deferred(deferred)
# Failure should propagate through.
deferred.errback(ValueError("abc"))
@ -337,6 +355,10 @@ class StopCancellationTests(TestCase):
self.failureResultOf(wrapper_deferred, ValueError)
self.assertIsNone(deferred.result, "`Failure` was not consumed")
class StopCancellationTests(TestCase):
"""Tests for the `stop_cancellation` function."""
def test_cancellation(self):
"""Test that cancellation of the new `Deferred` leaves the original running."""
deferred: "Deferred[str]" = Deferred()
@ -347,11 +369,101 @@ class StopCancellationTests(TestCase):
self.assertTrue(wrapper_deferred.called)
self.failureResultOf(wrapper_deferred, CancelledError)
self.assertFalse(
deferred.called, "Original `Deferred` was unexpectedly cancelled."
deferred.called, "Original `Deferred` was unexpectedly cancelled"
)
# Now make the inner `Deferred` fail.
# Now make the original `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")
class DelayCancellationTests(TestCase):
"""Tests for the `delay_cancellation` function."""
def test_cancellation(self):
"""Test that cancellation of the new `Deferred` waits for the original."""
deferred: "Deferred[str]" = Deferred()
wrapper_deferred = delay_cancellation(deferred)
# Cancel the new `Deferred`.
wrapper_deferred.cancel()
self.assertNoResult(wrapper_deferred)
self.assertFalse(
deferred.called, "Original `Deferred` was unexpectedly cancelled"
)
# Now make the original `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")
# Now that the original `Deferred` has failed, we should get a `CancelledError`.
self.failureResultOf(wrapper_deferred, CancelledError)
def test_suppresses_second_cancellation(self):
"""Test that a second cancellation is suppressed.
Identical to `test_cancellation` except the new `Deferred` is cancelled twice.
"""
deferred: "Deferred[str]" = Deferred()
wrapper_deferred = delay_cancellation(deferred)
# Cancel the new `Deferred`, twice.
wrapper_deferred.cancel()
wrapper_deferred.cancel()
self.assertNoResult(wrapper_deferred)
self.assertFalse(
deferred.called, "Original `Deferred` was unexpectedly cancelled"
)
# Now make the original `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")
# Now that the original `Deferred` has failed, we should get a `CancelledError`.
self.failureResultOf(wrapper_deferred, CancelledError)
def test_propagates_cancelled_error(self):
"""Test that a `CancelledError` from the original `Deferred` gets propagated."""
deferred: "Deferred[str]" = Deferred()
wrapper_deferred = delay_cancellation(deferred)
# Fail the original `Deferred` with a `CancelledError`.
cancelled_error = CancelledError()
deferred.errback(cancelled_error)
# The new `Deferred` should fail with exactly the same `CancelledError`.
self.assertTrue(wrapper_deferred.called)
self.assertIs(cancelled_error, self.failureResultOf(wrapper_deferred).value)
def test_preserves_logcontext(self):
"""Test that logging contexts are preserved."""
blocking_d: "Deferred[None]" = Deferred()
async def inner():
await make_deferred_yieldable(blocking_d)
async def outer():
with LoggingContext("c") as c:
try:
await delay_cancellation(defer.ensureDeferred(inner()))
self.fail("`CancelledError` was not raised")
except CancelledError:
self.assertEqual(c, current_context())
# Succeed with no error, unless the logging context is wrong.
# Run and block inside `inner()`.
d = defer.ensureDeferred(outer())
self.assertEqual(SENTINEL_CONTEXT, current_context())
d.cancel()
# Now unblock. `outer()` will consume the `CancelledError` and check the
# logging context.
blocking_d.callback(None)
self.successResultOf(d)