34 lines
1.3 KiB
Python
34 lines
1.3 KiB
Python
"""接口限流：Redis 固定窗口计数（Nginx 层之外的第二道，对齐《02-技术文档》4.5）"""
|
||
from datetime import datetime
|
||
|
||
import redis.asyncio as aioredis
|
||
|
||
from app.core import errors
|
||
from app.core.response import BizException
|
||
|
||
|
||
async def hit(redis: aioredis.Redis, key: str, limit: int, window_seconds: int, message: str | None = None) -> None:
|
||
"""固定窗口计数，超限抛 1005(HTTP 429)。key 需自带维度与窗口标识。"""
|
||
n = await redis.incr(key)
|
||
if n == 1:
|
||
await redis.expire(key, window_seconds)
|
||
if n > limit:
|
||
raise BizException(errors.RATE_LIMITED, message, http_status=429)
|
||
|
||
|
||
def window_key(window_seconds: int) -> str:
|
||
"""窗口标识：分钟 %Y%m%d%H%M / 小时 %Y%m%d%H / 天 %Y%m%d"""
|
||
now = datetime.now()
|
||
if window_seconds >= 86400:
|
||
return f"{now:%Y%m%d}"
|
||
if window_seconds >= 3600:
|
||
return f"{now:%Y%m%d%H}"
|
||
return f"{now:%Y%m%d%H%M}"
|
||
|
||
|
||
async def hit_subject(
|
||
redis: aioredis.Redis, subject: str, action: str, limit: int, window_seconds: int, message: str | None = None
|
||
) -> None:
|
||
"""按主体（用户ID/设备ID/IP）+ 动作限流：rl:{action}:{subject}:{窗口}"""
|
||
await hit(redis, f"rl:{action}:{subject}:{window_key(window_seconds)}", limit, window_seconds, message)
|