Expose debug API with tracemalloc

Runs along the provisioning API with same authentication.

Enabling PYTHONTRACEMALLOC is extremely heavy so we have our custom
env MAUTRIXTRACEMALLOC that enables tracing just when the bridge
starts to prevent counting all the python runtime init memory.

mautrix_telegram/__main__.py

@@ -16,6 +16,8 @@
 from __future__ import annotations
 
 from typing import Any
+from os import getenv
+import tracemalloc
 
 from telethon import __version__ as __telethon_version__
 
@@ -142,4 +144,7 @@ def manhole_banner_program_version(self) -> str:
         return f"{super().manhole_banner_program_version} and Telethon {__telethon_version__}"
 
 
+if getenv("MAUTRIXTRACEMALLOC"):
+    tracemalloc.start()
+
 TelegramBridge().run()

mautrix_telegram/config.py

@@ -83,6 +83,7 @@ def do_update(self, helper: ConfigUpdateHelper) -> None:
         copy("appservice.provisioning.shared_secret")
         if base["appservice.provisioning.shared_secret"] == "generate":
             base["appservice.provisioning.shared_secret"] = self._new_token()
+        copy("appservice.provisioning.debug_endpoints")
 
         if "pool_size" in base["appservice.database_opts"]:
             pool_size = base["appservice.database_opts"].pop("pool_size")

mautrix_telegram/example-config.yaml

@@ -74,6 +74,8 @@ appservice:
         # The shared secret to authorize users of the API.
         # Set to "generate" to generate and save a new token.
         shared_secret: generate
+        # Enable debug API at {prefix}/debug with provisioning authentication.
+        debug_endpoints: false
 
     # The unique ID of this appservice.
     id: telegram

mautrix_telegram/web/provisioning/__init__.py

@@ -20,6 +20,8 @@
 import datetime
 import json
 import logging
+import tracemalloc
+import pickle
 
 from aiohttp import web
 from telethon.errors import SessionPasswordNeededError
@@ -90,6 +92,14 @@ def __init__(self, bridge: "TelegramBridge") -> None:
 
         self.app.router.add_route("GET", "/v1/bridge", self.bridge_info)
 
+        if bridge.config["appservice.provisioning.debug_endpoints"]:
+            self.log.info("Enabling provisioning debug endpoints")
+            self.app.router.add_route("POST", "/debug/tracemalloc", self.post_tracemalloc)
+            self.app.router.add_route("DELETE", "/debug/tracemalloc", self.delete_tracemalloc)
+            self.app.router.add_route("GET", "/debug/tracemalloc", self.get_tracemalloc)
+            self.app.router.add_route("GET", "/debug/tracemalloc/snapshot", self.get_tracemalloc_snapshot)
+            self.app.router.add_route("GET", "/debug/tracemalloc/statistics/{key}", self.get_tracemalloc_statistics)
+
     async def get_portal_by_mxid(self, request: web.Request) -> web.Response:
         err = self.check_authorization(request)
         if err is not None:
@@ -790,3 +800,89 @@ async def get_user_request_info(
         user, err = await self.get_user(mxid, expect_logged_in, require_puppeting)
 
         return data, user, err
+
+    async def post_tracemalloc(self, request: web.Request) -> web.Response:
+        err = self.check_authorization(request)
+        if err is not None:
+            return err
+
+        tracemalloc.start()
+
+        return web.Response(
+            status=201,
+        )
+
+    async def delete_tracemalloc(self, request: web.Request) -> web.Response:
+        err = self.check_authorization(request)
+        if err is not None:
+            return err
+
+        tracemalloc.stop()
+
+        return web.Response(
+            status=201,
+        )
+
+    async def get_tracemalloc(self, request: web.Request) -> web.Response:
+        err = self.check_authorization(request)
+        if err is not None:
+            return err
+
+        current_memory, peak_memory = tracemalloc.get_traced_memory()
+
+        return web.json_response(
+            {
+                "is_tracing": tracemalloc.is_tracing(),
+                "traced_memory": {
+                    "current": current_memory,
+                    "peak": peak_memory,
+                },
+                "tracemalloc_memory": tracemalloc.get_tracemalloc_memory(),
+            },
+            status=200,
+        )
+
+    async def get_tracemalloc_statistics(self, request: web.Request) -> web.Response:
+        err = self.check_authorization(request)
+        if err is not None:
+            return err
+
+        try:
+            snapshot = tracemalloc.take_snapshot()
+        except RuntimeError:
+            return web.Response(status=400)
+
+        key = request.match_info["key"]
+
+        if key not in ["filename", "lineno", "traceback"]:
+            return web.Response(status=400)
+
+        stats = snapshot.statistics(key)
+
+        response = []
+        for stat in stats:
+            response.append({
+                "count": stat.count,
+                "size": stat.size,
+                "traceback": stat.traceback.format(),
+            })
+
+        return web.json_response(
+            response,
+            status=200,
+        )
+
+    async def get_tracemalloc_snapshot(self, request: web.Request) -> web.Response:
+        err = self.check_authorization(request)
+        if err is not None:
+            return err
+
+        try:
+            snapshot = tracemalloc.take_snapshot()
+        except RuntimeError:
+            return web.Response(status=400)
+
+        return web.Response(
+            body=pickle.dumps(snapshot, pickle.HIGHEST_PROTOCOL),
+            status=200,
+        )