From c6620ad70afce9407c16793de8ab4fea35523d81 Mon Sep 17 00:00:00 2001
From: retoor <retoor@molodetz.nl>
Date: Tue, 18 Feb 2025 13:29:33 +0100
Subject: [PATCH] Added profiler.

---
 src/snek/app.py             |  7 +++----
 src/snek/system/profiler.py | 42 +++++++++++++++++++++++++++++++++++++
 src/snek/view/rpc.py        |  5 ++++-
 3 files changed, 49 insertions(+), 5 deletions(-)
 create mode 100644 src/snek/system/profiler.py

diff --git a/src/snek/app.py b/src/snek/app.py
index ab337ba..0fc1016 100644
--- a/src/snek/app.py
+++ b/src/snek/app.py
@@ -2,6 +2,7 @@ import pathlib
 import asyncio
 
 import logging 
+
 logging.basicConfig(level=logging.DEBUG)
 
 from aiohttp import web
@@ -32,7 +33,7 @@ from snek.view.status import StatusView
 from snek.view.web import WebView
 from snek.view.upload import UploadView
 from snek.view.search_user import SearchUserView 
-
+from snek.system.profiler import profiler_handler
 
 SESSION_KEY = b"c79a0c5fda4b424189c427d28c9f7c34"
 
@@ -43,7 +44,6 @@ async def session_middleware(request, handler):
     response = await handler(request)
     return response
 
-
 class Application(BaseApplication):
 
     def __init__(self, *args, **kwargs):
@@ -77,6 +77,7 @@ class Application(BaseApplication):
             name="static",
             show_index=True,
         )
+        self.router.add_view("/profiler.html", profiler_handler)
         self.router.add_view("/about.html", AboutHTMLView)
         self.router.add_view("/about.md", AboutMDView)
         self.router.add_view("/logout.json", LogoutView)
@@ -128,8 +129,6 @@ class Application(BaseApplication):
 
 async def main():
     app = Application(db_path="sqlite:///snek.db")
-    loop = asyncio.get_event_loop()
-    loop.set_debug(True)
     await web._run_app(app, port=8081, host="0.0.0.0")
 
 if __name__ == "__main__":
diff --git a/src/snek/system/profiler.py b/src/snek/system/profiler.py
new file mode 100644
index 0000000..824208a
--- /dev/null
+++ b/src/snek/system/profiler.py
@@ -0,0 +1,42 @@
+import cProfile
+import pstats
+import sys 
+from aiohttp import web
+profiler = None
+import io 
+
+
+@web.middleware
+async def profile_middleware(request, handler):
+    global profiler 
+    if not profiler:
+        profiler = cProfile.Profile()
+    profiler.enable()
+    response = await handler(request)
+    profiler.disable()
+    stats = pstats.Stats(profiler, stream=sys.stdout)
+    stats.sort_stats('cumulative')
+    stats.print_stats() 
+    return response
+
+async def profiler_handler(request):
+    output = io.StringIO()
+    stats = pstats.Stats(profiler, stream=output)
+    stats.sort_stats('cumulative')
+    stats.print_stats()
+    return web.Response(text=output.getvalue())
+
+class Profiler:
+
+    def __init__(self):
+        global profiler 
+        if profiler is None:
+            profiler = cProfile.Profile()
+        self.profiler = profiler
+
+    async def __aenter__(self):
+        self.profiler.enable()    
+    
+    async def __aexit__(self, *args, **kwargs):
+        self.profiler.disable()
+
diff --git a/src/snek/view/rpc.py b/src/snek/view/rpc.py
index 07893e1..05e447d 100644
--- a/src/snek/view/rpc.py
+++ b/src/snek/view/rpc.py
@@ -12,6 +12,7 @@ from snek.system.view import BaseView
 import traceback
 import json
 from snek.system.model import now 
+from snek.system.profiler import Profiler
 
 class RPCView(BaseView):
 
@@ -169,8 +170,10 @@ class RPCView(BaseView):
         async for msg in ws:
             if msg.type == web.WSMsgType.TEXT:
                 try:
-                    await rpc(msg.json())
+                    async with Profiler():
+                        await rpc(msg.json())
                 except Exception as ex:
+                    print(ex, flush=True)
                     await self.services.socket.delete(ws)
                     break
             elif msg.type == web.WSMsgType.ERROR: