feat: 流式空闲超时、健康监控查询优化、限流桶内存上限与维护清理修复

Close #233

Co-authored-by: AAEE86 <ppk0227@hotmail.com>

- cli_monitor_mixin: 引入 STREAM_IDLE_TIMEOUT_SECONDS(可通过环境变量配置),
  流传输开始后若超出空闲窗口无新 chunk 则提前取消并返回 504,避免长时间挂起
- stream_context: 新增 managed_recorded_bodies 上下文管理器,确保 chunks 在
  telemetry 完成后及时释放;stream_telemetry 使用该接口统一管理 response body 构建
- health endpoint: 将状态聚合改为 GROUP BY 直接统计,事件列表按 api_format
  单独查询,避免单次 limit 拉取大量记录导致的遗漏与性能问题;同时过滤不活跃
  provider/endpoint,与公开健康接口保持一致
- endpoint health service: 修正时间线数据按 endpoint_id 而非 key_id 聚合
- token_bucket: 引入 max_buckets/bucket_expiry 上限与定时清理,防止内存无限增长;
  修复 refill_rate=0 时 get_reset_time 除零异常;新增 _is_unlimited_rate_limit 判断
- maintenance_scheduler: 调整清理顺序(先删整行再按窗口清理),新增 newer_than
  边界参数,避免同一行在同一轮中被重复改写
- sync_execute: 新增 create_pending_usage 开关,允许已预创建记录的调用方跳过重复创建
- quota_reader / provider_ops balance: 小幅修复与健壮性提升
- Dockerfile: 添加 MALLOC_ARENA_MAX=2 环境变量以降低 gunicorn worker RSS
- 补充相关测试覆盖
This commit is contained in:
fawney19
2026-03-18 23:38:26 +08:00
parent 3d5b6141a5
commit 1d72a8f9c1
37 changed files with 1787 additions and 607 deletions

View File

@@ -0,0 +1,111 @@
from __future__ import annotations
from datetime import datetime, timezone
import pytest
import src.plugins.rate_limit.token_bucket as token_bucket_module
from src.plugins.rate_limit.token_bucket import RedisTokenBucketBackend, TokenBucketStrategy
@pytest.mark.asyncio
async def test_token_bucket_cleans_up_expired_buckets(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("RATE_LIMIT_BACKEND", "memory")
strategy = TokenBucketStrategy()
strategy.configure({"bucket_expiry": 1, "cleanup_interval": 0})
await strategy.check_limit("api_key:stale")
strategy.buckets["api_key:stale"].last_access_time -= 3600
await strategy.check_limit("api_key:fresh")
assert "api_key:stale" not in strategy.buckets
assert "api_key:fresh" in strategy.buckets
@pytest.mark.asyncio
async def test_token_bucket_reconfigures_existing_bucket_when_rate_limit_changes(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("RATE_LIMIT_BACKEND", "memory")
strategy = TokenBucketStrategy()
await strategy.check_limit("user:42", rate_limit=120)
bucket = strategy.buckets["user:42"]
bucket.tokens = 90
await strategy.check_limit("user:42", rate_limit=30)
updated_bucket = strategy.buckets["user:42"]
assert updated_bucket.capacity == 30
assert updated_bucket.refill_rate == 0.5
assert updated_bucket.tokens <= 30
@pytest.mark.asyncio
async def test_token_bucket_treats_non_positive_dynamic_rate_limit_as_unlimited(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("RATE_LIMIT_BACKEND", "memory")
strategy = TokenBucketStrategy()
result = await strategy.check_limit("public_ip:test", rate_limit=0)
consumed = await strategy.consume("public_ip:test", amount=1, rate_limit=0)
assert result.allowed is True
assert consumed is True
assert "public_ip:test" not in strategy.buckets
class _FakeRedisClient:
async def hmget(self, _key: str, *_fields: str) -> list[None]:
return [None, None]
def register_script(self, _script: str): # type: ignore[no-untyped-def]
async def _runner(*args, **kwargs): # type: ignore[no-untyped-def]
return [1, 0, 0]
return _runner
@pytest.mark.asyncio
async def test_token_bucket_retries_redis_backend_probe_after_initial_miss(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("RATE_LIMIT_BACKEND", "auto")
strategy = TokenBucketStrategy()
strategy._redis_retry_interval = 0
fake_redis = _FakeRedisClient()
calls = {"count": 0}
def _fake_get_redis_client_sync(): # type: ignore[no-untyped-def]
calls["count"] += 1
if calls["count"] == 1:
return None
return fake_redis
monkeypatch.setattr(
token_bucket_module,
"get_redis_client_sync",
_fake_get_redis_client_sync,
)
await strategy.check_limit("public_ip:first")
assert strategy._redis_backend is None
await strategy.check_limit("public_ip:second")
assert strategy._redis_backend is not None
assert calls["count"] == 2
@pytest.mark.asyncio
async def test_redis_token_bucket_missing_bucket_reports_reset_now() -> None:
backend = RedisTokenBucketBackend(_FakeRedisClient())
result = await backend.peek("public_ip:test", capacity=60, refill_rate=1.0, amount=1)
assert result.allowed is True
assert result.remaining == 60
assert result.reset_at is not None
assert abs((result.reset_at - datetime.now(timezone.utc)).total_seconds()) < 2