Expose debug API with tracemalloc

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 pickle
+import tracemalloc
 
 from aiohttp import web
 from telethon.errors import SessionPasswordNeededError
@@ -90,6 +92,18 @@ 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 +804,91 @@ 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,
+        )