fix: prevent subscription worker deadlocks

This commit is contained in:
3252a8
2026-05-29 22:07:16 +03:00
parent 001e54cfe2
commit cab963dcdc
5 changed files with 157 additions and 9 deletions
+2
View File
@@ -14,6 +14,7 @@ from bot.middlewares.i18n import JsonI18n
from bot.services.panel_api_service import PanelApiService from bot.services.panel_api_service import PanelApiService
from bot.utils.text_sanitizer import panel_description_from_profile from bot.utils.text_sanitizer import panel_description_from_profile
from config.settings import Settings 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.dal import panel_sync_dal, subscription_dal, user_dal
from db.models import Subscription, User from db.models import Subscription, User
@@ -928,6 +929,7 @@ async def _perform_sync_impl(
total_panel_users = len(panel_users_data) total_panel_users = len(panel_users_data)
logging.info(f"Starting sync for {total_panel_users} panel users.") 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) sync_indexes = await _prefetch_sync_indexes(session, panel_users_data)
users_by_telegram_id = sync_indexes["users_by_telegram_id"] users_by_telegram_id = sync_indexes["users_by_telegram_id"]
users_by_user_id = sync_indexes["users_by_user_id"] users_by_user_id = sync_indexes["users_by_user_id"]
@@ -20,6 +20,7 @@ from bot.services.subscription_lifecycle_notifications import (
) )
from bot.services.subscription_service import SubscriptionService from bot.services.subscription_service import SubscriptionService
from config.settings import Settings from config.settings import Settings
from db.advisory_locks import acquire_subscription_background_sync_lock
from db.dal import subscription_dal from db.dal import subscription_dal
from db.models import Subscription from db.models import Subscription
@@ -67,6 +68,7 @@ class SubscriptionNotificationWorker:
else: else:
started = time.monotonic() started = time.monotonic()
async with self.session_factory() as session: async with self.session_factory() as session:
await acquire_subscription_background_sync_lock(session)
await self.expiry_tick(session) await self.expiry_tick(session)
await self.trial_traffic_tick(session) await self.trial_traffic_tick(session)
await session.commit() await session.commit()
+82 -9
View File
@@ -2,7 +2,7 @@ import asyncio
import logging import logging
import time import time
from datetime import datetime, timezone from datetime import datetime, timezone
from typing import Any, Optional from typing import Any, Awaitable, Callable, Optional
from aiogram import Bot from aiogram import Bot
from aiogram.types import InlineKeyboardButton, InlineKeyboardMarkup, WebAppInfo 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.date_utils import month_start
from bot.utils.mini_app_url import subscription_mini_app_topup_url from bot.utils.mini_app_url import subscription_mini_app_topup_url
from config.settings import Settings 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.dal import subscription_dal, tariff_dal, user_dal
from db.models import Subscription from db.models import Subscription
@@ -31,6 +32,10 @@ TARIFF_WORKER_BATCH_SIZE = 50
TARIFF_WORKER_PANEL_CONCURRENCY = 10 TARIFF_WORKER_PANEL_CONCURRENCY = 10
TARIFF_WORKER_BULK_PANEL_FETCH_THRESHOLD = 50 TARIFF_WORKER_BULK_PANEL_FETCH_THRESHOLD = 50
TARIFF_WORKER_SQUAD_CONFIRMATION_CACHE_TTL_SECONDS = 900 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: class TariffTrafficWorker:
@@ -111,12 +116,14 @@ class TariffTrafficWorker:
logging.info("TariffTrafficWorker tick skipped: Redis lock is held") logging.info("TariffTrafficWorker tick skipped: Redis lock is held")
else: else:
started = time.monotonic() started = time.monotonic()
async with self.session_factory() as session: await self._run_db_tick_with_retry(
await self.traffic_period_tick(session) "traffic_period",
await session.commit() self.traffic_period_tick,
async with self.session_factory() as session: )
await self.legacy_throttle_recovery_tick(session) await self._run_db_tick_with_retry(
await session.commit() "legacy_throttle_recovery",
self.legacy_throttle_recovery_tick,
)
logging.info( logging.info(
"metric worker_tick_duration_seconds=%.3f worker=tariff", "metric worker_tick_duration_seconds=%.3f worker=tariff",
time.monotonic() - started, time.monotonic() - started,
@@ -134,16 +141,80 @@ class TariffTrafficWorker:
def stop(self) -> None: def stop(self) -> None:
self._stopped.set() 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: async def traffic_period_tick(self, session: AsyncSession) -> None:
now = datetime.now(timezone.utc) now = datetime.now(timezone.utc)
self._premium_node_usage_tick_cache = {} self._premium_node_usage_tick_cache = {}
warning_period_start = month_start(now) warning_period_start = month_start(now)
result = await session.execute( result = await session.execute(
select(Subscription).where( select(Subscription)
.where(
Subscription.is_active == True, Subscription.is_active == True,
Subscription.end_date > now, Subscription.end_date > now,
Subscription.tariff_key.is_not(None), Subscription.tariff_key.is_not(None),
) )
.order_by(Subscription.subscription_id.asc())
) )
subs = list(result.scalars().all()) subs = list(result.scalars().all())
if not subs: if not subs:
@@ -1159,10 +1230,12 @@ class TariffTrafficWorker:
from Internal Squads. from Internal Squads.
""" """
result = await session.execute( result = await session.execute(
select(Subscription).where( select(Subscription)
.where(
Subscription.is_active == True, Subscription.is_active == True,
Subscription.is_throttled == True, Subscription.is_throttled == True,
) )
.order_by(Subscription.subscription_id.asc())
) )
for sub in result.scalars().all(): for sub in result.scalars().all():
try: try:
+12
View File
@@ -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},
)
+59
View File
@@ -34,6 +34,65 @@ def _tariffs_config_payload() -> dict:
class TariffWorkerTests(unittest.IsolatedAsyncioTestCase): 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): async def test_period_tariff_uses_panel_month_strategy_without_resetting(self):
with tempfile.TemporaryDirectory() as tmpdir: with tempfile.TemporaryDirectory() as tmpdir:
config_path = Path(tmpdir) / "tariffs.json" config_path = Path(tmpdir) / "tariffs.json"