feat(userbot): after a restart the userbot catches up on what it missed while down, from a saved cursor and within bounds

backend/src/userbot/modules/catchup/__init__.py

@@ -0,0 +1,3 @@
+from userbot.modules.catchup.service import SavedState, catch_up, read_saved_state
+
+__all__ = ["SavedState", "catch_up", "read_saved_state"]

backend/src/userbot/modules/catchup/repository.py

@@ -0,0 +1,41 @@
+from dataclasses import dataclass
+
+import asyncpg
+
+_LOAD = """
+SELECT catchup_pts, catchup_qts, catchup_date FROM accounts WHERE account_id = $1
+"""
+
+_SAVE = """
+UPDATE accounts SET catchup_pts = $2, catchup_qts = $3, catchup_date = $4
+WHERE account_id = $1
+"""
+
+_CLEAR = """
+UPDATE accounts SET catchup_pts = NULL, catchup_qts = NULL, catchup_date = NULL
+WHERE account_id = $1
+"""
+
+
+@dataclass(frozen=True)
+class UpdateCursor:
+    pts: int
+    qts: int
+    date: int
+
+
+async def load_cursor(pool: asyncpg.Pool, account_id: int) -> UpdateCursor | None:
+    row = await pool.fetchrow(_LOAD, account_id)
+    if row is None or row["catchup_pts"] is None:
+        return None
+    return UpdateCursor(row["catchup_pts"], row["catchup_qts"], row["catchup_date"])
+
+
+async def save_cursor(
+    pool: asyncpg.Pool, account_id: int, cursor: UpdateCursor
+) -> None:
+    await pool.execute(_SAVE, account_id, cursor.pts, cursor.qts, cursor.date)
+
+
+async def clear_cursor(pool: asyncpg.Pool, account_id: int) -> None:
+    await pool.execute(_CLEAR, account_id)

backend/src/userbot/modules/catchup/service.py

@@ -0,0 +1,255 @@
+import asyncio
+import sqlite3
+import time
+from contextlib import closing
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from pyrogram import raw, utils
+from pyrogram.errors import FloodPremiumWait, FloodWait, RPCError
+
+from userbot.modules.capture.chat_meta import meta_from_peer
+from userbot.modules.capture.context import CaptureContext
+from userbot.modules.catchup import repository
+from userbot.modules.catchup.repository import UpdateCursor
+from userbot.modules.client import PyroClient
+from utils.logging import logger
+from utils.policy.models import ChatKind
+
+BOX_SLICES = 50
+CHANNEL_SLICES = 10
+CHANNEL_LIMIT = 100
+MAX_CHANNELS = 50
+QUEUE_HIGH_WATER = 200
+QUEUE_WAIT_SECONDS = 0.5
+NO_QTS = -1
+MAX_AGE_SECONDS = 3 * 24 * 3600
+
+
+@dataclass(frozen=True)
+class SavedState:
+    box: UpdateCursor | None = None
+    channels: dict[int, int] = field(default_factory=dict)
+
+
+@dataclass
+class BoxWalk:
+    cursor: UpdateCursor
+    done: bool = False
+    updates: int = 0
+    channels: dict[int, raw.base.Chat] = field(default_factory=dict)
+
+
+def read_saved_state(path: Path) -> SavedState:
+    try:
+        uri = f"{path.resolve().as_uri()}?mode=ro"
+        with closing(sqlite3.connect(uri, uri=True)) as conn:
+            rows = conn.execute(
+                "SELECT id, pts, qts, date FROM update_state"
+            ).fetchall()
+    except Exception as exc:
+        logger.warning(f"[yellow]Session {path.stem}: no saved update state ({exc})[/]")
+        return SavedState()
+    box = None
+    channels: dict[int, int] = {}
+    for state_id, pts, qts, date in rows:
+        if pts is None:
+            continue
+        if state_id == 0:
+            box = UpdateCursor(
+                pts, NO_QTS if qts is None else qts, date or int(time.time())
+            )
+        else:
+            channels[state_id] = pts
+    return SavedState(box, channels)
+
+
+async def _dispatch(
+    client: PyroClient,
+    new_messages: list,
+    other_updates: list,
+    users: list,
+    chats: list,
+) -> int:
+    await client.fetch_peers(users)
+    await client.fetch_peers(chats)
+    user_map = {user.id: user for user in users}
+    chat_map = {chat.id: chat for chat in chats}
+    queue = client.dispatcher.updates_queue
+    for message in new_messages:
+        update = raw.types.UpdateNewMessage(message=message, pts=0, pts_count=0)
+        queue.put_nowait((update, user_map, chat_map))
+    for update in other_updates:
+        queue.put_nowait((update, user_map, chat_map))
+    while queue.qsize() > QUEUE_HIGH_WATER:  # noqa: ASYNC110
+        await asyncio.sleep(QUEUE_WAIT_SECONDS)
+    return len(new_messages) + len(other_updates)
+
+
+def _channel_hints(other_updates: list, chats: list) -> dict[int, raw.base.Chat]:
+    by_id = {chat.id: chat for chat in chats}
+    return {
+        utils.get_channel_id(update.channel_id): by_id[update.channel_id]
+        for update in other_updates
+        if isinstance(update, raw.types.UpdateChannelTooLong)
+        and update.channel_id in by_id
+    }
+
+
+async def _walk_box(
+    client: PyroClient, ctx: CaptureContext, cursor: UpdateCursor
+) -> BoxWalk:
+    walk = BoxWalk(cursor)
+    for _ in range(BOX_SLICES):
+        diff = await client.invoke(
+            raw.functions.updates.GetDifference(
+                pts=walk.cursor.pts, date=walk.cursor.date, qts=walk.cursor.qts
+            )
+        )
+        if isinstance(diff, raw.types.updates.DifferenceEmpty):
+            walk.done = True
+            return walk
+        if isinstance(diff, raw.types.updates.DifferenceTooLong):
+            logger.warning(
+                f"[yellow]Account {ctx.account_id}: too much missed to catch up "
+                f"(pts {walk.cursor.pts} to {diff.pts}), that gap stays[/]"
+            )
+            walk.done = True
+            return walk
+        walk.updates += await _dispatch(
+            client, diff.new_messages, diff.other_updates, diff.users, diff.chats
+        )
+        walk.channels.update(_channel_hints(diff.other_updates, diff.chats))
+        final = isinstance(diff, raw.types.updates.Difference)
+        state = diff.state if final else diff.intermediate_state
+        moved = UpdateCursor(state.pts, state.qts, state.date or walk.cursor.date)
+        if moved == walk.cursor and not final:
+            logger.warning(
+                f"[yellow]Account {ctx.account_id}: catch-up stopped advancing "
+                f"at pts={moved.pts}[/]"
+            )
+            return walk
+        walk.cursor = moved
+        await repository.save_cursor(ctx.pool, ctx.account_id, moved)
+        if final:
+            walk.done = True
+            return walk
+    return walk
+
+
+async def _walk_channel(client: PyroClient, peer_id: int, pts: int) -> int:
+    peer = await client.resolve_peer(peer_id)
+    if not isinstance(peer, raw.types.InputPeerChannel):
+        return 0
+    channel = raw.types.InputChannel(
+        channel_id=peer.channel_id, access_hash=peer.access_hash
+    )
+    updates = 0
+    for _ in range(CHANNEL_SLICES):
+        diff = await client.invoke(
+            raw.functions.updates.GetChannelDifference(
+                channel=channel,
+                filter=raw.types.ChannelMessagesFilterEmpty(),
+                pts=pts,
+                limit=CHANNEL_LIMIT,
+            )
+        )
+        if isinstance(diff, raw.types.updates.ChannelDifferenceEmpty):
+            return updates
+        if isinstance(diff, raw.types.updates.ChannelDifferenceTooLong):
+            logger.info(f"Channel {peer_id}: too much missed to catch up")
+            return updates
+        updates += await _dispatch(
+            client, diff.new_messages, diff.other_updates, diff.users, diff.chats
+        )
+        if diff.final:
+            return updates
+        pts = diff.pts
+    return updates
+
+
+def _captured_channels(
+    ctx: CaptureContext, hints: dict[int, raw.base.Chat], saved: dict[int, int]
+) -> list[int]:
+    wanted: list[tuple[bool, int]] = []
+    for peer_id, chat in hints.items():
+        if peer_id not in saved:
+            continue
+        meta = meta_from_peer(
+            raw.types.PeerChannel(channel_id=chat.id),
+            {chat.id: chat},
+            {},
+            ctx.contacts.ids,
+        )
+        toggles = ctx.resolve(meta)
+        if toggles.messages or toggles.track_edits_deletes:
+            wanted.append((meta.kind is ChatKind.CHANNEL, peer_id))
+    return [peer_id for _, peer_id in sorted(wanted)[:MAX_CHANNELS]]
+
+
+async def _walk_channels(
+    client: PyroClient, ctx: CaptureContext, walk: BoxWalk, saved: dict[int, int]
+) -> int:
+    peer_ids = _captured_channels(ctx, walk.channels, saved)
+    updates = 0
+    for peer_id in peer_ids:
+        try:
+            updates += await _walk_channel(client, peer_id, saved[peer_id])
+        except (FloodWait, FloodPremiumWait) as exc:
+            logger.warning(
+                f"[yellow]Account {ctx.account_id}: channel catch-up stopped "
+                f"by a flood wait:[/] {exc}"
+            )
+            break
+        except RPCError as exc:
+            logger.warning(
+                f"[yellow]Account {ctx.account_id}: channel {peer_id} "
+                f"catch-up failed:[/] {exc}"
+            )
+    skipped = len(walk.channels) - len(peer_ids)
+    if skipped:
+        logger.info(f"Account {ctx.account_id}: {skipped} channel gap(s) left as is")
+    return updates
+
+
+async def catch_up(client: PyroClient, saved: SavedState) -> None:
+    ctx = client.capture
+    if ctx is None:
+        return
+    try:
+        pending = await repository.load_cursor(ctx.pool, ctx.account_id)
+        cursor = pending or saved.box
+        if cursor is None:
+            logger.info(f"Account {ctx.account_id}: no saved update state to catch up")
+            return
+        if time.time() - cursor.date > MAX_AGE_SECONDS:
+            logger.warning(
+                f"[yellow]Account {ctx.account_id}: last update seen too long ago "
+                f"(pts={cursor.pts}), not catching up[/]"
+            )
+            await repository.clear_cursor(ctx.pool, ctx.account_id)
+            return
+        await repository.save_cursor(ctx.pool, ctx.account_id, cursor)
+        logger.info(
+            f"Account {ctx.account_id}: catching up from pts={cursor.pts}"
+            f"{' (resumed)' if pending else ''}"
+        )
+        walk = await _walk_box(client, ctx, cursor)
+        if not walk.done:
+            logger.warning(
+                f"[yellow]Account {ctx.account_id}: catch-up unfinished at "
+                f"pts={walk.cursor.pts},[/] {walk.updates} updates replayed, "
+                f"resumes on the next start"
+            )
+            return
+        from_channels = await _walk_channels(client, ctx, walk, saved.channels)
+        await repository.clear_cursor(ctx.pool, ctx.account_id)
+        logger.info(
+            f"[green]Account {ctx.account_id}: caught up,[/] {walk.updates} updates "
+            f"from the main box, {from_channels} from channels"
+        )
+    except Exception:
+        logger.exception(
+            f"[red]Account {ctx.account_id}: catch-up failed,[/] "
+            f"resumes on the next start"
+        )

backend/src/userbot/runner.py

@@ -11,6 +11,7 @@ import uvloop
 
 from dependencies.container import container
 from userbot.modules.capture import CaptureContext, build_capture_context
+from userbot.modules.catchup import SavedState, catch_up, read_saved_state
 from userbot.modules.client import PyroClient
 from userbot.modules.jobs import JobConsumer
 from utils.env import env
@@ -35,7 +36,7 @@ SESSION_RETRY_INTERVAL = 60
 @dataclass
 class RunningAccount:
     client: PyroClient
-    consumer_task: asyncio.Task
+    tasks: list[asyncio.Task]
     device_model: str | None
 
 
@@ -56,6 +57,7 @@ class AccountRegistry:
         self._pool = pool
         self._storage = storage
         self._running: dict[str, RunningAccount] = {}
+        self._saved: dict[str, SavedState] = {}
         self._lock = asyncio.Lock()
 
     @property
@@ -88,6 +90,8 @@ class AccountRegistry:
 
     async def _start(self, path: Path, device_model: str | None) -> None:
         session_name = path.stem
+        if session_name not in self._saved:
+            self._saved[session_name] = await asyncio.to_thread(read_saved_state, path)
         client = PyroClient(
             session_name, workdir=str(path.parent), device_model=device_model
         )
@@ -110,9 +114,12 @@ class AccountRegistry:
                 await client.stop()
             return
         consumer = JobConsumer(client, self._pool, account_id)
-        self._running[session_name] = RunningAccount(
-            client, asyncio.create_task(consumer.run()), device_model
-        )
+        saved = self._saved.pop(session_name)
+        tasks = [
+            asyncio.create_task(consumer.run()),
+            asyncio.create_task(catch_up(client, saved)),
+        ]
+        self._running[session_name] = RunningAccount(client, tasks, device_model)
         logger.info(f"[green]Client started:[/] {me.full_name} ({me.id})")
         await enqueue(self._pool, account_id, "sync_dialogs", {})
         await enqueue(self._pool, account_id, "sync_contacts", {})
@@ -121,7 +128,8 @@ class AccountRegistry:
         account = self._running.pop(session_name, None)
         if account is None:
             return
-        await _cancel(account.consumer_task)
+        for task in account.tasks:
+            await _cancel(task)
         with contextlib.suppress(Exception):
             await account.client.stop()
         logger.info(f"[yellow]Client stopped:[/] {session_name}")
@@ -129,7 +137,8 @@ class AccountRegistry:
     async def _log_out(self, path: Path) -> None:
         account = self._running.pop(path.stem, None)
         if account is not None:
-            await _cancel(account.consumer_task)
+            for task in account.tasks:
+                await _cancel(task)
             client = account.client
         else:
             client = PyroClient(