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

@@ -36,6 +36,22 @@ def test_codex_reader_preserves_summary_formats() -> None:
assert credits_reader.display_summary() == "积分 12.35"
def test_codex_reader_hides_reset_countdown_when_remaining_is_full() -> None:
reader = get_quota_reader(
"codex",
{
"codex": {
"primary_used_percent": 0.0,
"primary_reset_seconds": 266400,
"secondary_used_percent": 0.0,
"secondary_reset_seconds": 3600,
}
},
)
assert reader.display_summary() == "周剩余 100.0% | 5H剩余 100.0%"
def test_antigravity_reader_keeps_used_percent_fallbacks() -> None:
reader = get_quota_reader(
"antigravity",
@@ -88,6 +104,20 @@ def test_extract_reset_seconds_uses_codex_weekly_reset_for_codex_provider() -> N
assert extract_reset_seconds(key_obj) == pytest.approx(1800.0)
def test_extract_reset_seconds_codex_weekly_full_quota_returns_none() -> None:
key_obj = SimpleNamespace(
provider_type="codex",
upstream_metadata={
"codex": {
"primary_used_percent": 0.0,
"primary_reset_seconds": 1800.0,
}
},
)
assert extract_reset_seconds(key_obj) is None
def test_extract_reset_seconds_prefers_codex_reset_at(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr("src.services.provider.pool.dimensions._helpers.time.time", lambda: 1000.0)
key_obj = SimpleNamespace(

View File

@@ -1,6 +1,7 @@
from __future__ import annotations
import inspect
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from unittest.mock import MagicMock
@@ -124,44 +125,29 @@ async def test_candidate_cleanup_uses_dedicated_retention_and_batch_settings(
assert batch_two.closed is True
def test_cleanup_body_fields_loads_ids_then_processes_records_individually(
def test_cleanup_body_fields_batches_records_with_single_commit(
monkeypatch: pytest.MonkeyPatch,
) -> None:
scheduler = MaintenanceScheduler()
class _IdBatchSession:
def __init__(self, ids: list[str]) -> None:
self.ids = ids
self.closed = False
self.query_obj = MagicMock()
filtered = self.query_obj.filter.return_value
filtered.filter.return_value.order_by.return_value.limit.return_value.all.return_value = [
SimpleNamespace(id=value) for value in ids
]
def query(self, *args): # type: ignore[no-untyped-def]
self.query_args = args
return self.query_obj
def rollback(self) -> None:
raise AssertionError("rollback should not be called")
def close(self) -> None:
self.closed = True
class _RecordSession:
def __init__(self, record: SimpleNamespace) -> None:
self.record = record
class _BatchSession:
def __init__(self, records: list[SimpleNamespace]) -> None:
self.records = records
self.closed = False
self.committed = False
self.executed = 0
self.query_obj = MagicMock()
self.query_obj.filter.return_value.first.return_value = record
filtered = self.query_obj.filter.return_value
filtered.filter.return_value.order_by.return_value.limit.return_value.all.return_value = (
records
)
def query(self, *args): # type: ignore[no-untyped-def]
self.query_args = args
return self.query_obj
def execute(self, _statement): # type: ignore[no-untyped-def]
self.executed += 1
return SimpleNamespace(rowcount=1)
def commit(self) -> None:
@@ -173,27 +159,26 @@ def test_cleanup_body_fields_loads_ids_then_processes_records_individually(
def close(self) -> None:
self.closed = True
batch_one = _IdBatchSession(["usage-1", "usage-2"])
record_one = _RecordSession(
SimpleNamespace(
id="usage-1",
request_body={"hello": "world"},
response_body=None,
provider_request_body=None,
client_response_body=None,
)
batch_one = _BatchSession(
[
SimpleNamespace(
id="usage-1",
request_body={"hello": "world"},
response_body=None,
provider_request_body=None,
client_response_body=None,
),
SimpleNamespace(
id="usage-2",
request_body=None,
response_body={"ok": True},
provider_request_body=None,
client_response_body=None,
),
]
)
record_two = _RecordSession(
SimpleNamespace(
id="usage-2",
request_body=None,
response_body={"ok": True},
provider_request_body=None,
client_response_body=None,
)
)
batch_two = _IdBatchSession([])
sessions = iter([batch_one, record_one, record_two, batch_two])
batch_two = _BatchSession([])
sessions = iter([batch_one, batch_two])
monkeypatch.setattr(
maintenance_scheduler_module,
@@ -212,12 +197,173 @@ def test_cleanup_body_fields_loads_ids_then_processes_records_individually(
)
assert compressed == 2
assert batch_one.query_args == (maintenance_scheduler_module.Usage.id,)
assert len(record_one.query_args) == 5
assert len(record_two.query_args) == 5
assert len(batch_one.query_args) == 5
assert batch_one.executed == 2
assert batch_one.committed is True
assert batch_one.closed is True
assert batch_two.closed is True
assert record_one.committed is True
assert record_two.committed is True
assert record_one.closed is True
assert record_two.closed is True
def test_cleanup_header_fields_clears_client_response_headers(
monkeypatch: pytest.MonkeyPatch,
) -> None:
scheduler = MaintenanceScheduler()
class _BatchSession:
def __init__(self, ids: list[str]) -> None:
self.ids = ids
self.closed = False
self.committed = False
self.query_obj = MagicMock()
self.filtered_by_time = MagicMock()
self.filtered_by_headers = MagicMock()
self.query_obj.filter.return_value = self.filtered_by_time
self.filtered_by_time.filter.return_value = self.filtered_by_headers
self.filtered_by_headers.order_by.return_value.limit.return_value.all.return_value = [
SimpleNamespace(id=value) for value in ids
]
self.executed_statements: list[str] = []
def query(self, *args): # type: ignore[no-untyped-def]
self.query_args = args
return self.query_obj
def execute(self, statement): # type: ignore[no-untyped-def]
self.executed_statements.append(str(statement))
return SimpleNamespace(rowcount=len(self.ids))
def commit(self) -> None:
self.committed = True
def rollback(self) -> None:
raise AssertionError("rollback should not be called")
def close(self) -> None:
self.closed = True
batch_one = _BatchSession(["usage-1"])
batch_two = _BatchSession([])
sessions = iter([batch_one, batch_two])
monkeypatch.setattr(
maintenance_scheduler_module,
"create_session",
lambda: next(sessions),
)
cleaned = scheduler._cleanup_header_fields(
cutoff_time=SimpleNamespace(), # type: ignore[arg-type]
batch_size=1000,
)
header_filter = str(batch_one.filtered_by_time.filter.call_args.args[0])
assert cleaned == 1
assert batch_one.query_args == (maintenance_scheduler_module.Usage.id,)
assert "client_response_headers" in header_filter
assert "client_response_headers" in batch_one.executed_statements[0]
assert batch_one.committed is True
assert batch_one.closed is True
assert batch_two.closed is True
@pytest.mark.asyncio
async def test_perform_cleanup_deletes_first_and_uses_non_overlapping_windows(
monkeypatch: pytest.MonkeyPatch,
) -> None:
scheduler = MaintenanceScheduler()
fixed_now = datetime(2026, 3, 18, 3, 0, 0, tzinfo=timezone.utc)
class _FakeDateTime(datetime):
@classmethod
def now(cls, tz=None): # type: ignore[override]
if tz is None:
return fixed_now.replace(tzinfo=None)
return fixed_now.astimezone(tz)
class _FakeLoop:
async def run_in_executor(self, _executor, func): # type: ignore[no-untyped-def]
return func()
class _ConfigSession:
def close(self) -> None:
return None
calls: list[tuple[str, datetime, int, datetime | None]] = []
def _record(name: str, count: int):
def _inner(
cutoff_time: datetime,
batch_size: int,
*,
newer_than: datetime | None = None,
) -> int:
calls.append((name, cutoff_time, batch_size, newer_than))
return count
return _inner
config_values = {
"enable_auto_cleanup": True,
"detail_log_retention_days": 7,
"compressed_log_retention_days": 30,
"header_retention_days": 90,
"log_retention_days": 365,
"cleanup_batch_size": 123,
"auto_delete_expired_keys": False,
}
monkeypatch.setattr(maintenance_scheduler_module, "datetime", _FakeDateTime)
monkeypatch.setattr(
maintenance_scheduler_module.asyncio, "get_running_loop", lambda: _FakeLoop()
)
monkeypatch.setattr(
maintenance_scheduler_module,
"create_session",
lambda: _ConfigSession(),
)
monkeypatch.setattr(
maintenance_scheduler_module.SystemConfigService,
"get_config",
lambda _db, key, default=None: config_values.get(key, default),
)
monkeypatch.setattr(
scheduler,
"_delete_old_records",
lambda cutoff_time, batch_size: calls.append(("delete", cutoff_time, batch_size, None))
or 5,
)
monkeypatch.setattr(
scheduler,
"_cleanup_header_fields",
_record("header", 4),
)
monkeypatch.setattr(
scheduler,
"_cleanup_stale_body_fields",
_record("body", 3),
)
monkeypatch.setattr(
scheduler,
"_cleanup_body_fields",
_record("compress", 2),
)
monkeypatch.setattr(
maintenance_scheduler_module.ApiKeyService,
"cleanup_expired_keys",
lambda _db, auto_delete=False: 0,
)
await scheduler._perform_cleanup()
detail_cutoff = fixed_now - timedelta(days=7)
compressed_cutoff = fixed_now - timedelta(days=30)
header_cutoff = fixed_now - timedelta(days=90)
log_cutoff = fixed_now - timedelta(days=365)
assert calls == [
("delete", log_cutoff, 123, None),
("header", header_cutoff, 123, log_cutoff),
("body", compressed_cutoff, 123, log_cutoff),
("compress", detail_cutoff, 123, compressed_cutoff),
]