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
|
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,
|
|
||||||
)
|
|
||||||
@@ -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)
|
||||||
@@ -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)
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user