fix(watcher): resume the bot backlog from a saved cursor instead of restarting it

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)

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)