127 lines
4.0 KiB
Python
127 lines
4.0 KiB
Python
import asyncio
|
|
import logging
|
|
from datetime import UTC, datetime, timedelta
|
|
from typing import TYPE_CHECKING
|
|
|
|
from sqlalchemy import delete, select, update
|
|
from sqlalchemy.dialects.postgresql import insert
|
|
|
|
from config import cfg
|
|
from db.models import RWSyncOutbox, Subscription
|
|
from db.session import UnitOfWork, async_session
|
|
from external.rw import sync_subscription_by_telegram_id
|
|
|
|
if TYPE_CHECKING:
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
LOCK_DURATION = timedelta(minutes=5)
|
|
MAX_RETRY_DELAY = timedelta(hours=1)
|
|
|
|
|
|
async def enqueue_rw_sync(session: "AsyncSession", subscription_id: int) -> None:
|
|
now = datetime.now(UTC)
|
|
stmt = insert(RWSyncOutbox).values(
|
|
subscription_id=subscription_id,
|
|
revision=1,
|
|
next_attempt_at=now,
|
|
)
|
|
await session.execute(
|
|
stmt.on_conflict_do_update(
|
|
index_elements=[RWSyncOutbox.subscription_id],
|
|
set_={
|
|
"revision": RWSyncOutbox.revision + 1,
|
|
"next_attempt_at": now,
|
|
"locked_until": None,
|
|
},
|
|
)
|
|
)
|
|
|
|
|
|
async def enqueue_all_rw_syncs(session: "AsyncSession") -> int:
|
|
subscription_ids = await session.scalars(select(Subscription.id))
|
|
count = 0
|
|
for subscription_id in subscription_ids:
|
|
await enqueue_rw_sync(session, subscription_id)
|
|
count += 1
|
|
return count
|
|
|
|
|
|
async def process_rw_syncs(uow: UnitOfWork, batch_size: int = 50) -> int:
|
|
session = uow.session
|
|
now = datetime.now(UTC)
|
|
jobs = await session.scalars(
|
|
select(RWSyncOutbox)
|
|
.where(
|
|
RWSyncOutbox.next_attempt_at <= now,
|
|
(RWSyncOutbox.locked_until.is_(None)) | (RWSyncOutbox.locked_until <= now),
|
|
)
|
|
.order_by(RWSyncOutbox.next_attempt_at)
|
|
.with_for_update(skip_locked=True)
|
|
.limit(batch_size)
|
|
)
|
|
jobs = list(jobs)
|
|
for job in jobs:
|
|
job.locked_until = now + LOCK_DURATION
|
|
await uow.commit()
|
|
|
|
for job in jobs:
|
|
subscription = job.subscription
|
|
user = subscription.user
|
|
synced = await sync_subscription_by_telegram_id(
|
|
expires_at=subscription.expires_at,
|
|
devices=subscription.devices,
|
|
telegram_id=user.telegram_id,
|
|
username=user.username,
|
|
)
|
|
if synced:
|
|
await session.execute(
|
|
delete(RWSyncOutbox).where(
|
|
RWSyncOutbox.id == job.id,
|
|
RWSyncOutbox.revision == job.revision,
|
|
)
|
|
)
|
|
await uow.commit()
|
|
continue
|
|
|
|
delay = min(timedelta(minutes=2**job.attempts), MAX_RETRY_DELAY)
|
|
await session.execute(
|
|
update(RWSyncOutbox)
|
|
.where(RWSyncOutbox.id == job.id, RWSyncOutbox.revision == job.revision)
|
|
.values(
|
|
attempts=RWSyncOutbox.attempts + 1,
|
|
next_attempt_at=datetime.now(UTC) + delay,
|
|
locked_until=None,
|
|
)
|
|
)
|
|
await uow.commit()
|
|
logger.warning("RW sync failed for subscription_id=%s", subscription.id)
|
|
|
|
return len(jobs)
|
|
|
|
|
|
async def run_rw_sync_worker() -> None:
|
|
interval = cfg.rw_sync_interval_minutes * 60
|
|
while True:
|
|
try:
|
|
async with async_session() as session, UnitOfWork(session) as uow:
|
|
await process_rw_syncs(uow)
|
|
except Exception:
|
|
logger.exception("RW sync worker iteration failed")
|
|
await asyncio.sleep(interval)
|
|
|
|
|
|
async def run_rw_sync_reconciler() -> None:
|
|
interval = cfg.rw_sync_reconcile_interval_minutes * 60
|
|
while True:
|
|
try:
|
|
async with async_session() as session:
|
|
async with UnitOfWork(session) as uow:
|
|
count = await enqueue_all_rw_syncs(uow.session)
|
|
await uow.commit()
|
|
logger.info("Queued %s subscriptions for RW reconciliation", count)
|
|
except Exception:
|
|
logger.exception("RW reconciliation iteration failed")
|
|
await asyncio.sleep(interval)
|