From 23c15040010345430129546ab702f35cec31d08a Mon Sep 17 00:00:00 2001 From: h Date: Mon, 19 Oct 2026 02:21:42 +0000 Subject: [PATCH] feat(userbot): after a restart the userbot catches up on what it missed while down, from a saved cursor and within bounds --- ...py => f8c2d6e1a4b7_capture_reliability.py} | 14 +- .../src/userbot/modules/catchup/__init__.py | 3 + .../src/userbot/modules/catchup/repository.py | 41 +++ .../src/userbot/modules/catchup/service.py | 255 ++++++++++++++++++ backend/src/userbot/runner.py | 21 +- 5 files changed, 327 insertions(+), 7 deletions(-) rename backend/migrations/versions/{f8c2d6e1a4b7_job_run_after.py => f8c2d6e1a4b7_capture_reliability.py} (58%) create mode 100644 backend/src/userbot/modules/catchup/__init__.py create mode 100644 backend/src/userbot/modules/catchup/repository.py create mode 100644 backend/src/userbot/modules/catchup/service.py diff --git a/backend/migrations/versions/f8c2d6e1a4b7_job_run_after.py b/backend/migrations/versions/f8c2d6e1a4b7_capture_reliability.py similarity index 58% rename from backend/migrations/versions/f8c2d6e1a4b7_job_run_after.py rename to backend/migrations/versions/f8c2d6e1a4b7_capture_reliability.py index 14cef95..2f73e0b 100644 --- a/backend/migrations/versions/f8c2d6e1a4b7_job_run_after.py +++ b/backend/migrations/versions/f8c2d6e1a4b7_capture_reliability.py @@ -1,4 +1,4 @@ -"""job run after +"""capture reliability Revision ID: f8c2d6e1a4b7 Revises: c9e2a7f4b1d8 @@ -21,7 +21,19 @@ def upgrade() -> None: op.add_column( "jobs", sa.Column("run_after", sa.DateTime(timezone=True), nullable=True) ) + op.execute( + "ALTER TABLE accounts " + "ADD COLUMN catchup_pts integer, " + "ADD COLUMN catchup_qts integer, " + "ADD COLUMN catchup_date integer" + ) def downgrade() -> None: + op.execute( + "ALTER TABLE accounts " + "DROP COLUMN IF EXISTS catchup_pts, " + "DROP COLUMN IF EXISTS catchup_qts, " + "DROP COLUMN IF EXISTS catchup_date" + ) op.drop_column("jobs", "run_after") diff --git a/backend/src/userbot/modules/catchup/__init__.py b/backend/src/userbot/modules/catchup/__init__.py new file mode 100644 index 0000000..0a2fa3f --- /dev/null +++ b/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"] diff --git a/backend/src/userbot/modules/catchup/repository.py b/backend/src/userbot/modules/catchup/repository.py new file mode 100644 index 0000000..97de17e --- /dev/null +++ b/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) diff --git a/backend/src/userbot/modules/catchup/service.py b/backend/src/userbot/modules/catchup/service.py new file mode 100644 index 0000000..46ce185 --- /dev/null +++ b/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" + ) diff --git a/backend/src/userbot/runner.py b/backend/src/userbot/runner.py index 2982b4b..549bc1c 100644 --- a/backend/src/userbot/runner.py +++ b/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(