feat(rate-limit): 实现分层 RPM 限速,支持系统默认/用户/独立Key三级配置

- 新增用户级 rate_limit 字段,支持系统默认/用户自定义/不限制三种模式
- 独立 Key 的 rate_limit 语义调整:null=跟随系统默认,0=不限制,>0=自定义
- 实现 UserRpmLimiter 基于 Redis sliding window 的 RPM 限速引擎
- Pipeline 请求流程集成用户级 RPM 检查
- 管理后台和用户面板新增 RPM 限速配置与实时状态查看
- 系统设置新增全局默认 RPM 配置项
- 迁移脚本回填现有 API Key 的 rate_limit 默认值
- 新增用户/Key RPM 状态监控 API 和前端展示

Closes #231

Co-authored-by: LewisPen <LewisPen@nyadoo.com>
This commit is contained in:
fawney19
2026-03-15 14:22:59 +08:00
parent 920a383136
commit f92b0943b5
35 changed files with 2051 additions and 238 deletions

View File

@@ -0,0 +1,148 @@
from __future__ import annotations
from collections.abc import AsyncGenerator
from unittest.mock import AsyncMock, MagicMock
import pytest
from src.services.rate_limit.user_rpm_limiter import RpmCheckResult, UserRpmLimiter
@pytest.fixture
async def limiter(monkeypatch: pytest.MonkeyPatch) -> AsyncGenerator[UserRpmLimiter]:
limiter = UserRpmLimiter()
limiter._redis = None
limiter._memory_counts.clear()
await limiter.close()
# 阻止 initialize() 连接真实 Redis确保纯内存模式
monkeypatch.setattr(limiter, "initialize", AsyncMock())
yield limiter
limiter._memory_counts.clear()
limiter._redis = None
await limiter.close()
@pytest.mark.asyncio
async def test_check_and_consume_enforces_user_scope_in_memory(
limiter: UserRpmLimiter,
) -> None:
user_key = limiter.get_user_rpm_key("user-1")
key_key = limiter.get_key_rpm_key("key-1")
first = await limiter.check_and_consume(
user_rpm_key=user_key,
user_rpm_limit=1,
key_rpm_key=key_key,
key_rpm_limit=0,
)
second = await limiter.check_and_consume(
user_rpm_key=user_key,
user_rpm_limit=1,
key_rpm_key=key_key,
key_rpm_limit=0,
)
assert first == RpmCheckResult(allowed=True, remaining=0)
assert second.allowed is False
assert second.scope == "user"
assert second.limit == 1
assert second.remaining == 0
@pytest.mark.asyncio
async def test_check_and_consume_enforces_key_scope_in_memory(
limiter: UserRpmLimiter,
) -> None:
user_key = limiter.get_user_rpm_key("user-1")
key_key = limiter.get_key_rpm_key("key-1")
first = await limiter.check_and_consume(
user_rpm_key=user_key,
user_rpm_limit=3,
key_rpm_key=key_key,
key_rpm_limit=1,
)
second = await limiter.check_and_consume(
user_rpm_key=user_key,
user_rpm_limit=3,
key_rpm_key=key_key,
key_rpm_limit=1,
)
assert first.allowed is True
assert second.allowed is False
assert second.scope == "key"
assert second.limit == 1
@pytest.mark.asyncio
async def test_check_and_consume_redis_failure_falls_back_to_memory_when_fail_close(
monkeypatch: pytest.MonkeyPatch,
limiter: UserRpmLimiter,
) -> None:
fake_redis = MagicMock()
fake_redis.eval = AsyncMock(side_effect=RuntimeError("redis down"))
limiter._redis = fake_redis
monkeypatch.setattr(
"src.services.rate_limit.user_rpm_limiter.config.rate_limit_fail_open", False
)
user_key = limiter.get_user_rpm_key("user-1")
key_key = limiter.get_key_rpm_key("key-1")
first = await limiter.check_and_consume(
user_rpm_key=user_key,
user_rpm_limit=1,
key_rpm_key=key_key,
key_rpm_limit=0,
)
second = await limiter.check_and_consume(
user_rpm_key=user_key,
user_rpm_limit=1,
key_rpm_key=key_key,
key_rpm_limit=0,
)
assert first.allowed is True
assert second.allowed is False
assert second.scope == "user"
@pytest.mark.asyncio
async def test_check_and_consume_redis_failure_fail_open_allows_request(
monkeypatch: pytest.MonkeyPatch,
limiter: UserRpmLimiter,
) -> None:
fake_redis = MagicMock()
fake_redis.eval = AsyncMock(side_effect=RuntimeError("redis down"))
limiter._redis = fake_redis
monkeypatch.setattr(
"src.services.rate_limit.user_rpm_limiter.config.rate_limit_fail_open", True
)
result = await limiter.check_and_consume(
user_rpm_key=limiter.get_user_rpm_key("user-1"),
user_rpm_limit=1,
key_rpm_key=limiter.get_key_rpm_key("key-1"),
key_rpm_limit=0,
)
assert result.allowed is True
@pytest.mark.asyncio
async def test_check_and_consume_skips_when_all_limits_are_zero(
limiter: UserRpmLimiter,
) -> None:
result = await limiter.check_and_consume(
user_rpm_key=limiter.get_user_rpm_key("user-1"),
user_rpm_limit=0,
key_rpm_key=limiter.get_key_rpm_key("key-1"),
key_rpm_limit=0,
)
assert result.allowed is True
assert result.scope is None
assert limiter._memory_counts == {}