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

This commit is contained in:
hh committed 2026-09-19 16:18:40 +00:00
1 parent 59d3f214e8
commit 468f72ed11
4 files changed
+154 -30

No files matched your search

@@ -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"
)
+38 -5
View File
@@ -35,6 +35,28 @@ SET enabled = false, status = 'removed', status_detail = NULL, updated_at = now(
WHERE bot_id = $1 AND enabled 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 = """ _SET_STATUS = """
UPDATE watched_bots SET UPDATE watched_bots SET
status = $2, 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: async def mark_backfilled(pool: asyncpg.Pool, bot_id: int) -> None:
await pool.execute( await pool.execute(_MARK_BACKFILLED, bot_id)
"UPDATE watched_bots SET backfilled_at = now(), updated_at = now() "
"WHERE bot_id = $1",
bot_id,
)
+56 -21
View File
@@ -7,6 +7,8 @@ same queue Bot API ``getUpdates`` drains is reachable over MTProto through
""" """
import time import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from pyrogram import raw, types from pyrogram import raw, types
from pyrogram.errors import RPCError from pyrogram.errors import RPCError
@@ -19,7 +21,31 @@ from utils.logging import logger
EARLIEST_PTS = 1 EARLIEST_PTS = 1
EARLIEST_QTS = 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( async def _store(
@@ -50,37 +76,41 @@ async def _parse_all(
return parsed return parsed
async def drain_backlog(client: PyroClient, account_id: int) -> int | None: async def drain_backlog(
"""Walk the bot's pending queue into the archive. 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 ``save`` is handed every cursor the walk moves past, so an interrupted
be completed - the caller keeps the bot marked as still needing a drain. 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 ctx = client.capture
if ctx is None: if ctx is None:
return None return DrainResult(cursor, 0, done=False)
pts, qts, date = EARLIEST_PTS, EARLIEST_QTS, int(time.time())
stored = 0 stored = 0
for _ in range(MAX_SLICES): while True:
try: try:
diff = await client.invoke( 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: except RPCError as exc:
logger.warning( logger.warning(
f"[yellow]Account {account_id}: getDifference failed:[/] {exc}" f"[yellow]Account {account_id}: getDifference failed:[/] {exc}"
) )
return None return DrainResult(cursor, stored, done=False)
if isinstance(diff, raw.types.updates.DifferenceEmpty): if isinstance(diff, raw.types.updates.DifferenceEmpty):
logger.info(f"Account {account_id}: backlog empty at pts={pts}") logger.info(f"Account {account_id}: backlog empty at pts={cursor.pts}")
return stored return DrainResult(cursor, stored, done=True)
if isinstance(diff, raw.types.updates.DifferenceTooLong): if isinstance(diff, raw.types.updates.DifferenceTooLong):
logger.info( logger.info(
f"Account {account_id}: backlog too long, server pts={diff.pts}; " f"Account {account_id}: backlog too long, server pts={diff.pts}; "
f"nothing older is retrievable" f"nothing older is retrievable"
) )
return stored return DrainResult(cursor, stored, done=True)
for message in await _parse_all( for message in await _parse_all(
client, diff.new_messages, diff.users, diff.chats 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) stored += await _store(client, ctx, message)
if isinstance(diff, raw.types.updates.Difference): if isinstance(diff, raw.types.updates.Difference):
cursor = _advance(cursor, diff.state)
await save(cursor)
logger.info( 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 = ( moved = _advance(cursor, diff.intermediate_state)
diff.intermediate_state.pts, if moved == cursor:
diff.intermediate_state.qts, logger.warning(
diff.intermediate_state.date, f"[yellow]Account {account_id}: backlog stopped advancing[/] "
f"at pts={cursor.pts}"
) )
logger.warning(f"[yellow]Account {account_id}: backlog walk hit the slice cap[/]") return DrainResult(cursor, stored, done=False)
return stored cursor = moved
await save(cursor)
+26 -4
View File
@@ -22,11 +22,13 @@ from utils.read import accounts, bots
from utils.storage import ContentAddressedStorage from utils.storage import ContentAddressedStorage
from watcher import tokens from watcher import tokens
from watcher.context import build_bot_capture_context 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" DEVICE_MODEL = "Beavergram Bot"
TOKEN_MISSING = "Токен не найден на диске" # noqa: S105 TOKEN_MISSING = "Токен не найден на диске" # noqa: S105
INVALID_TOKEN_ERRORS = ("ACCESS_TOKEN_INVALID", "ACCESS_TOKEN_EXPIRED") INVALID_TOKEN_ERRORS = ("ACCESS_TOKEN_INVALID", "ACCESS_TOKEN_EXPIRED")
BACKLOG_ATTEMPTS = 3
BACKLOG_RETRY_PAUSE = 60
@dataclass @dataclass
@@ -184,10 +186,30 @@ class BotRegistry:
) )
async def _backlog(self, client: PyroClient, bot_id: int, account_id: int) -> None: async def _backlog(self, client: PyroClient, bot_id: int, account_id: int) -> None:
stored = await drain_backlog(client, account_id) async def save(at: DrainCursor) -> None:
if stored is None: await bots.save_backlog_cursor(self._pool, bot_id, at.pts, at.qts, at.date)
return
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) 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: if stored:
await accounts.notify_accounts_changed(self._pool) await accounts.notify_accounts_changed(self._pool)