diff --git a/backend/src/watcher/registry.py b/backend/src/watcher/registry.py index ebcf179..039f21c 100644 --- a/backend/src/watcher/registry.py +++ b/backend/src/watcher/registry.py @@ -36,6 +36,7 @@ class RunningBot: client: PyroClient account_id: int tasks: list[asyncio.Task] + backlog: asyncio.Task | None = None async def _cancel(task: asyncio.Task) -> None: @@ -105,6 +106,8 @@ class BotRegistry: await self._stop(bot_id) for bot_id in enabled - set(self._running): await self._start(bot_id) + for bot_id in enabled & set(self._running): + await self._resume_backlog(bot_id) for bot_id in await bots.removed_bot_ids(self._pool): tokens.drop(bot_id) @@ -177,14 +180,24 @@ class BotRegistry: await accounts.notify_accounts_changed(self._pool) consumer = JobConsumer(client, self._pool, account_id) - tasks = [asyncio.create_task(consumer.run())] - if await bots.needs_backfill(self._pool, bot_id): - tasks.append(asyncio.create_task(self._backlog(client, bot_id, account_id))) - self._running[bot_id] = RunningBot(client, account_id, tasks) + self._running[bot_id] = RunningBot( + client, account_id, [asyncio.create_task(consumer.run())] + ) + await self._resume_backlog(bot_id) logger.info( f"[green]Bot account {account_id} watched:[/] @{me.username} ({bot_id})" ) + async def _resume_backlog(self, bot_id: int) -> None: + running = self._running[bot_id] + if running.backlog is not None and not running.backlog.done(): + return + if not await bots.needs_backfill(self._pool, bot_id): + return + running.backlog = asyncio.create_task( + self._backlog(running.client, bot_id, running.account_id) + ) + async def _backlog(self, client: PyroClient, bot_id: int, account_id: int) -> None: async def save(at: DrainCursor) -> None: await bots.save_backlog_cursor(self._pool, bot_id, at.pts, at.qts, at.date) @@ -217,8 +230,9 @@ class BotRegistry: running = self._running.pop(bot_id, None) if running is None: return - for task in running.tasks: - await _cancel(task) + for task in (*running.tasks, running.backlog): + if task is not None: + await _cancel(task) with contextlib.suppress(Exception): await running.client.stop() logger.info(f"[yellow]Bot unwatched:[/] {bot_id}")