diff --git a/module/webui/app.py b/module/webui/app.py index abb82a610..2795d05b4 100644 --- a/module/webui/app.py +++ b/module/webui/app.py @@ -8,7 +8,7 @@ from functools import partial 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 @@ from pywebio.output import ( 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.base import Frame 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 @@ from module.webui.widgets import ( put_output, ) +patch_executor() task_handler = TaskHandler() diff --git a/module/webui/patch.py b/module/webui/patch.py new file mode 100644 index 000000000..09c0b0a71 --- /dev/null +++ b/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)