fix(watcher): pick an unfinished backlog back up without restarting the bot
This commit is contained in:
1 parent
468f72ed11
commit
b60815d042
1 file changed
+20
-6
@@ -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}")
|
||||
Reference in new issue
Block a user