diff --git a/backend/bot/handlers/admin/sync_admin.py b/backend/bot/handlers/admin/sync_admin.py index f192be1..bf79818 100644 --- a/backend/bot/handlers/admin/sync_admin.py +++ b/backend/bot/handlers/admin/sync_admin.py @@ -14,6 +14,7 @@ from bot.middlewares.i18n import JsonI18n from bot.services.panel_api_service import PanelApiService from bot.utils.text_sanitizer import panel_description_from_profile from config.settings import Settings +from db.advisory_locks import acquire_subscription_background_sync_lock from db.dal import panel_sync_dal, subscription_dal, user_dal from db.models import Subscription, User @@ -928,6 +929,7 @@ async def _perform_sync_impl( total_panel_users = len(panel_users_data) logging.info(f"Starting sync for {total_panel_users} panel users.") + await acquire_subscription_background_sync_lock(session) sync_indexes = await _prefetch_sync_indexes(session, panel_users_data) users_by_telegram_id = sync_indexes["users_by_telegram_id"] users_by_user_id = sync_indexes["users_by_user_id"] diff --git a/backend/bot/services/subscription_notification_worker.py b/backend/bot/services/subscription_notification_worker.py index 6fadb04..67d5f07 100644 --- a/backend/bot/services/subscription_notification_worker.py +++ b/backend/bot/services/subscription_notification_worker.py @@ -20,6 +20,7 @@ from bot.services.subscription_lifecycle_notifications import ( ) from bot.services.subscription_service import SubscriptionService from config.settings import Settings +from db.advisory_locks import acquire_subscription_background_sync_lock from db.dal import subscription_dal from db.models import Subscription @@ -67,6 +68,7 @@ class SubscriptionNotificationWorker: else: started = time.monotonic() async with self.session_factory() as session: + await acquire_subscription_background_sync_lock(session) await self.expiry_tick(session) await self.trial_traffic_tick(session) await session.commit() diff --git a/backend/bot/services/tariff_worker.py b/backend/bot/services/tariff_worker.py index 0711fd3..5764fb0 100644 --- a/backend/bot/services/tariff_worker.py +++ b/backend/bot/services/tariff_worker.py @@ -2,7 +2,7 @@ import asyncio import logging import time from datetime import datetime, timezone -from typing import Any, Optional +from typing import Any, Awaitable, Callable, Optional from aiogram import Bot from aiogram.types import InlineKeyboardButton, InlineKeyboardMarkup, WebAppInfo @@ -18,6 +18,7 @@ from bot.services.subscription_service import SubscriptionService from bot.utils.date_utils import month_start from bot.utils.mini_app_url import subscription_mini_app_topup_url from config.settings import Settings +from db.advisory_locks import acquire_subscription_background_sync_lock from db.dal import subscription_dal, tariff_dal, user_dal from db.models import Subscription @@ -31,6 +32,10 @@ TARIFF_WORKER_BATCH_SIZE = 50 TARIFF_WORKER_PANEL_CONCURRENCY = 10 TARIFF_WORKER_BULK_PANEL_FETCH_THRESHOLD = 50 TARIFF_WORKER_SQUAD_CONFIRMATION_CACHE_TTL_SECONDS = 900 +TARIFF_WORKER_DB_RETRY_ATTEMPTS = 3 +TARIFF_WORKER_DB_RETRY_BASE_SLEEP_SECONDS = 0.5 +POSTGRES_RETRYABLE_SQLSTATES = {"40001", "40P01"} +POSTGRES_RETRYABLE_ERROR_NAMES = {"DeadlockDetectedError", "SerializationError"} class TariffTrafficWorker: @@ -111,12 +116,14 @@ class TariffTrafficWorker: logging.info("TariffTrafficWorker tick skipped: Redis lock is held") else: started = time.monotonic() - async with self.session_factory() as session: - await self.traffic_period_tick(session) - await session.commit() - async with self.session_factory() as session: - await self.legacy_throttle_recovery_tick(session) - await session.commit() + await self._run_db_tick_with_retry( + "traffic_period", + self.traffic_period_tick, + ) + await self._run_db_tick_with_retry( + "legacy_throttle_recovery", + self.legacy_throttle_recovery_tick, + ) logging.info( "metric worker_tick_duration_seconds=%.3f worker=tariff", time.monotonic() - started, @@ -134,16 +141,80 @@ class TariffTrafficWorker: def stop(self) -> None: self._stopped.set() + async def _run_db_tick_with_retry( + self, + tick_name: str, + tick: Callable[[AsyncSession], Awaitable[None]], + ) -> None: + for attempt in range(1, TARIFF_WORKER_DB_RETRY_ATTEMPTS + 1): + async with self.session_factory() as session: + try: + await acquire_subscription_background_sync_lock(session) + await tick(session) + await session.commit() + return + except Exception as exc: + await session.rollback() + if ( + attempt < TARIFF_WORKER_DB_RETRY_ATTEMPTS + and self._is_retryable_db_exception(exc) + ): + delay = TARIFF_WORKER_DB_RETRY_BASE_SLEEP_SECONDS * attempt + logging.warning( + "TariffTrafficWorker %s retrying after database concurrency " + "error, attempt %s/%s: %s", + tick_name, + attempt + 1, + TARIFF_WORKER_DB_RETRY_ATTEMPTS, + exc, + ) + await asyncio.sleep(delay) + continue + raise + + @staticmethod + def _is_retryable_db_exception(exc: BaseException) -> bool: + pending: list[BaseException] = [exc] + seen: set[int] = set() + while pending: + current = pending.pop() + current_id = id(current) + if current_id in seen: + continue + seen.add(current_id) + + sqlstate = getattr(current, "sqlstate", None) or getattr(current, "pgcode", None) + if sqlstate in POSTGRES_RETRYABLE_SQLSTATES: + return True + + error_name = type(current).__name__ + message = str(current).lower() + if ( + error_name in POSTGRES_RETRYABLE_ERROR_NAMES + or "deadlock detected" in message + or "could not serialize access" in message + ): + return True + + for attr in ("orig", "__cause__", "__context__"): + nested = getattr(current, attr, None) + if isinstance(nested, BaseException): + pending.append(nested) + + return False + async def traffic_period_tick(self, session: AsyncSession) -> None: now = datetime.now(timezone.utc) self._premium_node_usage_tick_cache = {} warning_period_start = month_start(now) result = await session.execute( - select(Subscription).where( + select(Subscription) + .where( Subscription.is_active == True, Subscription.end_date > now, Subscription.tariff_key.is_not(None), ) + .order_by(Subscription.subscription_id.asc()) ) subs = list(result.scalars().all()) if not subs: @@ -1159,10 +1230,12 @@ class TariffTrafficWorker: from Internal Squads. """ result = await session.execute( - select(Subscription).where( + select(Subscription) + .where( Subscription.is_active == True, Subscription.is_throttled == True, ) + .order_by(Subscription.subscription_id.asc()) ) for sub in result.scalars().all(): try: diff --git a/backend/db/advisory_locks.py b/backend/db/advisory_locks.py new file mode 100644 index 0000000..a0c6305 --- /dev/null +++ b/backend/db/advisory_locks.py @@ -0,0 +1,12 @@ +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncSession + +# Serializes background jobs that rewrite subscription rows from panel state. +SUBSCRIPTION_BACKGROUND_SYNC_LOCK_ID = 817512404897421338 + + +async def acquire_subscription_background_sync_lock(session: AsyncSession) -> None: + await session.execute( + text("SELECT pg_advisory_xact_lock(:lock_id)"), + {"lock_id": SUBSCRIPTION_BACKGROUND_SYNC_LOCK_ID}, + ) diff --git a/tests/test_tariff_worker.py b/tests/test_tariff_worker.py index aa71c91..ae63e44 100644 --- a/tests/test_tariff_worker.py +++ b/tests/test_tariff_worker.py @@ -34,6 +34,65 @@ def _tariffs_config_payload() -> dict: class TariffWorkerTests(unittest.IsolatedAsyncioTestCase): + async def test_db_tick_retries_deadlock_once(self): + class FakeSession: + def __init__(self): + self.execute = AsyncMock() + self.commit = AsyncMock() + self.rollback = AsyncMock() + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + sessions = [] + + def session_factory(): + session = FakeSession() + sessions.append(session) + return session + + attempts = 0 + + async def tick(_session): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RuntimeError("deadlock detected") + + worker = TariffTrafficWorker( + settings=SimpleNamespace(), + session_factory=session_factory, + panel_service=SimpleNamespace(), + subscription_service=SimpleNamespace(), + ) + + with patch("bot.services.tariff_worker.asyncio.sleep", new=AsyncMock()) as sleep: + await worker._run_db_tick_with_retry("test", tick) + + self.assertEqual(attempts, 2) + self.assertEqual(len(sessions), 2) + sessions[0].rollback.assert_awaited_once() + sessions[0].commit.assert_not_awaited() + sessions[1].commit.assert_awaited_once() + sleep.assert_awaited_once() + + async def test_retryable_db_exception_detects_wrapped_sqlstate(self): + class PgError(Exception): + sqlstate = "40P01" + + class WrappedDbError(Exception): + def __init__(self, orig): + super().__init__("wrapped") + self.orig = orig + + self.assertTrue( + TariffTrafficWorker._is_retryable_db_exception(WrappedDbError(PgError())) + ) + self.assertFalse(TariffTrafficWorker._is_retryable_db_exception(RuntimeError("plain"))) + async def test_period_tariff_uses_panel_month_strategy_without_resetting(self): with tempfile.TemporaryDirectory() as tmpdir: config_path = Path(tmpdir) / "tariffs.json"