Implement cancellation support/protection for module callbacks (#12568)

There's no guarantee that module callbacks will handle cancellation
appropriately. Protect module callbacks with read semantics from
cancellation and avoid swallowing `CancelledError`s that arise.

Other module callbacks, such as the `on_*` callbacks, are presumed to
live on code paths that involve writes and aren't cancellation-friendly.
These module callbacks have been left alone.

Signed-off-by: Sean Quah <seanq@element.io>

1 files changed, 9 insertions, 3 deletions

diff --git a/synapse/events/presence_router.py b/synapse/events/presence_router.py
index 8437ce52dc..bb4a6bd957 100644
--- a/synapse/events/presence_router.py
+++ b/synapse/events/presence_router.py
@@ -28,8 +28,10 @@ from typing import (
 
 from typing_extensions import ParamSpec
 
+from twisted.internet.defer import CancelledError
+
 from synapse.api.presence import UserPresenceState
-from synapse.util.async_helpers import maybe_awaitable
+from synapse.util.async_helpers import delay_cancellation, maybe_awaitable
 
 if TYPE_CHECKING:
     from synapse.server import HomeServer
@@ -158,7 +160,9 @@ class PresenceRouter:
             try:
                 # Note: result is an object here, because we don't trust modules to
                 # return the types they're supposed to.
-                result: object = await callback(state_updates)
+                result: object = await delay_cancellation(callback(state_updates))
+            except CancelledError:
+                raise
             except Exception as e:
                 logger.warning("Failed to run module API callback %s: %s", callback, e)
                 continue
@@ -210,7 +214,9 @@ class PresenceRouter:
         # run all the callbacks for get_interested_users and combine the results
         for callback in self._get_interested_users_callbacks:
             try:
-                result = await callback(user_id)
+                result = await delay_cancellation(callback(user_id))
+            except CancelledError:
+                raise
             except Exception as e:
                 logger.warning("Failed to run module API callback %s: %s", callback, e)
                 continue