forked from MirrorHub/synapse
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:
parent
54f674f7a9
commit
90b2327066
1
changelog.d/12180.misc
Normal file
1
changelog.d/12180.misc
Normal 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`.
|
|
@ -686,12 +686,48 @@ def stop_cancellation(deferred: "defer.Deferred[T]") -> "defer.Deferred[T]":
|
||||||
Synapse logcontext rules.
|
Synapse logcontext rules.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
A new `Deferred`, which will contain the result of the original `Deferred`,
|
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 not propagate cancellation through to the original.
|
||||||
the new `Deferred` will fail with a `CancelledError` and will not follow the
|
When cancelled, the new `Deferred` will fail with a `CancelledError`.
|
||||||
Synapse logcontext rules. `make_deferred_yieldable` should be used to wrap
|
|
||||||
the new `Deferred`.
|
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)
|
deferred.chainDeferred(new_deferred)
|
||||||
return new_deferred
|
return new_deferred
|
||||||
|
|
|
@ -13,6 +13,8 @@
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
import traceback
|
import traceback
|
||||||
|
|
||||||
|
from parameterized import parameterized_class
|
||||||
|
|
||||||
from twisted.internet import defer
|
from twisted.internet import defer
|
||||||
from twisted.internet.defer import CancelledError, Deferred, ensureDeferred
|
from twisted.internet.defer import CancelledError, Deferred, ensureDeferred
|
||||||
from twisted.internet.task import Clock
|
from twisted.internet.task import Clock
|
||||||
|
@ -23,10 +25,12 @@ from synapse.logging.context import (
|
||||||
LoggingContext,
|
LoggingContext,
|
||||||
PreserveLoggingContext,
|
PreserveLoggingContext,
|
||||||
current_context,
|
current_context,
|
||||||
|
make_deferred_yieldable,
|
||||||
)
|
)
|
||||||
from synapse.util.async_helpers import (
|
from synapse.util.async_helpers import (
|
||||||
ObservableDeferred,
|
ObservableDeferred,
|
||||||
concurrently_execute,
|
concurrently_execute,
|
||||||
|
delay_cancellation,
|
||||||
stop_cancellation,
|
stop_cancellation,
|
||||||
timeout_deferred,
|
timeout_deferred,
|
||||||
)
|
)
|
||||||
|
@ -313,13 +317,27 @@ class ConcurrentlyExecuteTest(TestCase):
|
||||||
self.successResultOf(d2)
|
self.successResultOf(d2)
|
||||||
|
|
||||||
|
|
||||||
class StopCancellationTests(TestCase):
|
@parameterized_class(
|
||||||
"""Tests for the `stop_cancellation` function."""
|
("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):
|
def test_succeed(self):
|
||||||
"""Test that the new `Deferred` receives the result."""
|
"""Test that the new `Deferred` receives the result."""
|
||||||
deferred: "Deferred[str]" = Deferred()
|
deferred: "Deferred[str]" = Deferred()
|
||||||
wrapper_deferred = stop_cancellation(deferred)
|
wrapper_deferred = self.wrap_deferred(deferred)
|
||||||
|
|
||||||
# Success should propagate through.
|
# Success should propagate through.
|
||||||
deferred.callback("success")
|
deferred.callback("success")
|
||||||
|
@ -329,7 +347,7 @@ class StopCancellationTests(TestCase):
|
||||||
def test_failure(self):
|
def test_failure(self):
|
||||||
"""Test that the new `Deferred` receives the `Failure`."""
|
"""Test that the new `Deferred` receives the `Failure`."""
|
||||||
deferred: "Deferred[str]" = Deferred()
|
deferred: "Deferred[str]" = Deferred()
|
||||||
wrapper_deferred = stop_cancellation(deferred)
|
wrapper_deferred = self.wrap_deferred(deferred)
|
||||||
|
|
||||||
# Failure should propagate through.
|
# Failure should propagate through.
|
||||||
deferred.errback(ValueError("abc"))
|
deferred.errback(ValueError("abc"))
|
||||||
|
@ -337,6 +355,10 @@ class StopCancellationTests(TestCase):
|
||||||
self.failureResultOf(wrapper_deferred, ValueError)
|
self.failureResultOf(wrapper_deferred, ValueError)
|
||||||
self.assertIsNone(deferred.result, "`Failure` was not consumed")
|
self.assertIsNone(deferred.result, "`Failure` was not consumed")
|
||||||
|
|
||||||
|
|
||||||
|
class StopCancellationTests(TestCase):
|
||||||
|
"""Tests for the `stop_cancellation` function."""
|
||||||
|
|
||||||
def test_cancellation(self):
|
def test_cancellation(self):
|
||||||
"""Test that cancellation of the new `Deferred` leaves the original running."""
|
"""Test that cancellation of the new `Deferred` leaves the original running."""
|
||||||
deferred: "Deferred[str]" = Deferred()
|
deferred: "Deferred[str]" = Deferred()
|
||||||
|
@ -347,11 +369,101 @@ class StopCancellationTests(TestCase):
|
||||||
self.assertTrue(wrapper_deferred.called)
|
self.assertTrue(wrapper_deferred.called)
|
||||||
self.failureResultOf(wrapper_deferred, CancelledError)
|
self.failureResultOf(wrapper_deferred, CancelledError)
|
||||||
self.assertFalse(
|
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
|
# The `Failure` must be consumed, otherwise unwanted tracebacks will be printed
|
||||||
# in logs.
|
# in logs.
|
||||||
deferred.errback(ValueError("abc"))
|
deferred.errback(ValueError("abc"))
|
||||||
self.assertIsNone(deferred.result, "`Failure` was not consumed")
|
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)
|
||||||
|
|
Loading…
Reference in a new issue