62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from typing import Any
|
|
|
|
from app.enums.private_portrait import PRIVATE_PORTRAIT_ACTION_QPS_LIMITS, PrivatePortraitEventSource, PrivatePortraitEventStatus, PrivatePortraitEventType
|
|
from app.services.operation_log_service import log_operation_event
|
|
from app.utils.redis import get_redis
|
|
|
|
DOMAIN = "private_portrait"
|
|
|
|
|
|
class PrivatePortraitRateLimitExceeded(RuntimeError):
|
|
pass
|
|
|
|
|
|
async def acquire_private_portrait_action_token(
|
|
*,
|
|
action: str,
|
|
wait_timeout_seconds: float = 2.0,
|
|
for_celery: bool = False,
|
|
) -> bool:
|
|
"""Redis 分布式 QPS 限制。Redis 不可用时降级放行，避免影响主功能。"""
|
|
limit = int(PRIVATE_PORTRAIT_ACTION_QPS_LIMITS.get(action, 1))
|
|
if limit <= 0:
|
|
return True
|
|
deadline = time.monotonic() + max(0.0, wait_timeout_seconds)
|
|
while True:
|
|
ok = await _try_take(action, limit)
|
|
if ok:
|
|
return True
|
|
if time.monotonic() >= deadline:
|
|
event_type = PrivatePortraitEventType.ARK_API_RATE_LIMIT_WAIT.value if for_celery else PrivatePortraitEventType.ARK_API_RATE_LIMIT_REJECT.value
|
|
log_operation_event(
|
|
domain=DOMAIN,
|
|
event_type=event_type,
|
|
event_status=PrivatePortraitEventStatus.SKIPPED.value if for_celery else PrivatePortraitEventStatus.FAILED.value,
|
|
source=PrivatePortraitEventSource.CELERY.value if for_celery else PrivatePortraitEventSource.API.value,
|
|
remote_action=action,
|
|
detail={"limit": limit, "wait_timeout_seconds": wait_timeout_seconds},
|
|
message="火山私域真人素材 API 触发本地 QPS 限制",
|
|
)
|
|
if for_celery:
|
|
return False
|
|
raise PrivatePortraitRateLimitExceeded("请求过于频繁，请稍后再试")
|
|
await asyncio.sleep(0.05)
|
|
|
|
|
|
async def _try_take(action: str, limit: int) -> bool:
|
|
client = get_redis()
|
|
if client is None:
|
|
return True
|
|
key = f"private_portrait:qps:{action}:{int(time.time())}"
|
|
try:
|
|
count = await client.incr(key)
|
|
if count == 1:
|
|
await client.expire(key, 2)
|
|
return int(count) <= limit
|
|
except Exception:
|
|
return True
|