From 468f72ed11a74422324c3005b1dfe8e24ccb16ed Mon Sep 17 00:00:00 2001 From: h Date: Sat, 19 Sep 2026 16:18:40 +0000 Subject: [PATCH] fix(watcher): resume the bot backlog from a saved cursor instead of restarting it --- .../b6d1f4a82e57_bot_backlog_cursor.py | 34 ++++++++ backend/src/utils/read/bots.py | 43 ++++++++-- backend/src/watcher/drain.py | 79 +++++++++++++------ backend/src/watcher/registry.py | 32 ++++++-- 4 files changed, 156 insertions(+), 32 deletions(-) create mode 100644 backend/migrations/versions/b6d1f4a82e57_bot_backlog_cursor.py diff --git a/backend/migrations/versions/b6d1f4a82e57_bot_backlog_cursor.py b/backend/migrations/versions/b6d1f4a82e57_bot_backlog_cursor.py new file mode 100644 index 0000000..19b9f91 --- /dev/null +++ b/backend/migrations/versions/b6d1f4a82e57_bot_backlog_cursor.py @@ -0,0 +1,34 @@ +"""bot backlog cursor + +Revision ID: b6d1f4a82e57 +Revises: a7f3d2b81c04 +Create Date: 2026-09-19 10:00:00.000000 + +""" + +from collections.abc import Sequence + +from alembic import op + +revision: str = "b6d1f4a82e57" +down_revision: str | None = "a7f3d2b81c04" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.execute( + "ALTER TABLE watched_bots " + "ADD COLUMN backfill_pts integer, " + "ADD COLUMN backfill_qts integer, " + "ADD COLUMN backfill_date integer" + ) + + +def downgrade() -> None: + op.execute( + "ALTER TABLE watched_bots " + "DROP COLUMN IF EXISTS backfill_pts, " + "DROP COLUMN IF EXISTS backfill_qts, " + "DROP COLUMN IF EXISTS backfill_date" + ) diff --git a/backend/src/utils/read/bots.py b/backend/src/utils/read/bots.py index d1ca70a..f389787 100644 --- a/backend/src/utils/read/bots.py +++ b/backend/src/utils/read/bots.py @@ -35,6 +35,28 @@ SET enabled = false, status = 'removed', status_detail = NULL, updated_at = now( WHERE bot_id = $1 AND enabled """ +_BACKLOG_CURSOR = """ +SELECT backfill_pts, backfill_qts, backfill_date +FROM watched_bots +WHERE bot_id = $1 AND backfill_pts IS NOT NULL +""" + +_SAVE_BACKLOG_CURSOR = """ +UPDATE watched_bots SET + backfill_pts = $2, backfill_qts = $3, backfill_date = $4, updated_at = now() +WHERE bot_id = $1 +""" + +_MARK_BACKFILLED = """ +UPDATE watched_bots SET + backfilled_at = now(), + backfill_pts = NULL, + backfill_qts = NULL, + backfill_date = NULL, + updated_at = now() +WHERE bot_id = $1 +""" + _SET_STATUS = """ UPDATE watched_bots SET status = $2, @@ -116,9 +138,20 @@ async def needs_backfill(pool: asyncpg.Pool, bot_id: int) -> bool: ) +async def backlog_cursor( + pool: asyncpg.Pool, bot_id: int +) -> tuple[int, int, int] | None: + row = await pool.fetchrow(_BACKLOG_CURSOR, bot_id) + if row is None: + return None + return row["backfill_pts"], row["backfill_qts"], row["backfill_date"] + + +async def save_backlog_cursor( + pool: asyncpg.Pool, bot_id: int, pts: int, qts: int, date: int +) -> None: + await pool.execute(_SAVE_BACKLOG_CURSOR, bot_id, pts, qts, date) + + async def mark_backfilled(pool: asyncpg.Pool, bot_id: int) -> None: - await pool.execute( - "UPDATE watched_bots SET backfilled_at = now(), updated_at = now() " - "WHERE bot_id = $1", - bot_id, - ) + await pool.execute(_MARK_BACKFILLED, bot_id) diff --git a/backend/src/watcher/drain.py b/backend/src/watcher/drain.py index 0f49e07..634a69a 100644 --- a/backend/src/watcher/drain.py +++ b/backend/src/watcher/drain.py @@ -7,6 +7,8 @@ same queue Bot API ``getUpdates`` drains is reachable over MTProto through """ import time +from collections.abc import Awaitable, Callable +from dataclasses import dataclass from pyrogram import raw, types from pyrogram.errors import RPCError @@ -19,7 +21,31 @@ from utils.logging import logger EARLIEST_PTS = 1 EARLIEST_QTS = 1 -MAX_SLICES = 50 + + +@dataclass(frozen=True) +class DrainCursor: + pts: int + qts: int + date: int + + +@dataclass(frozen=True) +class DrainResult: + cursor: DrainCursor + stored: int + done: bool + + +SaveCursor = Callable[[DrainCursor], Awaitable[None]] + + +def first_cursor() -> DrainCursor: + return DrainCursor(EARLIEST_PTS, EARLIEST_QTS, int(time.time())) + + +def _advance(cursor: DrainCursor, state: raw.types.updates.State) -> DrainCursor: + return DrainCursor(state.pts, state.qts, state.date or cursor.date) async def _store( @@ -50,37 +76,41 @@ async def _parse_all( return parsed -async def drain_backlog(client: PyroClient, account_id: int) -> int | None: - """Walk the bot's pending queue into the archive. +async def drain_backlog( + client: PyroClient, account_id: int, cursor: DrainCursor, save: SaveCursor +) -> DrainResult: + """Walk the bot's pending queue into the archive, starting from ``cursor``. - Returns the number of messages stored, or ``None`` when the walk could not - be completed - the caller keeps the bot marked as still needing a drain. + ``save`` is handed every cursor the walk moves past, so an interrupted + walk resumes where it stopped instead of starting over. ``done`` is true + only when Telegram has nothing older left to hand out. """ ctx = client.capture if ctx is None: - return None - pts, qts, date = EARLIEST_PTS, EARLIEST_QTS, int(time.time()) + return DrainResult(cursor, 0, done=False) stored = 0 - for _ in range(MAX_SLICES): + while True: try: diff = await client.invoke( - raw.functions.updates.GetDifference(pts=pts, date=date, qts=qts) + raw.functions.updates.GetDifference( + pts=cursor.pts, date=cursor.date, qts=cursor.qts + ) ) except RPCError as exc: logger.warning( f"[yellow]Account {account_id}: getDifference failed:[/] {exc}" ) - return None + return DrainResult(cursor, stored, done=False) if isinstance(diff, raw.types.updates.DifferenceEmpty): - logger.info(f"Account {account_id}: backlog empty at pts={pts}") - return stored + logger.info(f"Account {account_id}: backlog empty at pts={cursor.pts}") + return DrainResult(cursor, stored, done=True) if isinstance(diff, raw.types.updates.DifferenceTooLong): logger.info( f"Account {account_id}: backlog too long, server pts={diff.pts}; " f"nothing older is retrievable" ) - return stored + return DrainResult(cursor, stored, done=True) for message in await _parse_all( client, diff.new_messages, diff.users, diff.chats @@ -88,15 +118,20 @@ async def drain_backlog(client: PyroClient, account_id: int) -> int | None: stored += await _store(client, ctx, message) if isinstance(diff, raw.types.updates.Difference): + cursor = _advance(cursor, diff.state) + await save(cursor) logger.info( - f"[green]Account {account_id}: backlog drained,[/] {stored} messages" + f"[green]Account {account_id}: backlog reached its end[/] " + f"at pts={cursor.pts}" ) - return stored + return DrainResult(cursor, stored, done=True) - pts, qts, date = ( - diff.intermediate_state.pts, - diff.intermediate_state.qts, - diff.intermediate_state.date, - ) - logger.warning(f"[yellow]Account {account_id}: backlog walk hit the slice cap[/]") - return stored + moved = _advance(cursor, diff.intermediate_state) + if moved == cursor: + logger.warning( + f"[yellow]Account {account_id}: backlog stopped advancing[/] " + f"at pts={cursor.pts}" + ) + return DrainResult(cursor, stored, done=False) + cursor = moved + await save(cursor) diff --git a/backend/src/watcher/registry.py b/backend/src/watcher/registry.py index 8a342db..ebcf179 100644 --- a/backend/src/watcher/registry.py +++ b/backend/src/watcher/registry.py @@ -22,11 +22,13 @@ from utils.read import accounts, bots from utils.storage import ContentAddressedStorage from watcher import tokens from watcher.context import build_bot_capture_context -from watcher.drain import drain_backlog +from watcher.drain import DrainCursor, drain_backlog, first_cursor DEVICE_MODEL = "Beavergram Bot" TOKEN_MISSING = "Токен не найден на диске" # noqa: S105 INVALID_TOKEN_ERRORS = ("ACCESS_TOKEN_INVALID", "ACCESS_TOKEN_EXPIRED") +BACKLOG_ATTEMPTS = 3 +BACKLOG_RETRY_PAUSE = 60 @dataclass @@ -184,10 +186,30 @@ class BotRegistry: ) async def _backlog(self, client: PyroClient, bot_id: int, account_id: int) -> None: - stored = await drain_backlog(client, account_id) - if stored is None: - return - await bots.mark_backfilled(self._pool, bot_id) + async def save(at: DrainCursor) -> None: + await bots.save_backlog_cursor(self._pool, bot_id, at.pts, at.qts, at.date) + + saved = await bots.backlog_cursor(self._pool, bot_id) + cursor = DrainCursor(*saved) if saved else first_cursor() + stored = 0 + for attempt in range(BACKLOG_ATTEMPTS): + if attempt: + await asyncio.sleep(BACKLOG_RETRY_PAUSE) + walk = await drain_backlog(client, account_id, cursor, save) + stored += walk.stored + cursor = walk.cursor + if walk.done: + await bots.mark_backfilled(self._pool, bot_id) + logger.info( + f"[green]Account {account_id}: backlog complete,[/] " + f"{stored} messages stored" + ) + break + else: + logger.warning( + f"[yellow]Account {account_id}: backlog unfinished at " + f"pts={cursor.pts},[/] {stored} messages stored, resumes on restart" + ) if stored: await accounts.notify_accounts_changed(self._pool)