Repository: matrix-org/synapse

Issue:
Run mypy over docker directory
There's a couple of scripts in there which it would be nice to typecheck.

- change mypy.ini's `files` list to include `docker`
- run `mypy` and see what errors crop up
- fix them until none remain

There's the `contrib` directory too, but we don't promise to support that diligently.

Discussion:

Patch:
diff --git a/changelog.d/12468.misc b/changelog.d/12468.misc
new file mode 100644
index 000000000..3d5d25247
--- /dev/null
+++ b/changelog.d/12468.misc
@@ -0,0 +1,1 @@
+Update `delay_cancellation` to accept any awaitable, rather than just `Deferred`s.
diff --git a/changelog.d/12492.doc b/changelog.d/12492.doc
new file mode 100644
index 000000000..4a3e2f4f0
--- /dev/null
+++ b/changelog.d/12492.doc
@@ -0,0 +1,1 @@
+Add some example configurations for workers and update architectural diagram.
diff --git a/changelog.d/12501.doc b/changelog.d/12501.doc
new file mode 100644
index 000000000..278193a69
--- /dev/null
+++ b/changelog.d/12501.doc
@@ -0,0 +1,1 @@
+Add HAProxy delegation example with CORS headers to docs.
diff --git a/changelog.d/12519.misc b/changelog.d/12519.misc
new file mode 100644
index 000000000..9c023d8e3
--- /dev/null
+++ b/changelog.d/12519.misc
@@ -0,0 +1,1 @@
+Refactor the relations code for clarity.
diff --git a/changelog.d/12528.misc b/changelog.d/12528.misc
new file mode 100644
index 000000000..f64b5d24b
--- /dev/null
+++ b/changelog.d/12528.misc
@@ -0,0 +1,1 @@
+Add type hints so `docker` and `stubs` directories pass `mypy --disallow-untyped-defs`.
diff --git a/docker/configure_workers_and_start.py b/docker/configure_workers_and_start.py
index 23cac18e8..3bda6c300 100755
--- a/docker/configure_workers_and_start.py
+++ b/docker/configure_workers_and_start.py
@@ -29,7 +29,7 @@
 import os
 import subprocess
 import sys
-from typing import Any, Dict, Mapping, Set
+from typing import Any, Dict, List, Mapping, MutableMapping, NoReturn, Set
 
 import jinja2
 import yaml
@@ -201,7 +201,7 @@ upstream {upstream_worker_type} {{
 
 
 # Utility functions
-def log(txt: str):
+def log(txt: str) -> None:
     """Log something to the stdout.
 
     Args:
@@ -210,7 +210,7 @@ def log(txt: str):
     print(txt)
 
 
-def error(txt: str):
+def error(txt: str) -> NoReturn:
     """Log something and exit with an error code.
 
     Args:
@@ -220,7 +220,7 @@ def error(txt: str):
     sys.exit(2)
 
 
-def convert(src: str, dst: str, **template_vars):
+def convert(src: str, dst: str, **template_vars: object) -> None:
     """Generate a file from a template
 
     Args:
@@ -290,7 +290,7 @@ def add_sharding_to_shared_config(
         shared_config.setdefault("media_instance_running_background_jobs", worker_name)
 
 
-def generate_base_homeserver_config():
+def generate_base_homeserver_config() -> None:
     """Starts Synapse and generates a basic homeserver config, which will later be
     modified for worker support.
 
@@ -302,12 +302,14 @@ def generate_base_homeserver_config():
     subprocess.check_output(["/usr/local/bin/python", "/start.py", "migrate_config"])
 
 
-def generate_worker_files(environ, config_path: str, data_dir: str):
+def generate_worker_files(
+    environ: Mapping[str, str], config_path: str, data_dir: str
+) -> None:
     """Read the desired list of workers from environment variables and generate
     shared homeserver, nginx and supervisord configs.
 
     Args:
-        environ: _Environ[str]
+        environ: os.environ instance.
         config_path: The location of the generated Synapse main worker config file.
         data_dir: The location of the synapse data directory. Where log and
             user-facing config files live.
@@ -369,13 +371,13 @@ def generate_worker_files(environ, config_path: str, data_dir: str):
     nginx_locations = {}
 
     # Read the desired worker configuration from the environment
-    worker_types = environ.get("SYNAPSE_WORKER_TYPES")
-    if worker_types is None:
+    worker_types_env = environ.get("SYNAPSE_WORKER_TYPES")
+    if worker_types_env is None:
         # No workers, just the main process
         worker_types = []
     else:
         # Split type names by comma
-        worker_types = worker_types.split(",")
+        worker_types = worker_types_env.split(",")
 
     # Create the worker configuration directory if it doesn't already exist
     os.makedirs("/conf/workers", exist_ok=True)
@@ -547,7 +549,7 @@ def generate_worker_log_config(
     return log_config_filepath
 
 
-def main(args, environ):
+def main(args: List[str], environ: MutableMapping[str, str]) -> None:
     config_dir = environ.get("SYNAPSE_CONFIG_DIR", "/data")
     config_path = environ.get("SYNAPSE_CONFIG_PATH", config_dir + "/homeserver.yaml")
     data_dir = environ.get("SYNAPSE_DATA_DIR", "/data")
diff --git a/docker/start.py b/docker/start.py
index ac62bbc8b..4ac8f0347 100755
--- a/docker/start.py
+++ b/docker/start.py
@@ -6,27 +6,28 @@ import os
 import platform
 import subprocess
 import sys
+from typing import Any, Dict, List, Mapping, MutableMapping, NoReturn, Optional
 
 import jinja2
 
 
 # Utility functions
-def log(txt):
+def log(txt: str) -> None:
     print(txt, file=sys.stderr)
 
 
-def error(txt):
+def error(txt: str) -> NoReturn:
     log(txt)
     sys.exit(2)
 
 
-def convert(src, dst, environ):
+def convert(src: str, dst: str, environ: Mapping[str, object]) -> None:
     """Generate a file from a template
 
     Args:
-        src (str): path to input file
-        dst (str): path to file to write
-        environ (dict): environment dictionary, for replacement mappings.
+        src: path to input file
+        dst: path to file to write
+        environ: environment dictionary, for replacement mappings.
     """
     with open(src) as infile:
         template = infile.read()
@@ -35,25 +36,30 @@ def convert(src, dst, environ):
         outfile.write(rendered)
 
 
-def generate_config_from_template(config_dir, config_path, environ, ownership):
+def generate_config_from_template(
+    config_dir: str,
+    config_path: str,
+    os_environ: Mapping[str, str],
+    ownership: Optional[str],
+) -> None:
     """Generate a homeserver.yaml from environment variables
 
     Args:
-        config_dir (str): where to put generated config files
-        config_path (str): where to put the main config file
-        environ (dict): environment dictionary
-        ownership (str|None): "<user>:<group>" string which will be used to set
+        config_dir: where to put generated config files
+        config_path: where to put the main config file
+        os_environ: environment mapping
+        ownership: "<user>:<group>" string which will be used to set
             ownership of the generated configs. If None, ownership will not change.
     """
     for v in ("SYNAPSE_SERVER_NAME", "SYNAPSE_REPORT_STATS"):
-        if v not in environ:
+        if v not in os_environ:
             error(
                 "Environment variable '%s' is mandatory when generating a config file."
                 % (v,)
             )
 
     # populate some params from data files (if they exist, else create new ones)
-    environ = environ.copy()
+    environ: Dict[str, Any] = dict(os_environ)
     secrets = {
         "registration": "SYNAPSE_REGISTRATION_SHARED_SECRET",
         "macaroon": "SYNAPSE_MACAROON_SECRET_KEY",
@@ -127,12 +133,12 @@ def generate_config_from_template(config_dir, config_path, environ, ownership):
     subprocess.check_output(args)
 
 
-def run_generate_config(environ, ownership):
+def run_generate_config(environ: Mapping[str, str], ownership: Optional[str]) -> None:
     """Run synapse with a --generate-config param to generate a template config file
 
     Args:
-        environ (dict): env var dict
-        ownership (str|None): "userid:groupid" arg for chmod. If None, ownership will not change.
+        environ: env vars from `os.enrivon`.
+        ownership: "userid:groupid" arg for chmod. If None, ownership will not change.
 
     Never returns.
     """
@@ -178,7 +184,7 @@ def run_generate_config(environ, ownership):
     os.execv(sys.executable, args)
 
 
-def main(args, environ):
+def main(args: List[str], environ: MutableMapping[str, str]) -> None:
     mode = args[1] if len(args) > 1 else "run"
 
     # if we were given an explicit user to switch to, do so
diff --git a/docs/reverse_proxy.md b/docs/reverse_proxy.md
index 5a0c84795..69caa8a73 100644
--- a/docs/reverse_proxy.md
+++ b/docs/reverse_proxy.md
@@ -206,6 +206,28 @@ backend matrix
   server matrix 127.0.0.1:8008
 ```
 
+
+[Delegation](delegate.md) example:
+```
+frontend https
+  acl matrix-well-known-client-path path /.well-known/matrix/client
+  acl matrix-well-known-server-path path /.well-known/matrix/server
+  use_backend matrix-well-known-client if matrix-well-known-client-path
+  use_backend matrix-well-known-server if matrix-well-known-server-path
+ 
+backend matrix-well-known-client
+  http-after-response set-header Access-Control-Allow-Origin "*"
+  http-after-response set-header Access-Control-Allow-Methods "GET, POST, PUT, DELETE, OPTIONS"
+  http-after-response set-header Access-Control-Allow-Headers "Origin, X-Requested-With, Content-Type, Accept, Authorization"
+  http-request return status 200 content-type application/json string '{"m.homeserver":{"base_url":"https://matrix.example.com"},"m.identity_server":{"base_url":"https://identity.example.com"}}'
+
+backend matrix-well-known-server
+  http-after-response set-header Access-Control-Allow-Origin "*"
+  http-after-response set-header Access-Control-Allow-Methods "GET, POST, PUT, DELETE, OPTIONS"
+  http-after-response set-header Access-Control-Allow-Headers "Origin, X-Requested-With, Content-Type, Accept, Authorization"
+  http-request return status 200 content-type application/json string '{"m.server":"matrix.example.com:443"}'
+```
+
 ### Relayd
 
 ```
diff --git a/docs/systemd-with-workers/workers/background_worker.yaml b/docs/systemd-with-workers/workers/background_worker.yaml
new file mode 100644
index 000000000..9fbfbda7d
--- /dev/null
+++ b/docs/systemd-with-workers/workers/background_worker.yaml
@@ -0,0 +1,8 @@
+worker_app: synapse.app.generic_worker
+worker_name: background_worker
+
+# The replication listener on the main synapse process.
+worker_replication_host: 127.0.0.1
+worker_replication_http_port: 9093
+
+worker_log_config: /etc/matrix-synapse/background-worker-log.yaml
diff --git a/docs/systemd-with-workers/workers/event_persister.yaml b/docs/systemd-with-workers/workers/event_persister.yaml
new file mode 100644
index 000000000..9bc6997ba
--- /dev/null
+++ b/docs/systemd-with-workers/workers/event_persister.yaml
@@ -0,0 +1,23 @@
+worker_app: synapse.app.generic_worker
+worker_name: event_persister1 
+
+# The replication listener on the main synapse process.
+worker_replication_host: 127.0.0.1
+worker_replication_http_port: 9093
+
+worker_listeners:
+  - type: http
+    port: 8034
+    resources:
+      - names: [replication]
+
+  # Enable listener if this stream writer handles endpoints for the `typing` or
+  # `to_device` streams. Uses a different port to the `replication` listener to
+  # avoid exposing the `replication` listener publicly.
+  #
+  #- type: http
+  #  port: 8035
+  #  resources:
+  #    - names: [client]
+
+worker_log_config: /etc/matrix-synapse/event-persister-log.yaml
diff --git a/docs/systemd-with-workers/workers/generic_worker.yaml b/docs/systemd-with-workers/workers/generic_worker.yaml
index 8561e2cda..a82f9c161 100644
--- a/docs/systemd-with-workers/workers/generic_worker.yaml
+++ b/docs/systemd-with-workers/workers/generic_worker.yaml
@@ -1,12 +1,13 @@
 worker_app: synapse.app.generic_worker
 worker_name: generic_worker1
 
+# The replication listener on the main synapse process.
 worker_replication_host: 127.0.0.1
 worker_replication_http_port: 9093
 
 worker_listeners:
   - type: http
-    port: 8011
+    port: 8083
     resources:
       - names: [client, federation]
 
diff --git a/docs/workers.md b/docs/workers.md
index 858411b15..afdcd785e 100644
--- a/docs/workers.md
+++ b/docs/workers.md
@@ -138,20 +138,7 @@ as the `listeners` option in the shared config.
 For example:
 
 ```yaml
-worker_app: synapse.app.generic_worker
-worker_name: worker1
-
-# The replication listener on the main synapse process.
-worker_replication_host: 127.0.0.1
-worker_replication_http_port: 9093
-
-worker_listeners:
-  - type: http
-    port: 8083
-    resources:
-      - names: [client, federation]
-
-worker_log_config: /home/matrix/synapse/config/worker1_log_config.yaml
+{{#include systemd-with-workers/workers/generic_worker.yaml}}
 ```
 
 ...is a full configuration for a generic worker instance, which will expose a
@@ -363,6 +350,12 @@ stream_writers:
     events: event_persister1
 ```
 
+An example for a stream writer instance:
+
+```yaml
+{{#include systemd-with-workers/workers/event_persister.yaml}}
+```
+
 Some of the streams have associated endpoints which, for maximum efficiency, should
 be routed to the workers handling that stream. See below for the currently supported
 streams and the endpoints associated with them:
@@ -436,6 +429,12 @@ run_background_tasks_on: background_worker
 You might also wish to investigate the `update_user_directory` and
 `media_instance_running_background_jobs` settings.
 
+An example for a dedicated background worker instance:
+
+```yaml
+{{#include systemd-with-workers/workers/background_worker.yaml}}
+```
+
 ### `synapse.app.pusher`
 
 Handles sending push notifications to sygnal and email. Doesn't handle any
@@ -615,14 +614,14 @@ The following shows an example setup using Redis and a reverse proxy:
 |   Main       |  |   Generic    |  |   Generic    |  |  Event       |
 |   Process    |  |   Worker 1   |  |   Worker 2   |  |  Persister   |
 +--------------+  +--------------+  +--------------+  +--------------+
-      ^    ^          |   ^   |         |   ^   |          ^    ^
-      |    |          |   |   |         |   |   |          |    |
-      |    |          |   |   |  HTTP   |   |   |          |    |
-      |    +----------+<--|---|---------+   |   |          |    |
-      |                   |   +-------------|-->+----------+    |
-      |                   |                 |                   |
-      |                   |                 |                   |
-      v                   v                 v                   v
-====================================================================
+      ^    ^          |   ^   |         |   ^   |         |   ^   ^
+      |    |          |   |   |         |   |   |         |   |   |
+      |    |          |   |   |  HTTP   |   |   |         |   |   |
+      |    +----------+<--|---|---------+<--|---|---------+   |   |
+      |                   |   +-------------|-->+-------------+   |
+      |                   |                 |                     |
+      |                   |                 |                     |
+      v                   v                 v                     v
+======================================================================
                                                          Redis pub/sub channel
 ```
diff --git a/stubs/sortedcontainers/sorteddict.pyi b/stubs/sortedcontainers/sorteddict.pyi
index 344d55cce..e18d61728 100644
--- a/stubs/sortedcontainers/sorteddict.pyi
+++ b/stubs/sortedcontainers/sorteddict.pyi
@@ -103,7 +103,7 @@ class SortedDict(Dict[_KT, _VT]):
         self,
         start: Optional[int] = ...,
         stop: Optional[int] = ...,
-        reverse=bool,
+        reverse: bool = ...,
     ) -> Iterator[_KT]: ...
     def bisect_left(self, value: _KT) -> int: ...
     def bisect_right(self, value: _KT) -> int: ...
diff --git a/stubs/sortedcontainers/sortedlist.pyi b/stubs/sortedcontainers/sortedlist.pyi
index f80a3a72c..403897e39 100644
--- a/stubs/sortedcontainers/sortedlist.pyi
+++ b/stubs/sortedcontainers/sortedlist.pyi
@@ -81,7 +81,7 @@ class SortedList(MutableSequence[_T]):
         self,
         start: Optional[int] = ...,
         stop: Optional[int] = ...,
-        reverse=bool,
+        reverse: bool = ...,
     ) -> Iterator[_T]: ...
     def _islice(
         self,
@@ -153,14 +153,14 @@ class SortedKeyList(SortedList[_T]):
         maximum: Optional[int] = ...,
         inclusive: Tuple[bool, bool] = ...,
         reverse: bool = ...,
-    ): ...
+    ) -> Iterator[_T]: ...
     def irange_key(
         self,
         min_key: Optional[Any] = ...,
         max_key: Optional[Any] = ...,
         inclusive: Tuple[bool, bool] = ...,
         reserve: bool = ...,
-    ): ...
+    ) -> Iterator[_T]: ...
     def bisect_left(self, value: _T) -> int: ...
     def bisect_right(self, value: _T) -> int: ...
     def bisect(self, value: _T) -> int: ...
diff --git a/stubs/sortedcontainers/sortedset.pyi b/stubs/sortedcontainers/sortedset.pyi
index f9c290838..43c860f42 100644
--- a/stubs/sortedcontainers/sortedset.pyi
+++ b/stubs/sortedcontainers/sortedset.pyi
@@ -103,7 +103,7 @@ class SortedSet(MutableSet[_T], Sequence[_T]):
         self,
         start: Optional[int] = ...,
         stop: Optional[int] = ...,
-        reverse=bool,
+        reverse: bool = ...,
     ) -> Iterator[_T]: ...
     def irange(
         self,
diff --git a/stubs/txredisapi.pyi b/stubs/txredisapi.pyi
index 2d8ca018f..695a2307c 100644
--- a/stubs/txredisapi.pyi
+++ b/stubs/txredisapi.pyi
@@ -18,6 +18,8 @@ from typing import Any, List, Optional, Type, Union
 
 from twisted.internet import protocol
 from twisted.internet.defer import Deferred
+from twisted.internet.interfaces import IAddress
+from twisted.python.failure import Failure
 
 class RedisProtocol(protocol.Protocol):
     def publish(self, channel: str, message: bytes) -> "Deferred[None]": ...
@@ -34,11 +36,14 @@ class RedisProtocol(protocol.Protocol):
     def get(self, key: str) -> "Deferred[Any]": ...
 
 class SubscriberProtocol(RedisProtocol):
-    def __init__(self, *args, **kwargs): ...
+    def __init__(self, *args: object, **kwargs: object): ...
     password: Optional[str]
-    def subscribe(self, channels: Union[str, List[str]]): ...
-    def connectionMade(self): ...
-    def connectionLost(self, reason): ...
+    def subscribe(self, channels: Union[str, List[str]]) -> "Deferred[None]": ...
+    def connectionMade(self) -> None: ...
+    # type-ignore: twisted.internet.protocol.Protocol provides a default argument for
+    # `reason`. txredisapi's LineReceiver Protocol doesn't. But that's fine: it's what's
+    # actually specified in twisted.internet.interfaces.IProtocol.
+    def connectionLost(self, reason: Failure) -> None: ...  # type: ignore[override]
 
 def lazyConnection(
     host: str = ...,
@@ -74,7 +79,7 @@ class RedisFactory(protocol.ReconnectingClientFactory):
         replyTimeout: Optional[int] = None,
         convertNumbers: Optional[int] = True,
     ): ...
-    def buildProtocol(self, addr) -> RedisProtocol: ...
+    def buildProtocol(self, addr: IAddress) -> RedisProtocol: ...
 
 class SubscriberFactory(RedisFactory):
     def __init__(self) -> None: ...
diff --git a/synapse/events/utils.py b/synapse/events/utils.py
index 2174b4a09..f8d3ba545 100644
--- a/synapse/events/utils.py
+++ b/synapse/events/utils.py
@@ -479,9 +479,9 @@ class EventClientSerializer:
         Args:
             event: The event being serialized.
             time_now: The current time in milliseconds
+            config: Event serialization config
             aggregations: The bundled aggregation to serialize.
             serialized_event: The serialized event which may be modified.
-            config: Event serialization config
             apply_edits: Whether the content of the event should be modified to reflect
                any replacement in `aggregations.replace`.
         """
diff --git a/synapse/handlers/relations.py b/synapse/handlers/relations.py
index 0be231957..5efb56127 100644
--- a/synapse/handlers/relations.py
+++ b/synapse/handlers/relations.py
@@ -256,64 +256,6 @@ class RelationsHandler:
 
         return filtered_results
 
-    async def _get_bundled_aggregation_for_event(
-        self, event: EventBase, ignored_users: FrozenSet[str]
-    ) -> Optional[BundledAggregations]:
-        """Generate bundled aggregations for an event.
-
-        Note that this does not use a cache, but depends on cached methods.
-
-        Args:
-            event: The event to calculate bundled aggregations for.
-            ignored_users: The users ignored by the requesting user.
-
-        Returns:
-            The bundled aggregations for an event, if bundled aggregations are
-            enabled and the event can have bundled aggregations.
-        """
-
-        # Do not bundle aggregations for an event which represents an edit or an
-        # annotation. It does not make sense for them to have related events.
-        relates_to = event.content.get("m.relates_to")
-        if isinstance(relates_to, (dict, frozendict)):
-            relation_type = relates_to.get("rel_type")
-            if relation_type in (RelationTypes.ANNOTATION, RelationTypes.REPLACE):
-                return None
-
-        event_id = event.event_id
-        room_id = event.room_id
-
-        # The bundled aggregations to include, a mapping of relation type to a
-        # type-specific value. Some types include the direct return type here
-        # while others need more processing during serialization.
-        aggregations = BundledAggregations()
-
-        annotations = await self.get_annotations_for_event(
-            event_id, room_id, ignored_users=ignored_users
-        )
-        if annotations:
-            aggregations.annotations = {"chunk": annotations}
-
-        references, next_token = await self.get_relations_for_event(
-            event_id,
-            event,
-            room_id,
-            RelationTypes.REFERENCE,
-            ignored_users=ignored_users,
-        )
-        if references:
-            aggregations.references = {
-                "chunk": [{"event_id": event.event_id} for event in references]
-            }
-
-            if next_token:
-                aggregations.references["next_batch"] = await next_token.to_string(
-                    self._main_store
-                )
-
-        # Store the bundled aggregations in the event metadata for later use.
-        return aggregations
-
     async def get_threads_for_events(
         self, event_ids: Collection[str], user_id: str, ignored_users: FrozenSet[str]
     ) -> Dict[str, _ThreadAggregation]:
@@ -435,11 +377,39 @@ class RelationsHandler:
 
         # Fetch other relations per event.
         for event in events_by_id.values():
-            event_result = await self._get_bundled_aggregation_for_event(
-                event, ignored_users
+            # Do not bundle aggregations for an event which represents an edit or an
+            # annotation. It does not make sense for them to have related events.
+            relates_to = event.content.get("m.relates_to")
+            if isinstance(relates_to, (dict, frozendict)):
+                relation_type = relates_to.get("rel_type")
+                if relation_type in (RelationTypes.ANNOTATION, RelationTypes.REPLACE):
+                    continue
+
+            annotations = await self.get_annotations_for_event(
+                event.event_id, event.room_id, ignored_users=ignored_users
             )
-            if event_result:
-                results[event.event_id] = event_result
+            if annotations:
+                results.setdefault(
+                    event.event_id, BundledAggregations()
+                ).annotations = {"chunk": annotations}
+
+            references, next_token = await self.get_relations_for_event(
+                event.event_id,
+                event,
+                event.room_id,
+                RelationTypes.REFERENCE,
+                ignored_users=ignored_users,
+            )
+            if references:
+                aggregations = results.setdefault(event.event_id, BundledAggregations())
+                aggregations.references = {
+                    "chunk": [{"event_id": ev.event_id} for ev in references]
+                }
+
+                if next_token:
+                    aggregations.references["next_batch"] = await next_token.to_string(
+                        self._main_store
+                    )
 
         # Fetch any edits (but not for redacted events).
         #
diff --git a/synapse/storage/database.py b/synapse/storage/database.py
index 5eb545c86..df1e9c1b8 100644
--- a/synapse/storage/database.py
+++ b/synapse/storage/database.py
@@ -41,7 +41,6 @@ from prometheus_client import Histogram
 from typing_extensions import Literal
 
 from twisted.enterprise import adbapi
-from twisted.internet import defer
 
 from synapse.api.errors import StoreError
 from synapse.config.database import DatabaseConnectionConfig
@@ -794,7 +793,7 @@ class DatabasePool:
         # We also wait until everything above is done before releasing the
         # `CancelledError`, so that logging contexts won't get used after they have been
         # finished.
-        return await delay_cancellation(defer.ensureDeferred(_runInteraction()))
+        return await delay_cancellation(_runInteraction())
 
     async def runWithConnection(
         self,
diff --git a/synapse/util/async_helpers.py b/synapse/util/async_helpers.py
index 650e44de2..e27c5d298 100644
--- a/synapse/util/async_helpers.py
+++ b/synapse/util/async_helpers.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import abc
+import asyncio
 import collections
 import inspect
 import itertools
@@ -25,6 +26,7 @@ from typing import (
     Awaitable,
     Callable,
     Collection,
+    Coroutine,
     Dict,
     Generic,
     Hashable,
@@ -701,27 +703,57 @@ def stop_cancellation(deferred: "defer.Deferred[T]") -> "defer.Deferred[T]":
     return new_deferred
 
 
-def delay_cancellation(deferred: "defer.Deferred[T]") -> "defer.Deferred[T]":
-    """Delay cancellation of a `Deferred` until it resolves.
+@overload
+def delay_cancellation(awaitable: "defer.Deferred[T]") -> "defer.Deferred[T]":
+    ...
+
+
+@overload
+def delay_cancellation(awaitable: Coroutine[Any, Any, T]) -> "defer.Deferred[T]":
+    ...
+
+
+@overload
+def delay_cancellation(awaitable: Awaitable[T]) -> Awaitable[T]:
+    ...
+
+
+def delay_cancellation(awaitable: Awaitable[T]) -> Awaitable[T]:
+    """Delay cancellation of a coroutine or `Deferred` awaitable 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.
+    resolve with a `CancelledError` until the original awaitable resolves.
 
     Args:
-        deferred: The `Deferred` to protect against cancellation. May optionally follow
-            the Synapse logcontext rules.
+        deferred: The coroutine or `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`.
+        A new `Deferred`, which will contain the result of the original coroutine or
+        `Deferred`. The new `Deferred` will not propagate cancellation through to the
+        original coroutine or `Deferred`.
 
-        The new `Deferred` will follow the Synapse logcontext rules if `deferred`
+        When cancelled, the new `Deferred` will wait until the original coroutine or
+        `Deferred` resolves before failing with a `CancelledError`.
+
+        The new `Deferred` will follow the Synapse logcontext rules if `awaitable`
         follows the Synapse logcontext rules. Otherwise the new `Deferred` should be
         wrapped with `make_deferred_yieldable`.
     """
 
+    # First, convert the awaitable into a `Deferred`.
+    if isinstance(awaitable, defer.Deferred):
+        deferred = awaitable
+    elif asyncio.iscoroutine(awaitable):
+        # Ideally we'd use `Deferred.fromCoroutine()` here, to save on redundant
+        # type-checking, but we'd need Twisted >= 21.2.
+        deferred = defer.ensureDeferred(awaitable)
+    else:
+        # We have no idea what to do with this awaitable.
+        # We assume it's already resolved, such as `DoneAwaitable`s or `Future`s from
+        # `make_awaitable`, and let the caller `await` it normally.
+        return awaitable
+
     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


Test patch:
diff --git a/tests/rest/client/test_relations.py b/tests/rest/client/test_relations.py
index 65743cdf6..39667e322 100644
--- a/tests/rest/client/test_relations.py
+++ b/tests/rest/client/test_relations.py
@@ -560,43 +560,6 @@ class RelationsTestCase(BaseRelationsTestCase):
                 {"event_id": edit_event_id, "sender": self.user_id}, m_replace_dict
             )
 
-    def test_edit_thread(self) -> None:
-        """Test that editing a thread works."""
-
-        # Create a thread and edit the last event.
-        channel = self._send_relation(
-            RelationTypes.THREAD,
-            "m.room.message",
-            content={"msgtype": "m.text", "body": "A threaded reply!"},
-        )
-        threaded_event_id = channel.json_body["event_id"]
-
-        new_body = {"msgtype": "m.text", "body": "I've been edited!"}
-        self._send_relation(
-            RelationTypes.REPLACE,
-            "m.room.message",
-            content={"msgtype": "m.text", "body": "foo", "m.new_content": new_body},
-            parent_id=threaded_event_id,
-        )
-
-        # Fetch the thread root, to get the bundled aggregation for the thread.
-        channel = self.make_request(
-            "GET",
-            f"/rooms/{self.room}/event/{self.parent_id}",
-            access_token=self.user_token,
-        )
-        self.assertEqual(200, channel.code, channel.json_body)
-
-        # We expect that the edit message appears in the thread summary in the
-        # unsigned relations section.
-        relations_dict = channel.json_body["unsigned"].get("m.relations")
-        self.assertIn(RelationTypes.THREAD, relations_dict)
-
-        thread_summary = relations_dict[RelationTypes.THREAD]
-        self.assertIn("latest_event", thread_summary)
-        latest_event_in_thread = thread_summary["latest_event"]
-        self.assertEqual(latest_event_in_thread["content"]["body"], "I've been edited!")
-
     def test_edit_edit(self) -> None:
         """Test that an edit cannot be edited."""
         new_body = {"msgtype": "m.text", "body": "Initial edit"}
@@ -1047,7 +1010,7 @@ class BundledAggregationsTestCase(BaseRelationsTestCase):
         channel = self._send_relation(RelationTypes.THREAD, "m.room.test")
         thread_2 = channel.json_body["event_id"]
 
-        def assert_annotations(bundled_aggregations: JsonDict) -> None:
+        def assert_thread(bundled_aggregations: JsonDict) -> None:
             self.assertEqual(2, bundled_aggregations.get("count"))
             self.assertTrue(bundled_aggregations.get("current_user_participated"))
             # The latest thread event has some fields that don't matter.
@@ -1066,7 +1029,38 @@ class BundledAggregationsTestCase(BaseRelationsTestCase):
                 bundled_aggregations.get("latest_event"),
             )
 
-        self._test_bundled_aggregations(RelationTypes.THREAD, assert_annotations, 9)
+        self._test_bundled_aggregations(RelationTypes.THREAD, assert_thread, 9)
+
+    def test_thread_edit_latest_event(self) -> None:
+        """Test that editing the latest event in a thread works."""
+
+        # Create a thread and edit the last event.
+        channel = self._send_relation(
+            RelationTypes.THREAD,
+            "m.room.message",
+            content={"msgtype": "m.text", "body": "A threaded reply!"},
+        )
+        threaded_event_id = channel.json_body["event_id"]
+
+        new_body = {"msgtype": "m.text", "body": "I've been edited!"}
+        channel = self._send_relation(
+            RelationTypes.REPLACE,
+            "m.room.message",
+            content={"msgtype": "m.text", "body": "foo", "m.new_content": new_body},
+            parent_id=threaded_event_id,
+        )
+
+        # Fetch the thread root, to get the bundled aggregation for the thread.
+        relations_dict = self._get_bundled_aggregations()
+
+        # We expect that the edit message appears in the thread summary in the
+        # unsigned relations section.
+        self.assertIn(RelationTypes.THREAD, relations_dict)
+
+        thread_summary = relations_dict[RelationTypes.THREAD]
+        self.assertIn("latest_event", thread_summary)
+        latest_event_in_thread = thread_summary["latest_event"]
+        self.assertEqual(latest_event_in_thread["content"]["body"], "I've been edited!")
 
     def test_aggregation_get_event_for_annotation(self) -> None:
         """Test that annotations do not get bundled aggregations included
@@ -1093,7 +1087,7 @@ class BundledAggregationsTestCase(BaseRelationsTestCase):
         channel = self._send_relation(RelationTypes.THREAD, "m.room.test")
         thread_id = channel.json_body["event_id"]
 
-        # Annotate the annotation.
+        # Annotate the thread.
         self._send_relation(
             RelationTypes.ANNOTATION, "m.reaction", "a", parent_id=thread_id
         )
diff --git a/tests/util/test_async_helpers.py b/tests/util/test_async_helpers.py
index e5bc416de..daacc54c7 100644
--- a/tests/util/test_async_helpers.py
+++ b/tests/util/test_async_helpers.py
@@ -382,7 +382,7 @@ class StopCancellationTests(TestCase):
 class DelayCancellationTests(TestCase):
     """Tests for the `delay_cancellation` function."""
 
-    def test_cancellation(self):
+    def test_deferred_cancellation(self):
         """Test that cancellation of the new `Deferred` waits for the original."""
         deferred: "Deferred[str]" = Deferred()
         wrapper_deferred = delay_cancellation(deferred)
@@ -403,6 +403,35 @@ class DelayCancellationTests(TestCase):
         # Now that the original `Deferred` has failed, we should get a `CancelledError`.
         self.failureResultOf(wrapper_deferred, CancelledError)
 
+    def test_coroutine_cancellation(self):
+        """Test that cancellation of the new `Deferred` waits for the original."""
+        blocking_deferred: "Deferred[None]" = Deferred()
+        completion_deferred: "Deferred[None]" = Deferred()
+
+        async def task():
+            await blocking_deferred
+            completion_deferred.callback(None)
+            # Raise an exception. Twisted should consume it, otherwise unwanted
+            # tracebacks will be printed in logs.
+            raise ValueError("abc")
+
+        wrapper_deferred = delay_cancellation(task())
+
+        # Cancel the new `Deferred`.
+        wrapper_deferred.cancel()
+        self.assertNoResult(wrapper_deferred)
+        self.assertFalse(
+            blocking_deferred.called, "Cancellation was propagated too deep"
+        )
+        self.assertFalse(completion_deferred.called)
+
+        # Unblock the task.
+        blocking_deferred.callback(None)
+        self.assertTrue(completion_deferred.called)
+
+        # Now that the original coroutine 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.
 
@@ -451,7 +480,7 @@ class DelayCancellationTests(TestCase):
         async def outer():
             with LoggingContext("c") as c:
                 try:
-                    await delay_cancellation(defer.ensureDeferred(inner()))
+                    await delay_cancellation(inner())
                     self.fail("`CancelledError` was not raised")
                 except CancelledError:
                     self.assertEqual(c, current_context())