From 2a40886fd60bfc34d178d9ca39e9aac0e30a3e40 Mon Sep 17 00:00:00 2001
From: Rongrong <i@rong.moe>
Date: Mon, 6 Nov 2023 02:49:11 +0800
Subject: [PATCH] feat: yappi profiling

Signed-off-by: Rongrong <i@rong.moe>
---
 requirements.txt |  4 ++++
 src/__init__.py  | 27 +++++++++++++++++++++++++--
 src/env.py       | 11 +++++++++++
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 0e732e372e..96e767214c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -38,3 +38,7 @@ cachetools==5.3.1
 CJKwrap==2.2
 typing-extensions==4.8.0
 uvloop==0.17.0; sys_platform!='win32' and sys_platform!='cygwin' and sys_platform!='cli'
+
+# dev requirements
+# TODO: drop requirements.txt and use pyproject.toml
+yappi==1.4.0
diff --git a/src/__init__.py b/src/__init__.py
index 548f84c8a8..e67d065d21 100644
--- a/src/__init__.py
+++ b/src/__init__.py
@@ -27,10 +27,22 @@
 from . import log, db, command
 from .i18n import i18n, ALL_LANGUAGES, get_commands_list
 from .parsing import tgraph
+from .compat import nullcontext
 
 # log
 logger = log.getLogger('RSStT')
 
+if env.PROFILING:
+    try:
+        import yappi
+        yappi.set_clock_type(env.PROFILING if isinstance(env.PROFILING, str) else 'cpu')
+        logger.info(f'Profiling mode, clock type: {yappi.get_clock_type()}')
+    except ImportError:
+        logger.warning('yappi is not installed, profiling is disabled.')
+        yappi = None
+else:
+    yappi = None
+
 loop = env.loop
 bot: Optional[TelegramClient] = None
 pre_tasks = []
@@ -308,9 +320,10 @@ def main():
                           trigger=CronTrigger(minute='*', second=env.CRON_SECOND, timezone='UTC'),
                           max_instances=10,
                           misfire_grace_time=10)
-        scheduler.start()
+        with yappi.run(profile_threads=False) if yappi else nullcontext():
+            scheduler.start()
+            loop.run_until_complete(bot.disconnected)
 
-        loop.run_until_complete(bot.disconnected)
     except (KeyboardInterrupt, SystemExit) as e:
         logger.error(f'Received {type(e).__name__}, exiting...', exc_info=e)
         exit_code = e.code if isinstance(e, SystemExit) and e.code is not None else 0
@@ -318,6 +331,16 @@ def main():
         logger.critical('Uncaught error:', exc_info=e)
         exit_code = 99
     finally:
+        if yappi:
+            try:
+                prof = yappi.get_func_stats()
+                prof.print_all()
+                import time
+                curr_time = int(time.time())
+                for typ in ('pstat', 'ystat'):
+                    prof.save(os.path.join(env.config_folder_path, f'prof_{curr_time}.{typ}'), typ)
+            except Exception as e:
+                logger.error('Error when exporting profiling stats:', exc_info=e)
         try:
             if getattr(signal, 'SIGALRM', None):
                 signal.alarm(15)
diff --git a/src/env.py b/src/env.py
index 91c69bca3c..0348bb3651 100644
--- a/src/env.py
+++ b/src/env.py
@@ -250,6 +250,7 @@ def __list_parser(var: Optional[str]) -> list[str]:
 LAZY_MEDIA_VALIDATION: Final = __bool_parser(os.environ.get('LAZY_MEDIA_VALIDATION'))
 NO_UVLOOP: Final = __bool_parser(os.environ.get('NO_UVLOOP'))
 MULTIPROCESSING: Final = __bool_parser(os.environ.get('MULTIPROCESSING'))
+
 DEBUG: Final = __bool_parser(os.environ.get('DEBUG'))
 __configure_logging(  # config twice to make .env file work
     level=colorlog.DEBUG if DEBUG else colorlog.INFO,
@@ -258,6 +259,16 @@ def __list_parser(var: Optional[str]) -> list[str]:
 if DEBUG:
     logger.debug('DEBUG mode enabled')
 
+_profiling = os.environ.get('PROFILING')
+if _profiling is None:
+    PROFILING: Final = False
+else:
+    _profiling = _profiling.lower()
+    if _profiling in ('cpu', 'wall'):
+        PROFILING: Final = _profiling
+    else:
+        PROFILING: Final = __bool_parser(_profiling)
+
 # ----- environment config -----
 RAILWAY_STATIC_URL: Final = os.environ.get('RAILWAY_STATIC_URL')
 PORT: Final = int(os.environ.get('PORT', 0)) or (8080 if RAILWAY_STATIC_URL else None)