microblog/app/utils/workers.py

80 lines
2.7 KiB
Python
Raw Normal View History

2022-08-10 18:39:19 +00:00
import asyncio
import signal
from typing import Generic
from typing import TypeVar
from loguru import logger
from app.database import AsyncSession
from app.database import async_session
T = TypeVar("T")
class Worker(Generic[T]):
def __init__(self, workers_count: int) -> None:
self._loop = asyncio.get_event_loop()
self._in_flight: set[int] = set()
self._queue: asyncio.Queue[T] = asyncio.Queue(maxsize=1)
self._stop_event = asyncio.Event()
self._workers_count = workers_count
async def _consumer(self, db_session: AsyncSession) -> None:
while not self._stop_event.is_set():
message = await self._queue.get()
try:
await self.process_message(db_session, message)
finally:
self._in_flight.remove(message.id) # type: ignore
self._queue.task_done()
async def _producer(self, db_session: AsyncSession) -> None:
while not self._stop_event.is_set():
next_message = await self.get_next_message(db_session)
if next_message:
self._in_flight.add(next_message.id) # type: ignore
await self._queue.put(next_message)
else:
await asyncio.sleep(1)
async def process_message(self, db_session: AsyncSession, message: T) -> None:
raise NotImplementedError
async def get_next_message(self, db_session: AsyncSession) -> T | None:
raise NotImplementedError
async def startup(self, db_session: AsyncSession) -> None:
return None
def in_flight_ids(self) -> set[int]:
return self._in_flight
async def run_forever(self) -> None:
signals = (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)
for s in signals:
self._loop.add_signal_handler(
s,
lambda s=s: asyncio.create_task(self._shutdown(s)),
)
async with async_session() as db_session:
await self.startup(db_session)
self._loop.create_task(self._producer(db_session))
for _ in range(self._workers_count):
self._loop.create_task(self._consumer(db_session))
await self._stop_event.wait()
logger.info("Waiting for tasks to finish")
await self._queue.join()
tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
logger.info(f"Cancelling {len(tasks)} tasks")
[task.cancel() for task in tasks]
await asyncio.gather(*tasks, return_exceptions=True)
logger.info("stopping loop")
async def _shutdown(self, sig: signal.Signals) -> None:
logger.info(f"Caught {signal=}")
self._stop_event.set()