perf: patch aiofiles to reduce thread pool size

module/webui/app.py

@@ -53,6 +53,7 @@ from module.webui.fake import (
 )
 from module.webui.fastapi import asgi_app
 from module.webui.lang import _t, t
+from module.webui.patch import patch_executor
 from module.webui.pin import put_input, put_select
 from module.webui.process_manager import ProcessManager
 from module.webui.remote_access import RemoteAccess
@@ -83,6 +84,7 @@ from module.webui.widgets import (
     put_output,
 )
 
+patch_executor()
 task_handler = TaskHandler()
 
 

module/webui/patch.py

@@ -0,0 +1,41 @@
+import asyncio
+from functools import partial, wraps
+
+from module.logger import logger
+from module.webui.setting import cached_class_property
+
+
+class CachedThreadPoolExecutor:
+    @cached_class_property
+    def executor(cls):
+        from concurrent.futures.thread import ThreadPoolExecutor
+        pool = ThreadPoolExecutor(max_workers=5)
+        logger.info('Patched ThreadPoolExecutor created')
+        return pool
+
+
+def wrap(func):
+    @wraps(func)
+    async def run(*args, loop=None, executor=None, **kwargs):
+        if loop is None:
+            loop = asyncio.get_event_loop()
+        if executor is None:
+            executor = CachedThreadPoolExecutor.executor
+        pfunc = partial(func, *args, **kwargs)
+        return await loop.run_in_executor(executor, pfunc)
+
+    return run
+
+
+def patch_executor():
+    """
+    Limit pool size in loop.run_in_executor
+    so starlette.staticfiles -> aiofiles won't create tons of threads
+    """
+    try:
+        import aiofiles
+    except ImportError:
+        return
+
+    loop = asyncio.get_event_loop()
+    loop.set_default_executor(CachedThreadPoolExecutor.executor)