Skip to content

Commit

Permalink
Opt: Patch aiofiles to reduce thread pool size
Browse files Browse the repository at this point in the history
  • Loading branch information
LmeSzinc committed Sep 11, 2024
1 parent 3d262a5 commit f362a63
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 10 deletions.
14 changes: 4 additions & 10 deletions module/webui/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -94,6 +87,7 @@
put_output,
)

patch_executor()
task_handler = TaskHandler()


Expand Down
36 changes: 36 additions & 0 deletions module/webui/patch.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit f362a63

Please sign in to comment.