From e0ad31aa88916daaf7628d71a0fcd1f53b162386 Mon Sep 17 00:00:00 2001 From: Rerowros <85008083+Rerowros@users.noreply.github.com> Date: Wed, 9 Sep 2026 21:23:42 +0400 Subject: [PATCH 1/4] fix(users): preserve scheduled traffic cycles across manual resets --- .../workflows/test-database-migrations.yml | 6 + app/db/compiles_types.py | 2 +- app/db/crud/bulk.py | 24 +- app/db/crud/user.py | 224 ++++++++- ...8e4e47b9f2c_add_user_usage_reset_source.py | 40 ++ app/db/models.py | 70 ++- app/jobs/reset_user_data_usage.py | 2 + app/models/user.py | 3 + app/operation/subscription.py | 3 + app/operation/user.py | 7 + app/templates/subscription/index.html | 6 + dashboard/public/statics/locales/en.json | 3 + dashboard/public/statics/locales/fa.json | 3 + dashboard/public/statics/locales/ru.json | 3 + dashboard/public/statics/locales/zh.json | 3 + .../src/features/users/dialogs/user-modal.tsx | 45 ++ dashboard/src/service/api/index.ts | 6 + tests/api/conftest.py | 1 + tests/api/test_user.py | 47 +- tests/api/test_user_usage_reset_cycle.py | 455 ++++++++++++++++++ .../test_user_usage_reset_source_migration.py | 99 ++++ 21 files changed, 1018 insertions(+), 34 deletions(-) create mode 100644 app/db/migrations/versions/b8e4e47b9f2c_add_user_usage_reset_source.py create mode 100644 tests/api/test_user_usage_reset_cycle.py create mode 100644 tests/test_user_usage_reset_source_migration.py diff --git a/.github/workflows/test-database-migrations.yml b/.github/workflows/test-database-migrations.yml index acc8795bf..a3c93ebe5 100644 --- a/.github/workflows/test-database-migrations.yml +++ b/.github/workflows/test-database-migrations.yml @@ -194,6 +194,9 @@ jobs: - name: Check Migrations run: make check-migrations + - name: Verify traffic reset races under REPEATABLE READ + run: .venv/bin/python -m pytest tests/api/test_user_usage_reset_cycle.py -k mysql_repeatable_read -q + - name: Run Tests run: make test @@ -317,6 +320,9 @@ jobs: - name: Check Migrations run: make check-migrations + - name: Verify traffic reset races under REPEATABLE READ + run: .venv/bin/python -m pytest tests/api/test_user_usage_reset_cycle.py -k mysql_repeatable_read -q + - name: Run Tests (multi-worker all-in-one) env: ROLE: all-in-one diff --git a/app/db/compiles_types.py b/app/db/compiles_types.py index 973587bab..7ee0e7161 100644 --- a/app/db/compiles_types.py +++ b/app/db/compiles_types.py @@ -114,7 +114,7 @@ class DateDiff(FunctionElement): inherit_cache = True def __init__(self, date1, date2, **kwargs): - super().__init__(**kwargs) + super().__init__(date1, date2, **kwargs) self.date1 = date1 self.date2 = date2 diff --git a/app/db/crud/bulk.py b/app/db/crud/bulk.py index 8c55b15fa..dcb8620ad 100644 --- a/app/db/crud/bulk.py +++ b/app/db/crud/bulk.py @@ -13,6 +13,7 @@ User, UserStatus, UserUsageResetLogs, + UserUsageResetSource, users_groups_association, ) from app.models.group import BulkGroup @@ -46,8 +47,9 @@ async def reset_all_users_data_usage( Operations performed: - Sets `used_traffic` to 0 for all target users. - - Sets `status` to `active` for all users, unless filtered by admin. - - Deletes all related `UserUsageResetLogs` and `NextPlan` entries. + - Reactivates limited users while preserving all other statuses. + - Adds a manual `UserUsageResetLogs` row for lifetime traffic accounting. + - Deletes all related `NextPlan` entries. - Deletes `NodeUserUsage` chart rows when `clean_chart_data` is enabled. Args: @@ -60,16 +62,28 @@ async def reset_all_users_data_usage( - This function assumes proper foreign key constraints and cascading rules are in place. - The function commits changes at the end of the operation. """ - user_ids_query = select(User.id).where(User.admin_id == admin.id) if admin else select(User.id) - user_ids = (await db.execute(user_ids_query)).scalars().all() + users_query = select(User.id, User.used_traffic).order_by(User.id).with_for_update() + if admin: + users_query = users_query.where(User.admin_id == admin.id) + user_rows = (await db.execute(users_query)).all() + user_ids = [row.id for row in user_rows] if not user_ids: return + db.add_all( + [ + UserUsageResetLogs( + user_id=row.id, + used_traffic_at_reset=row.used_traffic, + reset_source=UserUsageResetSource.manual.value, + ) + for row in user_rows + ] + ) reset_status = case((User.status == UserStatus.limited, UserStatus.active), else_=User.status) await db.execute(update(User).where(User.id.in_(user_ids)).values(used_traffic=0, status=reset_status)) - await db.execute(delete(UserUsageResetLogs).where(UserUsageResetLogs.user_id.in_(user_ids))) if clean_chart_data: await db.execute(delete(NodeUserUsage).where(NodeUserUsage.user_id.in_(user_ids))) await db.execute(delete(NextPlan).where(NextPlan.user_id.in_(user_ids))) diff --git a/app/db/crud/user.py b/app/db/crud/user.py index 46e5def36..44472dad1 100644 --- a/app/db/crud/user.py +++ b/app/db/crud/user.py @@ -6,6 +6,7 @@ from sqlalchemy import and_, case, delete, desc, func, literal, not_, or_, select, update from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import joinedload, selectinload, with_expression +from sqlalchemy.orm.attributes import set_committed_value from sqlalchemy.sql import Select from sqlalchemy.sql.functions import coalesce @@ -22,6 +23,7 @@ UserStatus, UserSubscriptionUpdate, UserUsageResetLogs, + UserUsageResetSource, users_groups_association, ) from app.models.proxy import ProxyTable @@ -68,6 +70,12 @@ _USER_AGENT_MAX_LEN = UserSubscriptionUpdate.__table__.columns.user_agent.type.length or 512 _SUBSCRIPTION_UPDATE_IP_MAX_LEN = UserSubscriptionUpdate.__table__.columns.ip.type.length or 64 _ONLINE_USERS_WINDOW = timedelta(minutes=2) +_TRAFFIC_RESET_DAYS = { + DataLimitResetStrategy.day: 1, + DataLimitResetStrategy.week: 7, + DataLimitResetStrategy.month: 30, + DataLimitResetStrategy.year: 365, +} def _user_reset_traffic_subquery(): @@ -79,6 +87,29 @@ def _user_reset_traffic_subquery(): ) +def _user_last_reset_subquery(sources: Sequence[str] | None = None): + stmt = select(func.max(UserUsageResetLogs.reset_at)).where(UserUsageResetLogs.user_id == User.id) + if sources is not None: + stmt = stmt.where(UserUsageResetLogs.reset_source.in_(sources)) + return stmt.correlate(User).scalar_subquery() + + +def _user_usage_summary_options(): + # Keep optimized API reads independent of the size of reset history. + # A loaded SQL NULL is a valid result, not a request to lazy-load logs. + return ( + with_expression(User._reseted_usage_query, _user_reset_traffic_subquery()), + with_expression(User._last_traffic_reset_at_query, _user_last_reset_subquery()), + with_expression( + User._last_cycle_traffic_reset_at_query, + coalesce( + _user_last_reset_subquery([UserUsageResetSource.scheduled.value, UserUsageResetSource.next_plan.value]), + _user_last_reset_subquery([UserUsageResetSource.legacy.value]), + ), + ), + ) + + def _safe_on_hold_expire_duration(duration: int | None) -> int | None: if duration is None or duration <= 0: return None @@ -123,7 +154,7 @@ def _build_user_select_stmt( if options: stmt = stmt.options(*options) if load_lifetime_used_traffic: - stmt = stmt.options(with_expression(User._reseted_usage_query, _user_reset_traffic_subquery())) + stmt = stmt.options(*_user_usage_summary_options()) return stmt @@ -360,7 +391,7 @@ async def get_users( if load_usage_logs: options.append(selectinload(User.usage_logs)) if load_lifetime_used_traffic: - options.append(with_expression(User._reseted_usage_query, _user_reset_traffic_subquery())) + options.extend(_user_usage_summary_options()) stmt = select(User).options(*options) filters = [] @@ -609,42 +640,65 @@ async def get_on_hold_to_active_users(db: AsyncSession) -> list[User]: return list((await db.execute(stmt)).unique().scalars().all()) -async def get_users_to_reset_data_usage(db: AsyncSession) -> list[User]: +async def get_users_to_reset_data_usage( + db: AsyncSession, + *, + user_ids: Sequence[int] | None = None, +) -> list[User]: """ Retrieves users whose data usage needs to be reset based on their reset strategy. """ - last_reset_subq = ( + last_scheduled_reset_subq = ( select( UserUsageResetLogs.user_id, func.max(UserUsageResetLogs.reset_at).label("last_reset_at"), ) + .where( + UserUsageResetLogs.reset_source.in_( + [UserUsageResetSource.scheduled.value, UserUsageResetSource.next_plan.value] + ) + ) .group_by(UserUsageResetLogs.user_id) .subquery() ) - last_reset_time = coalesce(last_reset_subq.c.last_reset_at, User.created_at) + last_legacy_reset_subq = ( + select( + UserUsageResetLogs.user_id, + func.max(UserUsageResetLogs.reset_at).label("last_reset_at"), + ) + .where(UserUsageResetLogs.reset_source == UserUsageResetSource.legacy.value) + .group_by(UserUsageResetLogs.user_id) + .subquery() + ) - reset_strategy_to_days = { - DataLimitResetStrategy.day: 1, - DataLimitResetStrategy.week: 7, - DataLimitResetStrategy.month: 30, - DataLimitResetStrategy.year: 365, - } + # Preserve the pre-upgrade cycle once by using ambiguous legacy history only + # until the scheduler (or next-plan application) records a trusted boundary. + last_reset_time = coalesce( + last_scheduled_reset_subq.c.last_reset_at, + last_legacy_reset_subq.c.last_reset_at, + User.created_at, + ) num_days_to_reset_case = case( - *((User.data_limit_reset_strategy == strategy, days) for strategy, days in reset_strategy_to_days.items()), + *((User.data_limit_reset_strategy == strategy, days) for strategy, days in _TRAFFIC_RESET_DAYS.items()), else_=None, ) stmt = ( _build_user_select_stmt() - .outerjoin(last_reset_subq, User.id == last_reset_subq.c.user_id) + .outerjoin(last_scheduled_reset_subq, User.id == last_scheduled_reset_subq.c.user_id) + .outerjoin(last_legacy_reset_subq, User.id == last_legacy_reset_subq.c.user_id) .where( User.status.in_([UserStatus.active, UserStatus.limited]), User.data_limit_reset_strategy != DataLimitResetStrategy.no_reset, DateDiff(func.now(), last_reset_time) >= num_days_to_reset_case, ) ) + if user_ids is not None: + if not user_ids: + return [] + stmt = stmt.where(User.id.in_(user_ids)) return list((await db.execute(stmt)).unique().scalars().all()) @@ -1147,12 +1201,18 @@ async def modify_user( return db_user -async def _reset_user_traffic_and_log(db: AsyncSession, db_user: User): +async def _reset_user_traffic_and_log( + db: AsyncSession, + db_user: User, + *, + reset_source: UserUsageResetSource = UserUsageResetSource.manual, +): """Helper to reset user traffic and log the action.""" await db_user.awaitable_attrs.next_plan usage_log = UserUsageResetLogs( user_id=db_user.id, used_traffic_at_reset=db_user.used_traffic, + reset_source=reset_source.value, ) db.add(usage_log) @@ -1163,6 +1223,103 @@ async def _reset_user_traffic_and_log(db: AsyncSession, db_user: User): db_user.used_traffic = 0 +async def lock_users_for_traffic_reset(db: AsyncSession, users: Sequence[User]) -> list[User]: + """Lock reset targets in a stable order and refresh stale ORM state. + + Every path that records/reset traffic must acquire these row locks before + reading ``used_traffic``. Otherwise a scheduler candidate loaded before a + manual reset can account the same traffic twice after it resumes. + """ + user_ids = sorted({user.id for user in users if user.id is not None}) + if not user_ids: + return [] + + locked_ids = list( + (await db.execute(select(User.id).where(User.id.in_(user_ids)).order_by(User.id).with_for_update())).scalars() + ) + users_by_id = {user.id: user for user in users if user.id is not None} + locked_users = [users_by_id[user_id] for user_id in locked_ids] + for user in locked_users: + # Refresh only reset-relevant state. A full populate_existing refresh + # would expire preloaded groups and other relationships, which callers + # may still need without triggering implicit async IO. + await db.refresh( + user, + attribute_names=[ + "used_traffic", + "status", + "data_limit", + "data_limit_reset_strategy", + "_expire", + "on_hold_expire_duration", + "on_hold_timeout", + "proxy_settings", + ], + # MySQL REPEATABLE READ keeps an earlier plain-read snapshot even + # after the row lock is acquired. Refresh with a current read too. + with_for_update=True, + ) + next_plans = ( + await db.scalars( + select(NextPlan) + .where(NextPlan.user_id.in_(locked_ids)) + .order_by(NextPlan.user_id) + .with_for_update() + .execution_options(populate_existing=True) + ) + ).all() + plans_by_user = {plan.user_id: plan for plan in next_plans} + for user in locked_users: + set_committed_value(user, "next_plan", plans_by_user.get(user.id)) + return locked_users + + +async def _recheck_locked_users_due_for_reset(db: AsyncSession, users: Sequence[User]) -> list[User]: + if not users: + return [] + + cycle_starts = {} + for user in users: + # Read the actual latest row under lock, not MAX() from an old MySQL + # snapshot. LIMIT bounds history loading regardless of account age. + latest = ( + select(UserUsageResetLogs.reset_at) + .where(UserUsageResetLogs.user_id == user.id) + .order_by(UserUsageResetLogs.reset_at.desc(), UserUsageResetLogs.id.desc()) + .limit(1) + .with_for_update() + ) + cycle_start = await db.scalar( + latest.where( + UserUsageResetLogs.reset_source.in_( + [UserUsageResetSource.scheduled.value, UserUsageResetSource.next_plan.value] + ) + ) + ) + if cycle_start is None: + cycle_start = await db.scalar( + latest.where(UserUsageResetLogs.reset_source == UserUsageResetSource.legacy.value) + ) + cycle_starts[user.id] = cycle_start or user.created_at + + # Retain the existing database-specific DateDiff/day-boundary behavior. + # Current-read User state also handles a concurrent disable/strategy change. + due_ids = set( + await db.scalars( + select(User.id) + .where( + User.id.in_(cycle_starts), + User.status.in_([UserStatus.active, UserStatus.limited]), + User.data_limit_reset_strategy != DataLimitResetStrategy.no_reset, + DateDiff(func.now(), case(cycle_starts, value=User.id)) + >= case(_TRAFFIC_RESET_DAYS, value=User.data_limit_reset_strategy, else_=0), + ) + .with_for_update() + ) + ) + return [user for user in users if user.id in due_ids] + + async def clear_user_node_usages(db: AsyncSession, user_id: int, *, before: datetime | None = None) -> None: stmt = delete(NodeUserUsage).where(NodeUserUsage.user_id == user_id) if before is not None: @@ -1171,7 +1328,12 @@ async def clear_user_node_usages(db: AsyncSession, user_id: int, *, before: date async def reset_user_data_usage( - db: AsyncSession, db_user: User, *, clean_chart_data: bool = False, commit: bool = True + db: AsyncSession, + db_user: User, + *, + clean_chart_data: bool = False, + reset_source: UserUsageResetSource = UserUsageResetSource.manual, + commit: bool = True, ) -> User: """ Resets the data usage of a user and logs the reset. @@ -1183,7 +1345,12 @@ async def reset_user_data_usage( Returns: User: The updated user object. """ - await _reset_user_traffic_and_log(db, db_user) + locked_users = await lock_users_for_traffic_reset(db, [db_user]) + if not locked_users: + raise LookupError(f"User {db_user.id} no longer exists") + db_user = locked_users[0] + + await _reset_user_traffic_and_log(db, db_user, reset_source=reset_source) await delete_user_passed_notification_reminders(db, db_user.id, ReminderType.data_usage, 0) if clean_chart_data: await clear_user_node_usages(db, db_user.id) @@ -1198,7 +1365,12 @@ async def reset_user_data_usage( async def bulk_reset_user_data_usage( - db: AsyncSession, users: list[User], *, clean_chart_data: bool = False, commit: bool = True + db: AsyncSession, + users: list[User], + *, + clean_chart_data: bool = False, + reset_source: UserUsageResetSource = UserUsageResetSource.manual, + commit: bool = True, ) -> list[User]: """ Resets the data usage for a list of users and logs the reset. @@ -1210,8 +1382,15 @@ async def bulk_reset_user_data_usage( Returns: list[User]: The updated list of user objects. """ + users = await lock_users_for_traffic_reset(db, users) + if reset_source is UserUsageResetSource.scheduled: + # Candidates are discovered before row locks are acquired. A different + # worker may have completed the scheduled boundary while this worker + # waited, so evaluate eligibility again inside the locked transaction. + users = await _recheck_locked_users_due_for_reset(db, users) + for db_user in users: - await _reset_user_traffic_and_log(db, db_user) + await _reset_user_traffic_and_log(db, db_user, reset_source=reset_source) await delete_user_passed_notification_reminders(db, db_user.id, ReminderType.data_usage, 0) if clean_chart_data: await clear_user_node_usages(db, db_user.id) @@ -1244,6 +1423,13 @@ async def reset_user_by_next(db: AsyncSession, db_user: User, *, clean_chart_dat Returns: User: The updated user object. """ + locked_users = await lock_users_for_traffic_reset(db, [db_user]) + if not locked_users: + raise LookupError(f"User {db_user.id} no longer exists") + db_user = locked_users[0] + if db_user.next_plan is None: + raise LookupError(f"User {db_user.id} no longer has a next plan") + remaining_traffic = (db_user.data_limit or 0) - db_user.used_traffic if db_user.next_plan.user_template_id is None: db_user.data_limit = db_user.next_plan.data_limit + ( @@ -1282,7 +1468,7 @@ async def reset_user_by_next(db: AsyncSession, db_user: User, *, clean_chart_dat db_user.proxy_settings = proxy_settings db_user.data_limit_reset_strategy = db_user.next_plan.user_template.data_limit_reset_strategy - await _reset_user_traffic_and_log(db, db_user) + await _reset_user_traffic_and_log(db, db_user, reset_source=UserUsageResetSource.next_plan) await delete_user_passed_notification_reminders(db, db_user.id, ReminderType.data_usage, 0) if clean_chart_data: await clear_user_node_usages(db, db_user.id) diff --git a/app/db/migrations/versions/b8e4e47b9f2c_add_user_usage_reset_source.py b/app/db/migrations/versions/b8e4e47b9f2c_add_user_usage_reset_source.py new file mode 100644 index 000000000..1ca71af3e --- /dev/null +++ b/app/db/migrations/versions/b8e4e47b9f2c_add_user_usage_reset_source.py @@ -0,0 +1,40 @@ +"""add user usage reset source + +Revision ID: b8e4e47b9f2c +Revises: 8e2f1a9c4b70 +Create Date: 2026-05-10 00:00:00.000000 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "b8e4e47b9f2c" +down_revision = "8e2f1a9c4b70" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.add_column( + "user_usage_logs", + sa.Column("reset_source", sa.String(length=32), nullable=False, server_default="legacy"), + ) + op.create_index( + "ix_user_usage_logs_user_id_source_reset_at", + "user_usage_logs", + ["user_id", "reset_source", "reset_at"], + ) + # Historical rows and rows written by an older process during a rolling + # upgrade may represent manual, scheduled, or next-plan resets. Do not guess: treating them as + # scheduled would let an old manual reset postpone the next cycle. + # ``legacy`` keeps those rows distinguishable so the scheduler can use the + # latest one only as a transitional anchor until a trusted reset exists. + # The server default backfills existing rows and protects old writers. + # Do not overwrite explicit sources from a new writer during the upgrade. + + +def downgrade() -> None: + op.drop_index("ix_user_usage_logs_user_id_source_reset_at", table_name="user_usage_logs") + op.drop_column("user_usage_logs", "reset_source") diff --git a/app/db/models.py b/app/db/models.py index 30ed85024..c86e6a102 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -1,5 +1,5 @@ import os -from datetime import UTC, datetime as dt +from datetime import UTC, datetime as dt, timedelta as td from enum import Enum from typing import Any @@ -190,6 +190,13 @@ class DataLimitResetStrategy(str, Enum): year = "year" +class UserUsageResetSource(str, Enum): + legacy = "legacy" + manual = "manual" + scheduled = "scheduled" + next_plan = "next_plan" + + class User(Base, CreatedAtUTCMixin): __tablename__ = "users" __table_args__ = ( @@ -209,7 +216,11 @@ class User(Base, CreatedAtUTCMixin): subscription_updates: Mapped[list[UserSubscriptionUpdate]] = relationship( back_populates="user", cascade="all, delete-orphan", init=False ) - usage_logs: Mapped[list[UserUsageResetLogs]] = relationship(back_populates="user", init=False) + usage_logs: Mapped[list[UserUsageResetLogs]] = relationship( + back_populates="user", + init=False, + order_by="UserUsageResetLogs.reset_at", + ) admin: Mapped[Admin] = relationship(back_populates="users", init=False) next_plan: Mapped[NextPlan | None] = relationship( uselist=False, back_populates="user", cascade="all, delete-orphan", init=False @@ -238,6 +249,8 @@ class User(Base, CreatedAtUTCMixin): edit_at: Mapped[dt | None] = mapped_column(DateTime(timezone=True), default=None) last_status_change: Mapped[dt | None] = mapped_column(DateTime(timezone=True), default=None) _reseted_usage_query: Mapped[int | None] = query_expression(repr=False) + _last_traffic_reset_at_query: Mapped[dt | None] = query_expression(repr=False) + _last_cycle_traffic_reset_at_query: Mapped[dt | None] = query_expression(repr=False) @hybrid_property def expire(self) -> dt | None: @@ -281,6 +294,48 @@ def lifetime_used_traffic(self) -> int: def last_traffic_reset_time(self): return self.usage_logs[-1].reset_at if self.usage_logs else self.created_at + @property + def last_traffic_reset_at(self): + if "_last_traffic_reset_at_query" in self.__dict__: + return self._last_traffic_reset_at_query + return self.usage_logs[-1].reset_at if self.usage_logs else None + + @property + def last_scheduled_traffic_reset_at(self): + scheduled_sources = {UserUsageResetSource.scheduled.value, UserUsageResetSource.next_plan.value} + for log in reversed(self.usage_logs): + if log.reset_source in scheduled_sources: + return log.reset_at + return None + + @property + def last_cycle_traffic_reset_at(self): + if "_last_cycle_traffic_reset_at_query" in self.__dict__: + return self._last_cycle_traffic_reset_at_query + if scheduled_reset_at := self.last_scheduled_traffic_reset_at: + return scheduled_reset_at + for log in reversed(self.usage_logs): + if log.reset_source == UserUsageResetSource.legacy.value: + return log.reset_at + return None + + @property + def next_traffic_reset_at(self): + reset_days = { + DataLimitResetStrategy.day: 1, + DataLimitResetStrategy.week: 7, + DataLimitResetStrategy.month: 30, + DataLimitResetStrategy.year: 365, + } + days = reset_days.get(self.data_limit_reset_strategy) + if days is None: + return None + + cycle_start = self.last_cycle_traffic_reset_at or self.created_at + if cycle_start.tzinfo is None: + cycle_start = cycle_start.replace(tzinfo=UTC) + return cycle_start + td(days=days) + async def inbounds(self) -> list[str]: """Returns a flat list of all included inbound tags for enabled groups.""" session = async_object_session(self) @@ -472,10 +527,21 @@ class UserUsageResetLogs(Base, IdMixin): __table_args__ = ( # Index for user-specific queries sorted by time Index("ix_user_usage_logs_user_id_reset_at", "user_id", "reset_at"), + # The reset scheduler groups the per-user history after filtering by + # source, so the original (user_id, reset_at) index cannot cover it. + Index("ix_user_usage_logs_user_id_source_reset_at", "user_id", "reset_source", "reset_at"), ) user_id: Mapped[int | None] = fk_id_column("users.id", ondelete="CASCADE", nullable=True) user: Mapped[User] = relationship(back_populates="usage_logs", init=False) used_traffic_at_reset: Mapped[int] = mapped_column(BigInteger, nullable=False) + reset_source: Mapped[str] = mapped_column( + String(32), + default=UserUsageResetSource.manual.value, + # Old application versions omit this column during a rolling upgrade. + # Mark those ambiguous writes as legacy; new code always sends the + # explicit Python default above. + server_default=UserUsageResetSource.legacy.value, + ) reset_at: Mapped[dt] = mapped_column(DateTime(timezone=True), default=lambda: dt.now(UTC), init=False) diff --git a/app/jobs/reset_user_data_usage.py b/app/jobs/reset_user_data_usage.py index 28e1998c2..d986e33db 100644 --- a/app/jobs/reset_user_data_usage.py +++ b/app/jobs/reset_user_data_usage.py @@ -4,6 +4,7 @@ from app import notification, scheduler from app.db import GetDB from app.db.crud.user import bulk_reset_user_data_usage, get_users_to_reset_data_usage +from app.db.models import UserUsageResetSource from app.jobs.dependencies import SYSTEM_ADMIN from app.operation import OperatorType from app.operation.user import UserOperation @@ -23,6 +24,7 @@ async def reset_data_usage(): db, users, clean_chart_data=usage_settings.reset_user_usage_clean_chart_data, + reset_source=UserUsageResetSource.scheduled, ) for db_user in updated_users: diff --git a/app/models/user.py b/app/models/user.py index 7af5d0756..0bf383bd3 100644 --- a/app/models/user.py +++ b/app/models/user.py @@ -124,6 +124,9 @@ class UserNotificationResponse(User): created_at: AwareDatetime edit_at: OptionalAwareDatetime = Field(default=None) online_at: OptionalAwareDatetime = Field(default=None) + last_traffic_reset_at: OptionalAwareDatetime = Field(default=None) + last_cycle_traffic_reset_at: OptionalAwareDatetime = Field(default=None) + next_traffic_reset_at: OptionalAwareDatetime = Field(default=None) subscription_url: str = Field(default="") admin: AdminContactInfo | None = Field(default=None) group_names: list[str] | None = Field(default_factory=list) diff --git a/app/operation/subscription.py b/app/operation/subscription.py index 4d345dce2..33e563db4 100644 --- a/app/operation/subscription.py +++ b/app/operation/subscription.py @@ -94,6 +94,9 @@ async def validated_user(db_user: User) -> UsersResponseWithInbounds: user.inbounds = await db_user.inbounds() user.expire = db_user.expire user.lifetime_used_traffic = db_user.lifetime_used_traffic + user.last_traffic_reset_at = db_user.last_traffic_reset_at + user.last_cycle_traffic_reset_at = db_user.last_cycle_traffic_reset_at + user.next_traffic_reset_at = db_user.next_traffic_reset_at return user diff --git a/app/operation/user.py b/app/operation/user.py index b79a888e3..1e57787b4 100644 --- a/app/operation/user.py +++ b/app/operation/user.py @@ -44,6 +44,7 @@ get_users_subscription_agent_stats, load_user_attrs, lock_admin_quota_row, + lock_users_for_traffic_reset, modify_user as crud_modify_user, remove_expired_users, remove_user, @@ -1857,6 +1858,12 @@ async def bulk_apply_template_to_users( ) prepared_updates.append((db_user, modify_user, validated_groups, original_status, emit_reset_status_change)) + if user_template.reset_usages: + # Acquire the whole batch in stable id order before per-user work; + # otherwise two overlapping template batches can deadlock while + # locking the same users in opposite request order. + await lock_users_for_traffic_reset(db, [item[0] for item in prepared_updates]) + modified_user_ids: list[int] = [] try: for db_user, modified_user_model, validated_groups, _, _ in prepared_updates: diff --git a/app/templates/subscription/index.html b/app/templates/subscription/index.html index b595019c8..6f32a02f5 100644 --- a/app/templates/subscription/index.html +++ b/app/templates/subscription/index.html @@ -577,6 +577,12 @@

Subscription Information

Data Used: {{ user.used_traffic | bytesformat }}{% if user.data_limit_reset_strategy != 'no_reset' %} (resets every {{ user.data_limit_reset_strategy.value }}){% endif %} + {% if user.next_traffic_reset_at %} +
+ Next Traffic Reset: + {{ user.next_traffic_reset_at | datetime }} +
+ {% endif %}
{% if user.status.value == 'on_hold' %} On Hold For: diff --git a/dashboard/public/statics/locales/en.json b/dashboard/public/statics/locales/en.json index 9760cd86b..d4a316e03 100644 --- a/dashboard/public/statics/locales/en.json +++ b/dashboard/public/statics/locales/en.json @@ -1,4 +1,7 @@ { + "lastTrafficResetAt": "Last traffic reset", + "lastCycleTrafficResetAt": "Last traffic cycle reset", + "nextTrafficResetAt": "Next traffic reset", "pasarguard": "PasarGuard", "dashboard": "Dashboard", "master": "Overall", diff --git a/dashboard/public/statics/locales/fa.json b/dashboard/public/statics/locales/fa.json index c0b3d2e9d..7ff50a9c2 100644 --- a/dashboard/public/statics/locales/fa.json +++ b/dashboard/public/statics/locales/fa.json @@ -1,4 +1,7 @@ { + "lastTrafficResetAt": "آخرین بازنشانی ترافیک", + "lastCycleTrafficResetAt": "آخرین بازنشانی چرخه ترافیک", + "nextTrafficResetAt": "بازنشانی بعدی ترافیک", "pasarguard": "پاسارگارد", "dashboard": "داشبورد", "master": "کلی", diff --git a/dashboard/public/statics/locales/ru.json b/dashboard/public/statics/locales/ru.json index 6c89c39bb..459cdb448 100644 --- a/dashboard/public/statics/locales/ru.json +++ b/dashboard/public/statics/locales/ru.json @@ -1,4 +1,7 @@ { + "lastTrafficResetAt": "Последний сброс трафика", + "lastCycleTrafficResetAt": "Последний сброс цикла трафика", + "nextTrafficResetAt": "Следующий сброс трафика", "pasarguard": "PasarGuard", "dashboard": "Панель управления", "master": "Общее", diff --git a/dashboard/public/statics/locales/zh.json b/dashboard/public/statics/locales/zh.json index 1a602679a..fd1523663 100644 --- a/dashboard/public/statics/locales/zh.json +++ b/dashboard/public/statics/locales/zh.json @@ -1,4 +1,7 @@ { + "lastTrafficResetAt": "上次流量重置", + "lastCycleTrafficResetAt": "上次流量周期重置", + "nextTrafficResetAt": "下次流量重置", "active": "激活", "platform": "平台", "dashboard": "仪表板", diff --git a/dashboard/src/features/users/dialogs/user-modal.tsx b/dashboard/src/features/users/dialogs/user-modal.tsx index cab52c3b7..422143328 100644 --- a/dashboard/src/features/users/dialogs/user-modal.tsx +++ b/dashboard/src/features/users/dialogs/user-modal.tsx @@ -1364,6 +1364,18 @@ function UserModal({ isDialogOpen, onOpenChange, form, editingUser, editingUserI return formatMetaDate(editingUserData?.edit_at) }, [editingUserData?.edit_at, formatMetaDate]) + const lastTrafficResetAtText = React.useMemo(() => { + return formatMetaDate(editingUserData?.last_traffic_reset_at) + }, [editingUserData?.last_traffic_reset_at, formatMetaDate]) + + const lastCycleTrafficResetAtText = React.useMemo(() => { + return formatMetaDate(editingUserData?.last_cycle_traffic_reset_at) + }, [editingUserData?.last_cycle_traffic_reset_at, formatMetaDate]) + + const nextTrafficResetAtText = React.useMemo(() => { + return formatMetaDate(editingUserData?.next_traffic_reset_at) + }, [editingUserData?.next_traffic_reset_at, formatMetaDate]) + const confirmResetUsage = async () => { if (!currentUserId || !currentUsername) return try { @@ -1421,6 +1433,39 @@ function UserModal({ isDialogOpen, onOpenChange, form, editingUser, editingUserI
)} + {lastTrafficResetAtText && ( +
+ + + {t('lastTrafficResetAt', { defaultValue: 'Last traffic reset' })} + + + {lastTrafficResetAtText} + +
+ )} + {lastCycleTrafficResetAtText && ( +
+ + + {t('lastCycleTrafficResetAt', { defaultValue: 'Last traffic cycle reset' })} + + + {lastCycleTrafficResetAtText} + +
+ )} + {nextTrafficResetAtText && ( +
+ + + {t('nextTrafficResetAt', { defaultValue: 'Next traffic reset' })} + + + {nextTrafficResetAtText} + +
+ )} diff --git a/dashboard/src/service/api/index.ts b/dashboard/src/service/api/index.ts index 6d7e8accf..83892bdab 100644 --- a/dashboard/src/service/api/index.ts +++ b/dashboard/src/service/api/index.ts @@ -2174,6 +2174,9 @@ export interface SubscriptionUserResponse { created_at: string; edit_at?: string | null; online_at?: string | null; + last_traffic_reset_at?: string | null; + last_cycle_traffic_reset_at?: string | null; + next_traffic_reset_at?: string | null; ip?: string | null; } @@ -2352,6 +2355,9 @@ export interface UserResponse { created_at: string; edit_at?: string | null; online_at?: string | null; + last_traffic_reset_at?: string | null; + last_cycle_traffic_reset_at?: string | null; + next_traffic_reset_at?: string | null; subscription_url?: string; admin?: AdminBase | null; } diff --git a/tests/api/conftest.py b/tests/api/conftest.py index 16a406b90..a18f07727 100644 --- a/tests/api/conftest.py +++ b/tests/api/conftest.py @@ -25,6 +25,7 @@ def mock_db_session(monkeypatch: pytest.MonkeyPatch): db_session = MagicMock(spec=TestSession) monkeypatch.setattr("app.settings.GetDB", db_session) monkeypatch.setattr("app.subscription.client_templates.GetDB", GetTestDB) + monkeypatch.setattr("app.utils.jwt.GetDB", GetTestDB) return db_session diff --git a/tests/api/test_user.py b/tests/api/test_user.py index 694dbf297..f5b7ea791 100644 --- a/tests/api/test_user.py +++ b/tests/api/test_user.py @@ -15,9 +15,7 @@ from sqlalchemy import delete, event, func, select, update from app.db.crud.hwid import register_user_hwid -from app.db.crud.user import get_user as get_db_user -from app.db.crud.user import get_users as get_db_users -from app.db.crud.user import update_users_status +from app.db.crud.user import get_user as get_db_user, get_users as get_db_users, update_users_status from app.db.models import NodeUserUsage, User, UserStatus, UserUsageResetLogs from app.models.settings import ConfigFormat, SubRule, Subscription from app.models.stats import Period, UserCountMetric, UserCountMetricStat, UserCountMetricStatsList @@ -967,13 +965,24 @@ def test_user_subscriptions(access_token): user = create_user( access_token, group_ids=[group["id"] for group in groups], - payload={"username": unique_name("test_user_subscriptions")}, + payload={ + "username": unique_name("test_user_subscriptions"), + "data_limit_reset_strategy": "month", + }, ) try: + reset_response = client.post( + f"/api/user/by-id/{user['id']}/reset", + headers={"Authorization": f"Bearer {access_token}"}, + ) + assert reset_response.status_code == status.HTTP_200_OK + for usf in user_subscription_formats: url = f"{user['subscription_url']}/{usf}" response = client.get(url, headers={"Accept": "text/html"} if usf == "" else None) assert response.status_code == status.HTTP_200_OK + if usf == "": + assert "Next Traffic Reset:" in response.text finally: delete_user(access_token, user["username"]) for host in hosts: @@ -1190,13 +1199,29 @@ def test_user_subscription_info_returns_request_ip(access_token): user = create_user( access_token, group_ids=[groups[0]["id"]], - payload={"username": unique_name("test_subscription_info_ip")}, + payload={ + "username": unique_name("test_subscription_info_ip"), + "data_limit_reset_strategy": "month", + }, ) try: + reset_response = client.post( + f"/api/user/by-id/{user['id']}/reset", + headers={"Authorization": f"Bearer {access_token}"}, + ) + assert reset_response.status_code == status.HTTP_200_OK + ip = "198.51.100.7" response = client.get(f"{user['subscription_url']}/info", headers={"X-Forwarded-For": ip}) assert response.status_code == status.HTTP_200_OK - assert response.json()["ip"] == ip + data = response.json() + assert data["ip"] == ip + assert data["last_traffic_reset_at"] is not None + # A manual reset updates the general timestamp only. The separate + # cycle timestamp must stay empty until the scheduler/next-plan path + # performs a reset, otherwise the UI would again conflate both clocks. + assert data["last_cycle_traffic_reset_at"] is None + assert data["next_traffic_reset_at"] is not None finally: delete_user(access_token, user["username"]) cleanup_groups(access_token, core, groups) @@ -2032,7 +2057,7 @@ def test_reset_by_next_user_usage(access_token): user = create_user( access_token, group_ids=[groups[0]["id"]], - payload={"username": unique_name("test_user_next_plan")}, + payload={"username": unique_name("test_user_next_plan"), "data_limit_reset_strategy": "month"}, ) try: update = client.put( @@ -2046,6 +2071,14 @@ def test_reset_by_next_user_usage(access_token): headers={"Authorization": f"Bearer {access_token}"}, ) assert response.status_code == status.HTTP_200_OK + data = response.json() + assert data["data_limit_reset_strategy"] == "month" + assert data["last_traffic_reset_at"] == data["last_cycle_traffic_reset_at"] + assert data["last_cycle_traffic_reset_at"] is not None + info = client.get(f"{user['subscription_url']}/info") + assert info.status_code == status.HTTP_200_OK + assert info.json()["last_cycle_traffic_reset_at"] == data["last_cycle_traffic_reset_at"] + assert info.json()["next_traffic_reset_at"] == data["next_traffic_reset_at"] finally: delete_user(access_token, user["username"]) cleanup_groups(access_token, core, groups) diff --git a/tests/api/test_user_usage_reset_cycle.py b/tests/api/test_user_usage_reset_cycle.py new file mode 100644 index 000000000..a70304134 --- /dev/null +++ b/tests/api/test_user_usage_reset_cycle.py @@ -0,0 +1,455 @@ +import asyncio +from datetime import UTC, datetime, timedelta +from uuid import uuid4 + +import pytest +from sqlalchemy import delete, event, func, select +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from app.db.crud.bulk import reset_all_users_data_usage +from app.db.crud.user import ( + bulk_reset_user_data_usage, + get_user, + get_users_to_reset_data_usage, + reset_user_by_next, + reset_user_data_usage, +) +from app.db.models import ( + Admin, + DataLimitResetStrategy, + NextPlan, + User, + UserStatus, + UserTemplate, + UserUsageResetLogs, + UserUsageResetSource, +) +from app.models.proxy import ProxyTable +from tests.api import DATABASE_URL, TestSession + + +@pytest.mark.asyncio +async def test_legacy_reset_is_only_used_until_a_scheduled_reset_exists(): + now = datetime.now(UTC) + + async with TestSession() as session: + + def make_user(label: str) -> User: + user = User( + username=f"{label}_reset_cycle_{uuid4().hex[:8]}", + data_limit_reset_strategy=DataLimitResetStrategy.month, + proxy_settings=ProxyTable().dict(no_obj=True), + ) + user.created_at = now - timedelta(days=61) + return user + + manual_reset_user = make_user("manual") + recent_scheduled_user = make_user("recent_scheduled") + recent_legacy_user = make_user("recent_legacy") + old_legacy_user = make_user("old_legacy") + scheduled_with_later_noise_user = make_user("scheduled_with_later_noise") + + session.add_all( + [ + manual_reset_user, + recent_scheduled_user, + recent_legacy_user, + old_legacy_user, + scheduled_with_later_noise_user, + ] + ) + await session.flush() + await session.execute( + delete(UserUsageResetLogs).where( + UserUsageResetLogs.user_id.in_( + [ + manual_reset_user.id, + recent_scheduled_user.id, + recent_legacy_user.id, + old_legacy_user.id, + scheduled_with_later_noise_user.id, + ] + ) + ) + ) + + def make_log(user: User, source: UserUsageResetSource, days_ago: int) -> UserUsageResetLogs: + log = UserUsageResetLogs( + user_id=user.id, + used_traffic_at_reset=1024, + reset_source=source.value, + ) + log.reset_at = now - timedelta(days=days_ago) + return log + + session.add_all( + [ + make_log(manual_reset_user, UserUsageResetSource.manual, 1), + make_log(recent_scheduled_user, UserUsageResetSource.scheduled, 1), + make_log(recent_legacy_user, UserUsageResetSource.legacy, 1), + make_log(old_legacy_user, UserUsageResetSource.legacy, 31), + make_log(scheduled_with_later_noise_user, UserUsageResetSource.scheduled, 31), + make_log(scheduled_with_later_noise_user, UserUsageResetSource.legacy, 1), + make_log(scheduled_with_later_noise_user, UserUsageResetSource.manual, 1), + ] + ) + await session.commit() + + users_to_reset = await get_users_to_reset_data_usage(session) + user_ids_to_reset = {user.id for user in users_to_reset} + + assert manual_reset_user.id in user_ids_to_reset + assert recent_scheduled_user.id not in user_ids_to_reset + assert recent_legacy_user.id not in user_ids_to_reset + assert old_legacy_user.id in user_ids_to_reset + assert scheduled_with_later_noise_user.id in user_ids_to_reset + + await session.refresh(recent_legacy_user, attribute_names=["usage_logs"]) + await session.refresh(scheduled_with_later_noise_user, attribute_names=["usage_logs"]) + assert recent_legacy_user.next_traffic_reset_at.date() == (now + timedelta(days=29)).date() + assert scheduled_with_later_noise_user.next_traffic_reset_at.date() == (now - timedelta(days=1)).date() + + +@pytest.mark.asyncio +async def test_stale_scheduler_user_does_not_double_count_manual_reset(): + username = f"reset_race_{uuid4().hex[:8]}" + + async with TestSession() as scheduler_session: + user = User( + username=username, + used_traffic=100, + data_limit_reset_strategy=DataLimitResetStrategy.day, + proxy_settings=ProxyTable().dict(no_obj=True), + ) + user.created_at = datetime.now(UTC) - timedelta(days=2) + scheduler_session.add(user) + await scheduler_session.commit() + stale_user = await scheduler_session.get(User, user.id) + await scheduler_session.commit() + assert stale_user.used_traffic == 100 + + async with TestSession() as manual_session: + manual_user = await manual_session.get(User, user.id) + await reset_user_data_usage(manual_session, manual_user) + + # The scheduler still holds the pre-reset ORM object. The row-locking + # reset must refresh it before lifetime traffic is logged. + await reset_user_data_usage( + scheduler_session, + stale_user, + reset_source=UserUsageResetSource.scheduled, + ) + + total_reset_traffic = await scheduler_session.scalar( + select(func.sum(UserUsageResetLogs.used_traffic_at_reset)).where(UserUsageResetLogs.user_id == user.id) + ) + assert total_reset_traffic == 100 + + +@pytest.mark.asyncio +async def test_scheduler_rechecks_due_state_after_locking_stale_candidates(): + username = f"scheduler_recheck_{uuid4().hex[:8]}" + + async with TestSession() as stale_scheduler_session: + user = User( + username=username, + used_traffic=100, + data_limit_reset_strategy=DataLimitResetStrategy.day, + proxy_settings=ProxyTable().dict(no_obj=True), + ) + user.created_at = datetime.now(UTC) - timedelta(days=2) + stale_scheduler_session.add(user) + await stale_scheduler_session.commit() + + stale_candidates = await get_users_to_reset_data_usage(stale_scheduler_session, user_ids=[user.id]) + assert [candidate.id for candidate in stale_candidates] == [user.id] + await stale_scheduler_session.commit() + + async with TestSession() as winning_scheduler_session: + winning_user = await winning_scheduler_session.get(User, user.id) + reset_users = await bulk_reset_user_data_usage( + winning_scheduler_session, + [winning_user], + reset_source=UserUsageResetSource.scheduled, + ) + assert [reset_user.id for reset_user in reset_users] == [user.id] + + reset_users = await bulk_reset_user_data_usage( + stale_scheduler_session, + stale_candidates, + reset_source=UserUsageResetSource.scheduled, + ) + assert reset_users == [] + + scheduled_log_count = await stale_scheduler_session.scalar( + select(func.count(UserUsageResetLogs.id)).where( + UserUsageResetLogs.user_id == user.id, + UserUsageResetLogs.reset_source == UserUsageResetSource.scheduled.value, + ) + ) + assert scheduled_log_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "strategy,days", + [ + (DataLimitResetStrategy.day, 1), + (DataLimitResetStrategy.week, 7), + (DataLimitResetStrategy.month, 30), + (DataLimitResetStrategy.year, 365), + ], +) +@pytest.mark.parametrize("bulk", [False, True]) +async def test_manual_reset_preserves_each_cycle_and_scheduled_reset_advances_it(strategy, days, bulk): + async with TestSession() as session: + user = User(username=f"cycle_{uuid4().hex[:16]}", used_traffic=123, data_limit_reset_strategy=strategy) + user.created_at = datetime.now(UTC) - timedelta(days=days + 2) + session.add(user) + await session.commit() + if bulk: + await bulk_reset_user_data_usage(session, [user]) + else: + await reset_user_data_usage(session, user) + assert user.used_traffic == 0 + assert user.lifetime_used_traffic == 123 + assert user.last_cycle_traffic_reset_at is None + assert user.last_traffic_reset_at is not None + assert user.next_traffic_reset_at.date() == (user.created_at + timedelta(days=days)).date() + assert [u.id for u in await get_users_to_reset_data_usage(session, user_ids=[user.id])] == [user.id] + reset = await bulk_reset_user_data_usage(session, [user], reset_source=UserUsageResetSource.scheduled) + assert [u.id for u in reset] == [user.id] + assert user.usage_logs[-1].reset_source == "scheduled" + assert user.last_cycle_traffic_reset_at == user.last_traffic_reset_at + assert user.next_traffic_reset_at.date() == (datetime.now(UTC) + timedelta(days=days)).date() + assert user.lifetime_used_traffic == 123 + assert await get_users_to_reset_data_usage(session, user_ids=[user.id]) == [] + + +@pytest.mark.asyncio +async def test_scoped_reset_all_retains_history_lifetime_and_status_without_advancing_cycle(): + async with TestSession() as session: + admin = Admin(username=f"reset_{uuid4().hex[:12]}", hashed_password="unused", role_id=3) + session.add(admin) + await session.flush() + user = User( + username=f"all_{uuid4().hex[:16]}", + admin_id=admin.id, + used_traffic=200, + status=UserStatus.disabled, + data_limit_reset_strategy=DataLimitResetStrategy.week, + ) + user.created_at = datetime.now(UTC) - timedelta(days=10) + outside = User(username=f"outside_{uuid4().hex[:16]}", used_traffic=55) + session.add_all([user, outside]) + await session.flush() + prior = UserUsageResetLogs(user_id=user.id, used_traffic_at_reset=100, reset_source="scheduled") + prior.reset_at = datetime.now(UTC) - timedelta(days=8) + session.add(prior) + session.add(NextPlan(user_id=user.id, user_template_id=None, data_limit=900)) + await session.commit() + await reset_all_users_data_usage(session, admin) + await session.refresh(user) + await user.awaitable_attrs.usage_logs + await user.awaitable_attrs.next_plan + await session.refresh(outside) + assert user.status == UserStatus.disabled + assert user.used_traffic == 0 + assert user.lifetime_used_traffic == 300 + assert [(log.used_traffic_at_reset, log.reset_source) for log in user.usage_logs] == [ + (100, "scheduled"), + (200, "manual"), + ] + assert user.last_cycle_traffic_reset_at == prior.reset_at + assert user.next_plan is None # Existing manual reset policy remains unchanged. + assert outside.used_traffic == 55 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("with_template", [False, True]) +async def test_next_plan_keeps_existing_policy_and_starts_a_new_cycle(with_template): + async with TestSession() as session: + user = User( + username=f"next_{uuid4().hex[:16]}", + data_limit=100, + used_traffic=10, + data_limit_reset_strategy=DataLimitResetStrategy.month, + proxy_settings=ProxyTable().dict(no_obj=True), + ) + user.created_at = datetime.now(UTC) - timedelta(days=40) + session.add(user) + await session.flush() + template_id = None + if with_template: + template = UserTemplate( + name=f"next_tpl_{uuid4().hex[:12]}", + username_prefix=None, + username_suffix=None, + extra_settings=None, + groups=[], + data_limit=1000, + data_limit_reset_strategy=DataLimitResetStrategy.week, + ) + session.add(template) + await session.flush() + template_id = template.id + session.add( + NextPlan(user_id=user.id, user_template_id=template_id, data_limit=1000, add_remaining_traffic=True) + ) + await session.commit() + user = await get_user(session, user.username) + await reset_user_by_next(session, user) + assert user.next_plan is None + assert user.used_traffic == 0 + assert user.data_limit == 1090 + assert user.lifetime_used_traffic == 10 + assert user.usage_logs[-1].reset_source == "next_plan" + assert user.last_cycle_traffic_reset_at == user.last_traffic_reset_at + assert user.data_limit_reset_strategy == ( + DataLimitResetStrategy.week if with_template else DataLimitResetStrategy.month + ) + assert ( + user.next_traffic_reset_at.date() == (datetime.now(UTC) + timedelta(days=7 if with_template else 30)).date() + ) + assert await get_users_to_reset_data_usage(session, user_ids=[user.id]) == [] + + +@pytest.mark.asyncio +async def test_optimized_reset_dates_match_history_and_refresh_after_manual_reset(): + async with TestSession() as session: + user = User(username=f"dates_{uuid4().hex[:16]}", data_limit_reset_strategy=DataLimitResetStrategy.month) + session.add(user) + await session.flush() + scheduled = UserUsageResetLogs(user_id=user.id, used_traffic_at_reset=10, reset_source="scheduled") + scheduled.reset_at = datetime.now(UTC) - timedelta(days=5) + session.add(scheduled) + await session.commit() + user = await get_user(session, user.username, load_usage_logs=False, load_lifetime_used_traffic=True) + assert user.last_cycle_traffic_reset_at == user.last_traffic_reset_at + assert user.last_cycle_traffic_reset_at.replace(tzinfo=UTC) == scheduled.reset_at.replace(tzinfo=UTC) + await reset_user_data_usage(session, user) + assert user.last_traffic_reset_at.replace(tzinfo=UTC) > user.last_cycle_traffic_reset_at.replace(tzinfo=UTC) + assert user.last_cycle_traffic_reset_at.replace(tzinfo=UTC) == scheduled.reset_at.replace(tzinfo=UTC) + + +@pytest.mark.asyncio +async def test_no_reset_strategy_has_no_next_date_and_is_never_scheduled(): + async with TestSession() as session: + user = User(username=f"no_cycle_{uuid4().hex[:16]}", used_traffic=50) + user.created_at = datetime.now(UTC) - timedelta(days=400) + session.add(user) + await session.commit() + await reset_user_data_usage(session, user) + assert user.next_traffic_reset_at is None + assert await get_users_to_reset_data_usage(session, user_ids=[user.id]) == [] + + +@pytest.mark.asyncio +@pytest.mark.skipif(not DATABASE_URL.startswith("mysql"), reason="requires MySQL/MariaDB REPEATABLE READ") +@pytest.mark.parametrize("winner_source", [UserUsageResetSource.manual, UserUsageResetSource.scheduled]) +async def test_mysql_repeatable_read_scheduler_waits_for_current_reset(winner_source): + engine = create_async_engine(DATABASE_URL, isolation_level="REPEATABLE READ") + sessions = async_sessionmaker(engine, expire_on_commit=False, autoflush=False) + user_id = None + try: + async with sessions() as seed: + user = User( + username=f"rr_reset_{uuid4().hex[:16]}", + used_traffic=100, + data_limit_reset_strategy=DataLimitResetStrategy.day, + ) + user.created_at = datetime.now(UTC) - timedelta(days=2) + seed.add(user) + await seed.commit() + user_id = user.id + async with sessions() as stale, sessions() as winner: + candidates = await get_users_to_reset_data_usage(stale, user_ids=[user_id]) + assert len(candidates) == 1 + # Keep this transaction/snapshot open while the winner changes both + # User and usage history. The loser must really wait for its lock. + winning_user = await winner.get(User, user_id) + await reset_user_data_usage(winner, winning_user, reset_source=winner_source, commit=False) + await winner.flush() + locking_read_started = asyncio.Event() + + def on_statement(_conn, _cursor, statement, _params, _context, _many): + if "FOR UPDATE" in statement.upper(): + locking_read_started.set() + + event.listen(engine.sync_engine, "before_cursor_execute", on_statement) + pending = asyncio.create_task( + bulk_reset_user_data_usage(stale, candidates, reset_source=UserUsageResetSource.scheduled) + ) + try: + await asyncio.wait_for(locking_read_started.wait(), timeout=10) + assert not pending.done() + await winner.commit() + reset = await asyncio.wait_for(pending, timeout=15) + finally: + event.remove(engine.sync_engine, "before_cursor_execute", on_statement) + if not pending.done(): + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + assert [u.id for u in reset] == ([user_id] if winner_source is UserUsageResetSource.manual else []) + async with sessions() as verify: + logs = ( + await verify.scalars( + select(UserUsageResetLogs) + .where(UserUsageResetLogs.user_id == user_id) + .order_by(UserUsageResetLogs.id) + ) + ).all() + assert sum(log.used_traffic_at_reset for log in logs) == 100 + assert sum(log.reset_source == "scheduled" for log in logs) == 1 + assert [log.reset_source for log in logs] == ( + ["manual", "scheduled"] if winner_source is UserUsageResetSource.manual else ["scheduled"] + ) + finally: + if user_id is not None: + async with sessions() as cleanup: + await cleanup.execute(delete(User).where(User.id == user_id)) + await cleanup.commit() + await engine.dispose() + + +@pytest.mark.asyncio +@pytest.mark.skipif(not DATABASE_URL.startswith("mysql"), reason="requires MySQL/MariaDB REPEATABLE READ") +async def test_mysql_repeatable_read_next_plan_activation_uses_current_values(): + engine = create_async_engine(DATABASE_URL, isolation_level="REPEATABLE READ") + sessions = async_sessionmaker(engine, expire_on_commit=False, autoflush=False) + user_id = None + try: + async with sessions() as seed: + user = User( + username=f"rr_next_{uuid4().hex[:16]}", + data_limit=100, + used_traffic=10, + data_limit_reset_strategy=DataLimitResetStrategy.month, + ) + seed.add(user) + await seed.flush() + user_id = user.id + seed.add(NextPlan(user_id=user.id, user_template_id=None, data_limit=1000, add_remaining_traffic=True)) + await seed.commit() + async with sessions() as stale: + old_user = await get_user(stale, user.username) + assert old_user.next_plan.data_limit == 1000 + async with sessions() as writer: + current = await get_user(writer, user.username) + current.used_traffic = 20 + current.next_plan.data_limit = 2000 + await writer.commit() + # Ordinary reads still see the old snapshot and identity map. + assert (await get_user(stale, user.username)).next_plan.data_limit == 1000 + result = await reset_user_by_next(stale, old_user) + assert result.data_limit == 2080 + assert result.used_traffic == 0 + assert result.usage_logs[-1].used_traffic_at_reset == 20 + assert result.usage_logs[-1].reset_source == "next_plan" + assert result.next_plan is None + finally: + if user_id is not None: + async with sessions() as cleanup: + await cleanup.execute(delete(User).where(User.id == user_id)) + await cleanup.commit() + await engine.dispose() diff --git a/tests/test_user_usage_reset_source_migration.py b/tests/test_user_usage_reset_source_migration.py new file mode 100644 index 000000000..63666d806 --- /dev/null +++ b/tests/test_user_usage_reset_source_migration.py @@ -0,0 +1,99 @@ +from alembic.command import downgrade, upgrade +from alembic.config import Config +from alembic.script import ScriptDirectory +from sqlalchemy import create_engine, inspect, text + +PARENT_REVISION = "8e2f1a9c4b70" +RESET_SOURCE_REVISION = "b8e4e47b9f2c" + + +def _upgrade(database_url: str, revision: str) -> None: + config = Config("alembic.ini") + config.set_main_option("sqlalchemy.url", database_url) + upgrade(config, revision) + + +def test_legacy_usage_logs_are_not_backfilled_as_scheduled(tmp_path): + database_url = f"sqlite+aiosqlite:///{tmp_path / 'legacy-reset-source.db'}" + _upgrade(database_url, PARENT_REVISION) + + engine = create_engine(database_url.replace("sqlite+aiosqlite", "sqlite")) + with engine.begin() as connection: + connection.execute( + text( + "INSERT INTO user_usage_logs " + "(user_id, used_traffic_at_reset, reset_at) " + "VALUES (NULL, 1024, CURRENT_TIMESTAMP)" + ) + ) + + _upgrade(database_url, RESET_SOURCE_REVISION) + + with engine.begin() as connection: + # This is how an old process writes during a rolling upgrade: it knows + # nothing about reset_source and relies on the database default. + connection.execute( + text( + "INSERT INTO user_usage_logs " + "(user_id, used_traffic_at_reset, reset_at) " + "VALUES (NULL, 2048, CURRENT_TIMESTAMP)" + ) + ) + reset_sources = connection.execute(text("SELECT reset_source FROM user_usage_logs ORDER BY id")).scalars().all() + + assert reset_sources == ["legacy", "legacy"] + + +def test_standalone_head_and_downgrade_preserve_usage_history(tmp_path): + config = Config("alembic.ini") + assert ScriptDirectory.from_config(config).get_heads() == [RESET_SOURCE_REVISION] + database_url = f"sqlite+aiosqlite:///{tmp_path / 'reset-source-roundtrip.db'}" + config.set_main_option("sqlalchemy.url", database_url) + upgrade(config, "head") + engine = create_engine(database_url.replace("sqlite+aiosqlite", "sqlite")) + with engine.begin() as connection: + for source in ("legacy", "manual", "scheduled", "next_plan"): + connection.execute( + text( + "INSERT INTO user_usage_logs (user_id, used_traffic_at_reset, reset_at, reset_source) " + "VALUES (NULL, 123, '2026-01-02 03:04:05', :source)" + ), + {"source": source}, + ) + before = connection.execute( + text("SELECT id, used_traffic_at_reset, reset_at FROM user_usage_logs ORDER BY id") + ).all() + assert "ix_user_usage_logs_user_id_source_reset_at" in { + i["name"] for i in inspect(connection).get_indexes("user_usage_logs") + } + upgrade(config, "head") + with engine.connect() as connection: + assert connection.execute(text("SELECT reset_source FROM user_usage_logs ORDER BY id")).scalars().all() == [ + "legacy", + "manual", + "scheduled", + "next_plan", + ] + downgrade(config, PARENT_REVISION) + with engine.connect() as connection: + assert ( + connection.execute( + text("SELECT id, used_traffic_at_reset, reset_at FROM user_usage_logs ORDER BY id") + ).all() + == before + ) + assert "reset_source" not in {c["name"] for c in inspect(connection).get_columns("user_usage_logs")} + upgrade(config, "head") + with engine.connect() as connection: + assert ( + connection.execute( + text("SELECT id, used_traffic_at_reset, reset_at FROM user_usage_logs ORDER BY id") + ).all() + == before + ) + # Downgrade necessarily removes provenance; a re-upgrade must not guess it. + assert ( + connection.execute(text("SELECT reset_source FROM user_usage_logs ORDER BY id")).scalars().all() + == ["legacy"] * 4 + ) + engine.dispose() From 9b1c74b5396cb8d0e1437dab4f613c1469272b02 Mon Sep 17 00:00:00 2001 From: Rerowros <85008083+Rerowros@users.noreply.github.com> Date: Wed, 9 Sep 2026 21:34:37 +0400 Subject: [PATCH 2/4] test: compare reset dates at database precision --- tests/api/test_user_usage_reset_cycle.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/api/test_user_usage_reset_cycle.py b/tests/api/test_user_usage_reset_cycle.py index a70304134..2036625b1 100644 --- a/tests/api/test_user_usage_reset_cycle.py +++ b/tests/api/test_user_usage_reset_cycle.py @@ -324,12 +324,20 @@ async def test_optimized_reset_dates_match_history_and_refresh_after_manual_rese scheduled.reset_at = datetime.now(UTC) - timedelta(days=5) session.add(scheduled) await session.commit() + persisted_scheduled_at = await session.scalar( + select(UserUsageResetLogs.reset_at).where(UserUsageResetLogs.id == scheduled.id) + ) + assert persisted_scheduled_at is not None user = await get_user(session, user.username, load_usage_logs=False, load_lifetime_used_traffic=True) assert user.last_cycle_traffic_reset_at == user.last_traffic_reset_at - assert user.last_cycle_traffic_reset_at.replace(tzinfo=UTC) == scheduled.reset_at.replace(tzinfo=UTC) + assert user.last_cycle_traffic_reset_at.replace(tzinfo=UTC) == persisted_scheduled_at.replace(tzinfo=UTC) await reset_user_data_usage(session, user) + persisted_scheduled_after_reset = await session.scalar( + select(UserUsageResetLogs.reset_at).where(UserUsageResetLogs.id == scheduled.id) + ) + assert persisted_scheduled_after_reset == persisted_scheduled_at assert user.last_traffic_reset_at.replace(tzinfo=UTC) > user.last_cycle_traffic_reset_at.replace(tzinfo=UTC) - assert user.last_cycle_traffic_reset_at.replace(tzinfo=UTC) == scheduled.reset_at.replace(tzinfo=UTC) + assert user.last_cycle_traffic_reset_at.replace(tzinfo=UTC) == persisted_scheduled_after_reset.replace(tzinfo=UTC) @pytest.mark.asyncio From 44e085b7024822fbcc105d5b4ea744383b423b29 Mon Sep 17 00:00:00 2001 From: Rerowros <85008083+Rerowros@users.noreply.github.com> Date: Wed, 9 Sep 2026 21:47:26 +0400 Subject: [PATCH 3/4] fix(db): refresh persisted usage reset timestamps --- app/db/crud/user.py | 16 ++++++++++++++++ tests/api/test_user_usage_reset_cycle.py | 19 +++++++++++++++++-- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/app/db/crud/user.py b/app/db/crud/user.py index 44472dad1..c4e682710 100644 --- a/app/db/crud/user.py +++ b/app/db/crud/user.py @@ -190,6 +190,22 @@ async def refresh_and_load_user( load_groups: bool = True, ): await db.refresh(user) + if load_usage_logs: + # A committed DateTime can differ from the Python value retained in the + # identity map (for example, MySQL/MariaDB DATETIME second precision). + # Re-populate existing history instances before exposing properties + # which fall back from optimized query expressions to usage_logs. + usage_logs = list( + ( + await db.scalars( + select(UserUsageResetLogs) + .where(UserUsageResetLogs.user_id == user.id) + .order_by(UserUsageResetLogs.reset_at, UserUsageResetLogs.id) + .execution_options(populate_existing=True) + ) + ).all() + ) + set_committed_value(user, "usage_logs", usage_logs) await load_user_attrs( user, load_admin=load_admin, diff --git a/tests/api/test_user_usage_reset_cycle.py b/tests/api/test_user_usage_reset_cycle.py index 2036625b1..04db5e9f2 100644 --- a/tests/api/test_user_usage_reset_cycle.py +++ b/tests/api/test_user_usage_reset_cycle.py @@ -3,7 +3,7 @@ from uuid import uuid4 import pytest -from sqlalchemy import delete, event, func, select +from sqlalchemy import delete, event, func, select, update from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from app.db.crud.bulk import reset_all_users_data_usage @@ -320,10 +320,24 @@ async def test_optimized_reset_dates_match_history_and_refresh_after_manual_rese user = User(username=f"dates_{uuid4().hex[:16]}", data_limit_reset_strategy=DataLimitResetStrategy.month) session.add(user) await session.flush() + # Local SQLite test databases may reuse user IDs while retaining rows + # from earlier interrupted runs with foreign-key enforcement disabled. + await session.execute(delete(UserUsageResetLogs).where(UserUsageResetLogs.user_id == user.id)) scheduled = UserUsageResetLogs(user_id=user.id, used_traffic_at_reset=10, reset_source="scheduled") - scheduled.reset_at = datetime.now(UTC) - timedelta(days=5) + scheduled.reset_at = (datetime.now(UTC) - timedelta(days=5)).replace(microsecond=900_000) session.add(scheduled) await session.commit() + # Reproduce a backend which stores DATETIME without fractional seconds + # while expire_on_commit=False retains the original Python value. + canonical_scheduled_at = scheduled.reset_at.replace(microsecond=0) + timedelta(seconds=1) + await session.execute( + update(UserUsageResetLogs) + .where(UserUsageResetLogs.id == scheduled.id) + .values(reset_at=canonical_scheduled_at) + .execution_options(synchronize_session=False) + ) + await session.commit() + assert scheduled.reset_at != canonical_scheduled_at persisted_scheduled_at = await session.scalar( select(UserUsageResetLogs.reset_at).where(UserUsageResetLogs.id == scheduled.id) ) @@ -336,6 +350,7 @@ async def test_optimized_reset_dates_match_history_and_refresh_after_manual_rese select(UserUsageResetLogs.reset_at).where(UserUsageResetLogs.id == scheduled.id) ) assert persisted_scheduled_after_reset == persisted_scheduled_at + assert scheduled.reset_at.replace(tzinfo=UTC) == persisted_scheduled_after_reset.replace(tzinfo=UTC) assert user.last_traffic_reset_at.replace(tzinfo=UTC) > user.last_cycle_traffic_reset_at.replace(tzinfo=UTC) assert user.last_cycle_traffic_reset_at.replace(tzinfo=UTC) == persisted_scheduled_after_reset.replace(tzinfo=UTC) From b9360a91e448fde57dbffb475494328d135fdd5e Mon Sep 17 00:00:00 2001 From: Rerowros <85008083+Rerowros@users.noreply.github.com> Date: Wed, 9 Sep 2026 21:48:35 +0400 Subject: [PATCH 4/4] ci: verify persisted reset dates alongside concurrency regressions --- .github/workflows/test-database-migrations.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test-database-migrations.yml b/.github/workflows/test-database-migrations.yml index a3c93ebe5..f9b69ea4e 100644 --- a/.github/workflows/test-database-migrations.yml +++ b/.github/workflows/test-database-migrations.yml @@ -195,7 +195,7 @@ jobs: run: make check-migrations - name: Verify traffic reset races under REPEATABLE READ - run: .venv/bin/python -m pytest tests/api/test_user_usage_reset_cycle.py -k mysql_repeatable_read -q + run: .venv/bin/python -m pytest tests/api/test_user_usage_reset_cycle.py -k "mysql_repeatable_read or optimized_reset_dates" -q - name: Run Tests run: make test @@ -321,7 +321,7 @@ jobs: run: make check-migrations - name: Verify traffic reset races under REPEATABLE READ - run: .venv/bin/python -m pytest tests/api/test_user_usage_reset_cycle.py -k mysql_repeatable_read -q + run: .venv/bin/python -m pytest tests/api/test_user_usage_reset_cycle.py -k "mysql_repeatable_read or optimized_reset_dates" -q - name: Run Tests (multi-worker all-in-one) env: