feat: 添加用户配额自动重置功能及相关配置

This commit is contained in:
YorhaL
2026-02-03 22:04:31 +08:00
parent 6f363a7703
commit bac0e7fb20
5 changed files with 584 additions and 0 deletions

View File

@@ -624,6 +624,16 @@ class AdminSetSystemConfigAdapter(AdminApiAdapter):
except Exception as e:
logger.warning(f"更新签到任务时间失败: {e}")
# 如果更新的是用户配额重置任务时间,动态更新调度器
if self.key == "user_quota_reset_time" and value:
try:
from src.services.system.maintenance_scheduler import get_maintenance_scheduler
scheduler = get_maintenance_scheduler()
scheduler.update_user_quota_reset_time(value)
except Exception as e:
logger.warning(f"更新用户配额重置任务时间失败: {e}")
# 返回时不暴露加密后的值
display_value = "********" if self.key in self.ENCRYPTED_KEYS else config.value

View File

@@ -123,6 +123,18 @@ class SystemConfigService:
"value": "01:05",
"description": "Provider 自动签到执行时间HH:MM 格式24小时制",
},
"enable_user_quota_reset": {
"value": False,
"description": "是否启用用户配额自动重置任务(按配置时间触发,按周期执行)",
},
"user_quota_reset_time": {
"value": "05:00",
"description": "用户配额自动重置执行时间HH:MM 格式24小时制",
},
"user_quota_reset_interval_days": {
"value": 1,
"description": "用户配额重置周期(天数)",
},
"provider_priority_mode": {
"value": "provider",
"description": "优先级策略provider(提供商优先模式) 或 global_key(全局Key优先模式)",

View File

@@ -38,6 +38,8 @@ class MaintenanceScheduler:
# 签到任务的 job_id
CHECKIN_JOB_ID = "provider_checkin"
# 用户配额重置任务的 job_id
USER_QUOTA_RESET_JOB_ID = "user_quota_reset"
def __init__(self) -> None:
self.running = False
@@ -57,6 +59,19 @@ class MaintenanceScheduler:
finally:
db.close()
def _get_user_quota_reset_time(self) -> tuple[int, int]:
"""获取用户配额重置任务的执行时间
Returns:
(hour, minute) 元组
"""
db = create_session()
try:
time_str = SystemConfigService.get_config(db, "user_quota_reset_time", "05:00")
return self._parse_user_quota_reset_time_string(time_str)
finally:
db.close()
@staticmethod
def _parse_time_string(time_str: str) -> tuple[int, int]:
"""解析时间字符串为 (hour, minute) 元组
@@ -80,6 +95,26 @@ class MaintenanceScheduler:
except (ValueError, IndexError):
return (1, 5)
@staticmethod
def _parse_user_quota_reset_time_string(time_str: str) -> tuple[int, int]:
"""解析用户配额重置时间字符串为 (hour, minute) 元组
Returns:
(hour, minute) 元组,解析失败返回默认值 (5, 0)
"""
try:
if not time_str or ":" not in time_str:
return (5, 0)
parts = time_str.split(":")
hour = int(parts[0])
minute = int(parts[1])
# 验证范围
if 0 <= hour <= 23 and 0 <= minute <= 59:
return (hour, minute)
return (5, 0)
except (ValueError, IndexError):
return (5, 0)
def update_checkin_time(self, time_str: str) -> bool:
"""更新签到任务的执行时间
@@ -103,6 +138,29 @@ class MaintenanceScheduler:
return success
def update_user_quota_reset_time(self, time_str: str) -> bool:
"""更新用户配额重置任务的执行时间
Args:
time_str: HH:MM 格式的时间字符串
Returns:
是否成功更新
"""
hour, minute = self._parse_user_quota_reset_time_string(time_str)
scheduler = get_scheduler()
success = scheduler.reschedule_cron_job(
self.USER_QUOTA_RESET_JOB_ID,
hour=hour,
minute=minute,
)
if success:
logger.info(f"用户配额重置任务时间已更新为: {hour:02d}:{minute:02d}")
return success
def get_checkin_job_info(self) -> dict | None:
"""获取签到任务的信息
@@ -193,6 +251,16 @@ class MaintenanceScheduler:
name="Provider签到",
)
# 用户配额重置任务 - 根据配置时间执行(按周期配置决定是否执行)
quota_reset_hour, quota_reset_minute = self._get_user_quota_reset_time()
scheduler.add_cron_job(
self._scheduled_user_quota_reset,
hour=quota_reset_hour,
minute=quota_reset_minute,
job_id=self.USER_QUOTA_RESET_JOB_ID,
name="用户配额自动重置",
)
# 启动时执行一次初始化任务
asyncio.create_task(self._run_startup_tasks())
@@ -260,6 +328,10 @@ class MaintenanceScheduler:
"""Provider 签到任务(定时调用)"""
await self._perform_provider_checkin()
async def _scheduled_user_quota_reset(self) -> None:
"""用户配额重置任务(定时调用)"""
await self._perform_user_quota_reset()
# ========== 实际任务实现 ==========
async def _perform_stats_aggregation(self, backfill: bool = False) -> None:
@@ -685,6 +757,111 @@ class MaintenanceScheduler:
if db is not None:
db.close()
async def _perform_user_quota_reset(self) -> None:
"""执行用户配额自动重置任务
适用范围:
- 普通用户role=user
- 未删除is_deleted=false
- 仅对 quota_usd != NULL 的用户生效
"""
db = create_session()
try:
# 检查是否启用用户配额重置
if not SystemConfigService.get_config(db, "enable_user_quota_reset", False):
logger.info("用户配额自动重置已禁用,跳过任务")
return
# 重置周期(天数),不限制上限
interval_value = SystemConfigService.get_config(db, "user_quota_reset_interval_days", 1)
try:
interval_days = int(interval_value)
except Exception:
interval_days = 1
if interval_days < 1:
interval_days = 1
# 滚动计算根据上次执行日APP_TIMEZONE判断是否到期
last_reset_at = SystemConfigService.get_config(db, "user_quota_last_reset_at")
should_run = True
if last_reset_at:
last_dt: datetime | None = None
try:
if isinstance(last_reset_at, str):
last_dt = datetime.fromisoformat(last_reset_at)
except Exception:
last_dt = None
if last_dt is None:
logger.warning("user_quota_last_reset_at 格式无效,视为需要执行一次")
else:
if last_dt.tzinfo is None:
last_dt = last_dt.replace(tzinfo=timezone.utc)
from zoneinfo import ZoneInfo
from src.services.system.scheduler import APP_TIMEZONE
tz = ZoneInfo(APP_TIMEZONE)
now_local = datetime.now(tz)
last_local_date = last_dt.astimezone(tz).date()
days_since_reset = (now_local.date() - last_local_date).days
if days_since_reset < 0:
logger.warning(
"user_quota_last_reset_at 在未来,跳过本次用户配额自动重置"
)
should_run = False
elif days_since_reset < interval_days:
logger.info(
f"用户配额自动重置未到周期,跳过任务({days_since_reset}/{interval_days}天)"
)
should_run = False
if not should_run:
return
from src.core.enums import UserRole
from src.models.database import User as DBUser
now_utc = datetime.now(timezone.utc)
reset_count = (
db.query(DBUser)
.filter(
DBUser.role == UserRole.USER,
DBUser.is_deleted.is_(False),
DBUser.quota_usd.isnot(None),
)
.update(
{
DBUser.used_usd: 0.0,
DBUser.updated_at: now_utc,
},
synchronize_session=False,
)
)
db.commit()
# 记录 last_reset_at成功执行后更新滚动计算用
SystemConfigService.set_config(
db,
"user_quota_last_reset_at",
now_utc.isoformat(),
"用户配额自动重置的上次执行时间UTC内部使用",
)
logger.info(f"用户配额自动重置完成: interval_days={interval_days}, 重置用户数={reset_count}")
except Exception as e:
logger.exception(f"用户配额自动重置任务执行失败: {e}")
try:
db.rollback()
except Exception:
pass
finally:
db.close()
async def _perform_cleanup(self) -> None:
"""执行清理任务"""
db = create_session()