256 lines
8.9 KiB
Python
256 lines
8.9 KiB
Python
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"
|
|
)
|