feat: add traffic periods topups and throttling

2026-04-28 14:34:00 +03:00
parent 26af92b5ac
commit 9deb98bbe7
2 changed files with 207 additions and 0 deletions
@@ -30,6 +30,7 @@ from bot.routers import build_root_router
 from bot.services.yookassa_service import YooKassaService
 from bot.services.panel_api_service import PanelApiService
 from bot.services.subscription_service import SubscriptionService
+from bot.services.tariff_worker import TariffTrafficWorker
 from bot.services.referral_service import ReferralService
 from bot.services.promo_code_service import PromoCodeService
 from bot.services.stars_service import StarsService
@@ -273,6 +274,15 @@ async def run_bot(settings_param: Settings):
        dp[key] = service
    dp["panel_service"] = services["panel_service"]
    dp["async_session_factory"] = local_async_session_factory
+    tariff_worker = TariffTrafficWorker(
+        settings_param,
+        local_async_session_factory,
+        services["panel_service"],
+        services["subscription_service"],
+        bot,
+        i18n_instance,
+    )
+    dp["tariff_worker"] = tariff_worker

    # Wrap startup/shutdown handlers to satisfy aiogram event signature (no args passed)
    async def _on_startup_wrapper():
@@ -306,6 +316,8 @@ async def run_bot(settings_param: Settings):
        await build_and_start_web_app(dp, bot, settings_param, local_async_session_factory)

    main_tasks.append(asyncio.create_task(web_server_task(), name="AIOHTTPServerTask"))
+    if settings_param.tariffs_config:
+        main_tasks.append(asyncio.create_task(tariff_worker.run(), name="TariffTrafficWorker"))

    # Recurring billing moved to panel webhook (24h before expiry). No periodic task needed here.

@@ -0,0 +1,195 @@
+import asyncio
+import logging
+from datetime import datetime, timedelta, timezone
+from typing import Optional
+
+from aiogram import Bot
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import sessionmaker
+
+from bot.middlewares.i18n import JsonI18n
+from bot.services.panel_api_service import PanelApiService
+from bot.services.subscription_service import SubscriptionService
+from config.settings import Settings
+from db.dal import subscription_dal, tariff_dal
+from db.models import Subscription
+
+
+class TariffTrafficWorker:
+    def __init__(
+        self,
+        settings: Settings,
+        session_factory: sessionmaker,
+        panel_service: PanelApiService,
+        subscription_service: SubscriptionService,
+        bot: Optional[Bot] = None,
+        i18n: Optional[JsonI18n] = None,
+    ):
+        self.settings = settings
+        self.session_factory = session_factory
+        self.panel_service = panel_service
+        self.subscription_service = subscription_service
+        self.bot = bot
+        self.i18n = i18n
+        self._stopped = asyncio.Event()
+
+    async def run(self) -> None:
+        if not self.settings.tariffs_config:
+            return
+        while not self._stopped.is_set():
+            try:
+                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.throttle_recovery_tick(session)
+                    await session.commit()
+            except Exception:
+                logging.exception("TariffTrafficWorker tick failed")
+            try:
+                await asyncio.wait_for(self._stopped.wait(), timeout=300)
+            except asyncio.TimeoutError:
+                pass
+
+    def stop(self) -> None:
+        self._stopped.set()
+
+    async def traffic_period_tick(self, session: AsyncSession) -> None:
+        now = datetime.now(timezone.utc)
+        result = await session.execute(
+            select(Subscription).where(
+                Subscription.is_active == True,
+                Subscription.end_date > now,
+                Subscription.tariff_key.is_not(None),
+            )
+        )
+        for sub in result.scalars().all():
+            try:
+                tariff = self.settings.tariffs_config.require(sub.tariff_key)
+            except Exception:
+                continue
+            panel_data = await self.panel_service.get_user_by_uuid(sub.panel_user_uuid, log_response=False) or {}
+            used, limit, _ = self.subscription_service._extract_panel_traffic_details(panel_data)
+            if used is not None and used != sub.traffic_used_bytes:
+                sub.traffic_used_bytes = used
+            if limit is not None and limit != sub.traffic_limit_bytes:
+                sub.traffic_limit_bytes = limit
+
+            if tariff.billing_model == "period":
+                await self._maybe_reset_period(session, sub, tariff, used)
+            await self._maybe_warn_or_throttle(session, sub, tariff, used, limit)
+
+    async def _maybe_reset_period(self, session: AsyncSession, sub: Subscription, tariff, used: Optional[int]) -> None:
+        if not sub.period_start_at:
+            sub.period_start_at = datetime.now(timezone.utc)
+            return
+        now = datetime.now(timezone.utc)
+        next_reset = sub.period_start_at + timedelta(days=30)
+        if now < next_reset or sub.end_date <= now:
+            return
+
+        used_now = int(used or sub.traffic_used_bytes or 0)
+        baseline = int(sub.tier_baseline_bytes or tariff.monthly_bytes)
+        topup_used = max(0, used_now - baseline)
+        new_topup = max(0, int(sub.topup_balance_bytes or 0) - topup_used)
+        await self.panel_service.reset_user_traffic(sub.panel_user_uuid)
+        new_limit = baseline + new_topup
+        payload = self.subscription_service._build_panel_update_payload(
+            panel_user_uuid=sub.panel_user_uuid,
+            expire_at=sub.end_date,
+            status="ACTIVE",
+            traffic_limit_bytes=new_limit,
+        )
+        payload["activeInternalSquads"] = tariff.squad_uuids
+        await self.panel_service.update_user_details_on_panel(sub.panel_user_uuid, payload, log_response=False)
+        await subscription_dal.update_subscription(
+            session,
+            sub.subscription_id,
+            {
+                "period_start_at": next_reset,
+                "traffic_used_bytes": 0,
+                "traffic_limit_bytes": new_limit,
+                "topup_balance_bytes": new_topup,
+                "is_throttled": False,
+            },
+        )
+        await tariff_dal.clear_period_warnings(session, sub.subscription_id)
+
+    async def _maybe_warn_or_throttle(
+        self,
+        session: AsyncSession,
+        sub: Subscription,
+        tariff,
+        used: Optional[int],
+        limit: Optional[int],
+    ) -> None:
+        used_val = int(used or sub.traffic_used_bytes or 0)
+        limit_val = int(limit or sub.traffic_limit_bytes or 0)
+        if limit_val <= 0:
+            return
+        ratio = used_val / limit_val
+        for level, threshold in ((80, 0.8), (95, 0.95), (100, 1.0)):
+            if ratio < threshold:
+                continue
+            warning = await tariff_dal.get_warning(
+                session,
+                subscription_id=sub.subscription_id,
+                period_start_at=sub.period_start_at if tariff.billing_model == "period" else None,
+                level=level,
+                traffic_limit_bytes=limit_val if tariff.billing_model == "traffic" else None,
+            )
+            if warning:
+                continue
+            await tariff_dal.create_warning(
+                session,
+                subscription_id=sub.subscription_id,
+                period_start_at=sub.period_start_at if tariff.billing_model == "period" else None,
+                level=level,
+                traffic_limit_bytes=limit_val if tariff.billing_model == "traffic" else None,
+            )
+            if self.bot:
+                try:
+                    left_pct = max(0, 100 - level)
+                    if level < 100:
+                        text = f"Трафик тарифа {tariff.name(self.settings.DEFAULT_LANGUAGE)} почти закончился. Осталось около {left_pct}%."
+                    else:
+                        text = "Трафик закончился. Доступ временно ограничен до сброса или докупки пакета."
+                    await self.bot.send_message(sub.user_id, text)
+                except Exception:
+                    logging.exception("Failed to send traffic warning to user %s", sub.user_id)
+            if level == 100:
+                await self._throttle(session, sub, tariff)
+
+    async def _throttle(self, session: AsyncSession, sub: Subscription, tariff) -> None:
+        if sub.is_throttled:
+            return
+        for squad_uuid in tariff.squad_uuids:
+            await self.panel_service.remove_users_from_internal_squad(squad_uuid, [sub.panel_user_uuid])
+        await subscription_dal.update_subscription(
+            session,
+            sub.subscription_id,
+            {"is_throttled": True, "status_from_panel": "THROTTLED_BY_BOT"},
+        )
+
+    async def throttle_recovery_tick(self, session: AsyncSession) -> None:
+        result = await session.execute(
+            select(Subscription).where(
+                Subscription.is_active == True,
+                Subscription.is_throttled == True,
+            )
+        )
+        for sub in result.scalars().all():
+            try:
+                tariff = self.settings.tariffs_config.require(sub.tariff_key)
+            except Exception:
+                continue
+            if int(sub.traffic_limit_bytes or 0) <= int(sub.traffic_used_bytes or 0):
+                continue
+            for squad_uuid in tariff.squad_uuids:
+                await self.panel_service.add_users_to_internal_squad(squad_uuid, [sub.panel_user_uuid])
+            await subscription_dal.update_subscription(
+                session,
+                sub.subscription_id,
+                {"is_throttled": False, "status_from_panel": "ACTIVE"},
+            )