Add wrapper for creating background tasks

The wrapper ensures that the reference to the task is not lost before
it completes, and also logs uncaught errors.

Author: tulir

mautrix/appservice/as_handler.py

@@ -27,6 +27,7 @@
     SerializerError,
     UserID,
 )
+from mautrix.util import background_task
 
 HandlerFunc = Callable[[Event], Awaitable]
 
@@ -314,7 +315,7 @@ async def try_handle(handler_func: HandlerFunc):
 
         for handler in self.event_handlers:
             # TODO add option to handle events synchronously
-            asyncio.create_task(try_handle(handler))
+            background_task.create(try_handle(handler))
 
     def matrix_event_handler(self, func: HandlerFunc) -> HandlerFunc:
         self.event_handlers.append(func)

mautrix/bridge/custom_puppet.py

@@ -41,6 +41,7 @@
     SyncToken,
     UserID,
 )
+from mautrix.util import background_task
 
 from .. import bridge as br
 
@@ -409,7 +410,7 @@ def _handle_sync(self, sync_resp: dict) -> None:
 
         # Deserialize and handle all events
         for event in chain(ephemeral_events, presence_events):
-            asyncio.create_task(self.mx.try_handle_sync_event(Event.deserialize(event)))
+            background_task.create(self.mx.try_handle_sync_event(Event.deserialize(event)))
 
     async def _try_sync(self) -> None:
         try:

mautrix/bridge/matrix.py

@@ -56,7 +56,7 @@
     Version,
     VersionsResponse,
 )
-from mautrix.util import markdown
+from mautrix.util import background_task, markdown
 from mautrix.util.logging import TraceLogger
 from mautrix.util.message_send_checkpoint import (
     CHECKPOINT_TYPES,
@@ -382,7 +382,7 @@ async def handle_puppet_nonportal_invite(
             members = await intent.get_room_members(room_id)
         except MatrixError:
             self.log.exception(f"Failed to get state after joining {room_id} as {intent.mxid}")
-            asyncio.create_task(intent.leave_room(room_id, reason="Internal error"))
+            background_task.create(intent.leave_room(room_id, reason="Internal error"))
             return
         if create_evt.type == RoomType.SPACE:
             await self.handle_puppet_space_invite(room_id, puppet, invited_by, evt)
@@ -798,7 +798,7 @@ async def _handle_encrypted_wait(
             f"Couldn't find session {err.session_id} trying to decrypt {evt.event_id},"
             " waiting even longer"
         )
-        asyncio.create_task(
+        background_task.create(
             self.e2ee.crypto.request_room_key(
                 evt.room_id,
                 evt.content.sender_key,
@@ -875,7 +875,7 @@ def _send_message_checkpoint(
             info=str(err) if err else None,
             retry_num=retry_num,
         )
-        asyncio.create_task(checkpoint.send(endpoint, self.az.as_token, self.log))
+        background_task.create(checkpoint.send(endpoint, self.az.as_token, self.log))
 
     allowed_event_classes: tuple[type, ...] = (
         MessageEvent,

mautrix/bridge/portal.py

@@ -30,6 +30,7 @@
     TextMessageEventContent,
     UserID,
 )
+from mautrix.util import background_task
 from mautrix.util.logging import TraceLogger
 from mautrix.util.simple_lock import SimpleLock
 
@@ -402,7 +403,7 @@ async def restart_scheduled_disappearing(cls) -> None:
         for msg in msgs:
             portal = await cls.bridge.get_portal(msg.room_id)
             if portal and portal.mxid:
-                asyncio.create_task(portal._disappear_event(msg))
+                background_task.create(portal._disappear_event(msg))
             else:
                 await msg.delete()
 
@@ -418,7 +419,7 @@ async def schedule_disappearing(self) -> None:
             for msg in msgs:
                 msg.start_timer()
                 await msg.update()
-                asyncio.create_task(self._disappear_event(msg))
+                background_task.create(self._disappear_event(msg))
 
     async def _send_message(
         self,
@@ -431,7 +432,7 @@ async def _send_message(
             event_type, content = await self.matrix.e2ee.encrypt(self.mxid, event_type, content)
         event_id = await intent.send_message_event(self.mxid, event_type, content, **kwargs)
         if intent.api.is_real_user:
-            asyncio.create_task(intent.mark_read(self.mxid, event_id))
+            background_task.create(intent.mark_read(self.mxid, event_id))
         return event_id
 
     @property

mautrix/bridge/user.py

@@ -16,6 +16,7 @@
 from mautrix.appservice import AppService
 from mautrix.errors import MNotFound
 from mautrix.types import EventID, EventType, Membership, MessageType, RoomID, UserID
+from mautrix.util import background_task
 from mautrix.util.bridge_state import BridgeState, BridgeStateEvent
 from mautrix.util.logging import TraceLogger
 from mautrix.util.message_send_checkpoint import (
@@ -244,7 +245,7 @@ def send_remote_checkpoint(
         """
         if not self.bridge.config["homeserver.message_send_checkpoint_endpoint"]:
             return WrappedTask(task=None)
-        task = asyncio.create_task(
+        task = background_task.create(
             MessageSendCheckpoint(
                 event_id=event_id,
                 room_id=room_id,

mautrix/client/api/modules/media_repository.py

@@ -20,6 +20,7 @@
     MXOpenGraph,
     SerializerError,
 )
+from mautrix.util import background_task
 from mautrix.util.async_body import async_iter_bytes
 from mautrix.util.opt_prometheus import Histogram
 
@@ -157,7 +158,7 @@ async def _try_upload():
                 except Exception as e:
                     self.log.error(f"Failed to upload {mxc}: {type(e).__name__}: {e}")
 
-            asyncio.create_task(_try_upload())
+            background_task.create(_try_upload())
             return mxc
         else:
             with self._observe_upload_time(size):

mautrix/crypto/decrypt_olm.py

@@ -19,6 +19,7 @@
     ToDeviceEvent,
     UserID,
 )
+from mautrix.util import background_task
 
 from .base import BaseOlmMachine
 from .sessions import Session
@@ -74,19 +75,19 @@ async def _decrypt_olm_ciphertext(
                 f"Found matching session yet decryption failed for sender {sender}"
                 f" with key {sender_key}"
             )
-            asyncio.create_task(self._unwedge_session(sender, sender_key))
+            background_task.create(self._unwedge_session(sender, sender_key))
             raise
 
         if not plaintext:
             if message.type != OlmMsgType.PREKEY:
-                asyncio.create_task(self._unwedge_session(sender, sender_key))
+                background_task.create(self._unwedge_session(sender, sender_key))
                 raise DecryptionError("Decryption failed for normal message")
 
             self.log.trace(f"Trying to create inbound session for {sender}/{sender_key}")
             try:
                 session = await self._create_inbound_session(sender_key, message.body)
             except olm.OlmSessionError as e:
-                asyncio.create_task(self._unwedge_session(sender, sender_key))
+                background_task.create(self._unwedge_session(sender, sender_key))
                 raise DecryptionError("Failed to create new session from prekey message") from e
             self.log.debug(
                 f"Created inbound session {session.id} for {sender} (sender key: {sender_key})"

mautrix/crypto/machine.py

@@ -24,6 +24,7 @@
     TrustState,
     UserID,
 )
+from mautrix.util import background_task
 from mautrix.util.logging import TraceLogger
 
 from .account import OlmAccount
@@ -109,7 +110,7 @@ async def handle_as_otk_counts(
                     self.log.warning(f"Got OTK count for unknown device {user_id}/{device_id}")
 
     async def handle_as_device_lists(self, device_lists: DeviceLists) -> None:
-        asyncio.create_task(self.handle_device_lists(device_lists))
+        background_task.create(self.handle_device_lists(device_lists))
 
     async def handle_as_to_device_event(self, evt: ASToDeviceEvent) -> None:
         if evt.to_user_id != self.client.mxid or evt.to_device_id != self.client.device_id:

mautrix/util/__init__.py

@@ -8,6 +8,7 @@
     # File modules
     "async_body",
     "async_getter_lock",
+    "background_task",
     "bridge_state",
     "color_log",
     "ffmpeg",

mautrix/util/background_task.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2023 Tulir Asokan
+#
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+from __future__ import annotations
+
+from typing import Coroutine
+import asyncio
+import logging
+
+_tasks = set()
+log = logging.getLogger("mau.background_task")
+
+
+async def catch(coro: Coroutine, caller: str) -> None:
+    try:
+        await coro
+    except Exception:
+        log.exception(f"Uncaught error in background task (created in {caller})")
+
+
+# Logger.findCaller finds the 3rd stack frame, so add an intermediate function
+# to get the caller of create().
+def _find_caller() -> tuple[str, int, str, None]:
+    return log.findCaller()
+
+
+def create(coro: Coroutine, *, name: str | None = None, catch_errors: bool = True) -> asyncio.Task:
+    """
+    Create a background asyncio task safely, ensuring a reference is kept until the task completes.
+    It also catches and logs uncaught errors (unless disabled via the parameter).
+
+    Args:
+        coro: The coroutine to wrap in a task and execute.
+        name: An optional name for the created task.
+        catch_errors: Should the task be wrapped in a try-except block to log any uncaught errors?
+
+    Returns:
+        An asyncio Task object wrapping the given coroutine.
+    """
+    if catch_errors:
+        try:
+            file_name, line_number, function_name, _ = _find_caller()
+            caller = f"{function_name} at {file_name}:{line_number}"
+        except ValueError:
+            caller = "unknown function"
+        task = asyncio.create_task(catch(coro, caller), name=name)
+    else:
+        task = asyncio.create_task(coro, name=name)
+    _tasks.add(task)
+    task.add_done_callback(_tasks.discard)
+    return task