Skip to content

Commit

Permalink
Clean up typing for executor pool
Browse files Browse the repository at this point in the history
  • Loading branch information
linuxdaemon committed Apr 9, 2024
1 parent 14985bc commit 00db07d
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 12 deletions.
6 changes: 5 additions & 1 deletion cloudbot/bot.py
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,11 @@ def __init__(
self.db_engine = create_engine(db_path)
database.configure(self.db_engine)
self.db_executor_pool = ExecutorPool(
50, max_workers=1, thread_name_prefix="cloudbot-db", loop=self.loop
50,
max_workers=1,
thread_name_prefix="cloudbot-db",
loop=self.loop,
executor_type=ThreadPoolExecutor,
)

logger.debug("Database system initialised.")
Expand Down
29 changes: 19 additions & 10 deletions cloudbot/util/executor_pool.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
from asyncio import AbstractEventLoop
import logging
import os
import random
from concurrent.futures import ThreadPoolExecutor
from asyncio import AbstractEventLoop
from concurrent.futures import Executor
from typing import Generic, List, Optional, Type, TypeVar

from cloudbot.util.async_util import create_future

Expand All @@ -26,9 +27,17 @@ def __del__(self):
self.release()


class ExecutorPool:
T = TypeVar("T", bound=Executor)


class ExecutorPool(Generic[T]):
def __init__(
self, max_executors=None, executor_type=ThreadPoolExecutor, *, loop:AbstractEventLoop, **kwargs
self,
max_executors: Optional[int] = None,
*,
executor_type: Type[T],
loop: AbstractEventLoop,
**kwargs,
) -> None:
if max_executors is None:
max_executors = (os.cpu_count() or 1) * 5
Expand All @@ -40,14 +49,14 @@ def __init__(
self._exec_class = executor_type
self._exec_args = kwargs

self._executors = []
self._free_executors = []
self._executors: List[T] = []
self._free_executors: List[T] = []
self._executor_waiter = create_future(loop)

def get(self):
def get(self) -> ExecutorWrapper:
return ExecutorWrapper(self, self._get())

def _get(self):
def _get(self) -> T:
if not self._free_executors:
if len(self._executors) < self._max:
return self._add_executor()
Expand All @@ -56,10 +65,10 @@ def _get(self):

return self._free_executors.pop()

def release_executor(self, executor):
def release_executor(self, executor: T) -> None:
self._free_executors.append(executor)

def _add_executor(self):
def _add_executor(self) -> T:
exc = self._exec_class(**self._exec_args)
self._executors.append(exc)

Expand Down
7 changes: 6 additions & 1 deletion tests/util/mock_bot.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import logging
from concurrent.futures import ThreadPoolExecutor
from typing import Awaitable, Dict, Optional

from watchdog.observers import Observer
Expand All @@ -23,7 +24,11 @@ def __init__(
):
if loop:
self.db_executor_pool = ExecutorPool(
50, max_workers=1, thread_name_prefix="cloudbot-db", loop=loop
50,
max_workers=1,
thread_name_prefix="cloudbot-db",
loop=loop,
executor_type=ThreadPoolExecutor,
)
else:
self.db_executor_pool = None
Expand Down

0 comments on commit 00db07d

Please sign in to comment.