Opt: Patch aiofiles to reduce thread pool size

Author: LmeSzinc

module/webui/app.py

@@ -8,7 +8,7 @@
 from typing import Dict, List, Optional
 
 from pywebio import config as webconfig
-from pywebio.input import file_upload, input_group, input, select
+from pywebio.input import file_upload, input, input_group, select
 from pywebio.output import (
     Output,
     clear,
@@ -32,15 +32,7 @@
     use_scope,
 )
 from pywebio.pin import pin, pin_on_change
-from pywebio.session import (
-    go_app,
-    info,
-    local,
-    register_thread,
-    run_js,
-    set_env,
-    download,
-)
+from pywebio.session import (download, go_app, info, local, register_thread, run_js, set_env)
 
 import module.webui.lang as lang
 from module.config.config import AzurLaneConfig, Function
@@ -64,6 +56,7 @@
 from module.webui.discord_presence import close_discord_rpc, init_discord_rpc
 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
@@ -94,6 +87,7 @@
     put_output,
 )
 
+patch_executor()
 task_handler = TaskHandler()
 
 

module/webui/patch.py

@@ -0,0 +1,36 @@
+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
+    """
+    loop = asyncio.get_event_loop()
+    loop.set_default_executor(CachedThreadPoolExecutor.executor)