fix(watcher): resume the bot backlog from a saved cursor instead of restarting it
This commit is contained in:
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"
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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,
|
||||
moved = _advance(cursor, diff.intermediate_state)
|
||||
if moved == cursor:
|
||||
logger.warning(
|
||||
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 stored
|
||||
return DrainResult(cursor, stored, done=False)
|
||||
cursor = moved
|
||||
await save(cursor)
|
||||
@@ -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
|
||||
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)
|
||||
|
||||
|
||||
Reference in new issue
Block a user