i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-03 09:50:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

3
src/plugins/manager.py
View File

@@ -49,7 +49,8 @@ class PluginManager:
# notification 默认不加载，避免未配置插件（如 email）在启动时初始化失败并占用内存。
DEFAULT_ENABLED_PLUGIN_MODULES: dict[str, tuple[str, ...]] = {
"auth": ("api_key",),
"rate_limit": ("sliding_window",),
# 默认切到 token_bucket；全局 Redis 就绪时会自动切到分布式后端。
"rate_limit": ("token_bucket",),
"cache": ("memory",),
"monitor": ("prometheus",),
"token": ("claude",),

181
src/plugins/rate_limit/token_bucket.py
View File

@@ -29,6 +29,7 @@ class TokenBucket:
self.refill_rate = refill_rate
self.tokens = capacity
self.last_refill = time.time()
self.last_access_time = self.last_refill
def _refill(self) -> None:
"""补充令牌"""
@@ -36,6 +37,7 @@ class TokenBucket:
time_passed = now - self.last_refill
tokens_to_add = time_passed * self.refill_rate
self.last_access_time = now
if tokens_to_add > 0:
self.tokens = min(self.capacity, self.tokens + tokens_to_add)
self.last_refill = now
@@ -64,7 +66,7 @@ class TokenBucket:
def get_reset_time(self) -> datetime:
"""获取下次完全恢复的时间"""
if self.tokens >= self.capacity:
if self.tokens >= self.capacity or self.refill_rate <= 0:
return datetime.now(timezone.utc)
tokens_needed = self.capacity - self.tokens
@@ -82,6 +84,10 @@ class TokenBucketStrategy(RateLimitStrategy):
- 适合处理不均匀的流量模式
"""
DEFAULT_MAX_BUCKETS = 10000
DEFAULT_BUCKET_EXPIRY = 3600
DEFAULT_REDIS_RETRY_INTERVAL = 30.0
def __init__(self) -> None:
super().__init__("token_bucket")
self.buckets: dict[str, TokenBucket] = {}
@@ -90,11 +96,46 @@ class TokenBucketStrategy(RateLimitStrategy):
# 默认配置
self.default_capacity = 100 # 默认桶容量
self.default_refill_rate = 10 # 默认每秒补充10个令牌
self.max_buckets = self.DEFAULT_MAX_BUCKETS
self.bucket_expiry = self.DEFAULT_BUCKET_EXPIRY
self._last_cleanup_time: float = time.time()
self._cleanup_interval = 300 # 每 5 分钟检查一次清理
# 可选的 Redis 后端
self._redis_backend: RedisTokenBucketBackend | None = None
self._redis_checked = False
self._backend_mode = os.getenv("RATE_LIMIT_BACKEND", "auto").lower()
self._redis_retry_interval = self.DEFAULT_REDIS_RETRY_INTERVAL
self._next_redis_probe_time = 0.0
@staticmethod
def _is_unlimited_rate_limit(rate_limit: Any) -> bool:
"""显式传入 0/负数时，按“不限流”处理。"""
if rate_limit is None:
return False
try:
return int(rate_limit) <= 0
except (TypeError, ValueError):
return False
def _resolve_bucket_config(self, key: str, rate_limit: int | None = None) -> tuple[int, float]:
"""解析指定 key 当前应使用的桶容量和补充速率。"""
if rate_limit is not None:
normalized_rate_limit = int(rate_limit)
if normalized_rate_limit <= 0:
return 0, 0.0
return normalized_rate_limit, normalized_rate_limit / 60.0
if key.startswith("api_key:"):
return (
self.config.get("api_key_capacity", self.default_capacity),
self.config.get("api_key_refill_rate", self.default_refill_rate),
)
if key.startswith("user:"):
return (
self.config.get("user_capacity", self.default_capacity * 2),
self.config.get("user_refill_rate", self.default_refill_rate * 2),
)
return self.default_capacity, self.default_refill_rate
def _get_bucket(self, key: str, rate_limit: int | None = None) -> TokenBucket:
"""
@@ -107,41 +148,88 @@ class TokenBucketStrategy(RateLimitStrategy):
Returns:
令牌桶实例
"""
if key not in self.buckets:
# 如果提供了rate_limit参数（来自数据库），优先使用
if rate_limit is not None:
# rate_limit 是每分钟请求数，转换为令牌桶参数
capacity = rate_limit # 桶容量等于每分钟限制
refill_rate = rate_limit / 60.0 # 每秒补充的令牌数
# 否则根据key的不同前缀使用不同的配置
elif key.startswith("api_key:"):
capacity = self.config.get("api_key_capacity", self.default_capacity)
refill_rate = self.config.get("api_key_refill_rate", self.default_refill_rate)
elif key.startswith("user:"):
capacity = self.config.get("user_capacity", self.default_capacity * 2)
refill_rate = self.config.get("user_refill_rate", self.default_refill_rate * 2)
else:
capacity = self.default_capacity
refill_rate = self.default_refill_rate
capacity, refill_rate = self._resolve_bucket_config(key, rate_limit)
bucket = self.buckets.get(key)
if bucket is None:
bucket = TokenBucket(capacity, refill_rate)
self.buckets[key] = bucket
return bucket
self.buckets[key] = TokenBucket(capacity, refill_rate)
if bucket.capacity != capacity or bucket.refill_rate != refill_rate:
bucket._refill()
bucket.capacity = capacity
bucket.refill_rate = refill_rate
bucket.tokens = min(bucket.tokens, capacity)
return self.buckets[key]
return bucket
def _cleanup_expired_buckets(self) -> int:
"""清理长时间未访问的桶，避免 key 集合无限增长。"""
current_time = time.time()
expired_keys = [
key
for key, bucket in self.buckets.items()
if current_time - bucket.last_access_time > self.bucket_expiry
]
for key in expired_keys:
del self.buckets[key]
if expired_keys:
logger.info("清理了 {} 个过期的令牌桶", len(expired_keys))
return len(expired_keys)
def _evict_lru_buckets(self, count: int) -> int:
"""达到容量上限时淘汰最久未使用的桶。"""
if not self.buckets or count <= 0:
return 0
sorted_keys = sorted(self.buckets, key=lambda key: self.buckets[key].last_access_time)
evicted = 0
for key in sorted_keys[:count]:
del self.buckets[key]
evicted += 1
if evicted:
logger.warning("LRU 淘汰了 {} 个令牌桶（达到容量上限）", evicted)
return evicted
async def _maybe_cleanup(self) -> None:
"""定期清理或淘汰桶，控制进程内桶数量。"""
current_time = time.time()
if current_time - self._last_cleanup_time > self._cleanup_interval:
self._cleanup_expired_buckets()
self._last_cleanup_time = current_time
if len(self.buckets) >= self.max_buckets:
evict_count = max(1, self.max_buckets // 10)
self._evict_lru_buckets(evict_count)
def _want_redis_backend(self) -> bool:
return self._backend_mode in {"auto", "redis"}
async def _ensure_backend(self) -> None:
if self._redis_checked:
if self._redis_backend is not None:
return
self._redis_checked = True
if not self._want_redis_backend():
self._redis_checked = True
return
current_time = time.time()
if self._redis_checked and current_time < self._next_redis_probe_time:
return
redis_client = get_redis_client_sync()
if redis_client:
self._redis_backend = RedisTokenBucketBackend(redis_client)
self._redis_checked = True
self._next_redis_probe_time = 0.0
logger.info("速率限制改用 Redis 令牌桶后端")
elif self._backend_mode == "redis":
return
self._redis_checked = True
self._next_redis_probe_time = current_time + self._redis_retry_interval
if self._backend_mode == "redis":
logger.warning("RATE_LIMIT_BACKEND=redis 但 Redis 客户端不可用，回退到内存桶")
async def check_limit(self, key: str, **kwargs: Any) -> RateLimitResult:
@@ -155,11 +243,14 @@ class TokenBucketStrategy(RateLimitStrategy):
Returns:
速率限制检查结果
"""
await self._ensure_backend()
rate_limit = kwargs.get("rate_limit")
amount = kwargs.get("amount", 1)
if self._is_unlimited_rate_limit(rate_limit):
return RateLimitResult(allowed=True, remaining=0)
await self._ensure_backend()
if self._redis_backend:
return await self._redis_backend.peek(
key=key,
@@ -169,6 +260,7 @@ class TokenBucketStrategy(RateLimitStrategy):
)
async with self._lock:
await self._maybe_cleanup()
bucket = self._get_bucket(key, rate_limit)
remaining = bucket.get_remaining()
reset_at = bucket.get_reset_time()
@@ -203,13 +295,17 @@ class TokenBucketStrategy(RateLimitStrategy):
Returns:
是否成功消费
"""
rate_limit = kwargs.get("rate_limit")
if self._is_unlimited_rate_limit(rate_limit):
return True
await self._ensure_backend()
if self._redis_backend:
success, remaining = await self._redis_backend.consume(
key=key,
capacity=self._resolve_capacity(key, kwargs.get("rate_limit")),
refill_rate=self._resolve_refill_rate(key, kwargs.get("rate_limit")),
capacity=self._resolve_capacity(key, rate_limit),
refill_rate=self._resolve_refill_rate(key, rate_limit),
amount=amount,
)
if success:
@@ -219,7 +315,8 @@ class TokenBucketStrategy(RateLimitStrategy):
return success
async with self._lock:
bucket = self._get_bucket(key)
await self._maybe_cleanup()
bucket = self._get_bucket(key, rate_limit)
success = bucket.consume(amount)
if success:
@@ -270,6 +367,7 @@ class TokenBucketStrategy(RateLimitStrategy):
)
async with self._lock:
await self._maybe_cleanup()
bucket = self._get_bucket(key)
return {
"strategy": "token_bucket",
@@ -293,24 +391,17 @@ class TokenBucketStrategy(RateLimitStrategy):
super().configure(config)
self.default_capacity = config.get("default_capacity", self.default_capacity)
self.default_refill_rate = config.get("default_refill_rate", self.default_refill_rate)
self.max_buckets = int(config.get("max_buckets", self.max_buckets))
self.bucket_expiry = int(config.get("bucket_expiry", self.bucket_expiry))
self._cleanup_interval = int(config.get("cleanup_interval", self._cleanup_interval))
def _resolve_capacity(self, key: str, rate_limit: int | None = None) -> int:
if rate_limit is not None:
return rate_limit
if key.startswith("api_key:"):
return self.config.get("api_key_capacity", self.default_capacity)
if key.startswith("user:"):
return self.config.get("user_capacity", self.default_capacity * 2)
return self.default_capacity
capacity, _ = self._resolve_bucket_config(key, rate_limit)
return capacity
def _resolve_refill_rate(self, key: str, rate_limit: int | None = None) -> float:
if rate_limit is not None:
return rate_limit / 60.0
if key.startswith("api_key:"):
return self.config.get("api_key_refill_rate", self.default_refill_rate)
if key.startswith("user:"):
return self.config.get("user_refill_rate", self.default_refill_rate * 2)
return self.default_refill_rate
_, refill_rate = self._resolve_bucket_config(key, rate_limit)
return refill_rate
class RedisTokenBucketBackend:
@@ -365,6 +456,9 @@ class RedisTokenBucketBackend:
refill_rate: float,
amount: int,
) -> RateLimitResult:
if capacity <= 0 or refill_rate <= 0:
return RateLimitResult(allowed=True, remaining=0)
bucket_key = self._redis_key(key)
data = await self.redis.hmget(bucket_key, "tokens", "timestamp")
tokens = data[0]
@@ -372,7 +466,7 @@ class RedisTokenBucketBackend:
if tokens is None or last_refill is None:
remaining = capacity
reset_at = datetime.now(timezone.utc) + timedelta(seconds=capacity / refill_rate)
reset_at = datetime.now(timezone.utc)
else:
tokens_value = float(tokens)
last_refill_value = float(last_refill)
@@ -407,6 +501,9 @@ class RedisTokenBucketBackend:
refill_rate: float,
amount: int,
) -> tuple[bool, int]:
if capacity <= 0 or refill_rate <= 0:
return True, 0
result = await self._consume_script(
keys=[self._redis_key(key)],
args=[time.time(), capacity, refill_rate, amount],