diff --git a/src/polygateway/client.py b/src/polygateway/client.py index d980274..49319d6 100644 --- a/src/polygateway/client.py +++ b/src/polygateway/client.py @@ -483,7 +483,10 @@ def _build_telemetry(settings: GatewaySettings) -> TelemetryRecorder | None: assert settings.telemetry_pg_dsn is not None # 内部不变量: _validate_telemetry 已保证 return PostgresRecorder( - settings.telemetry_pg_dsn, auto_migrate=settings.telemetry_auto_migrate + settings.telemetry_pg_dsn, + auto_migrate=settings.telemetry_auto_migrate, + pool_max=settings.telemetry_pg_pool_max, + write_timeout_s=settings.telemetry_pg_write_timeout_s, ) from polygateway.telemetry.sqlite import SQLiteRecorder diff --git a/src/polygateway/config.py b/src/polygateway/config.py index 19bb13f..ca04023 100644 --- a/src/polygateway/config.py +++ b/src/polygateway/config.py @@ -63,6 +63,13 @@ _SCHEMA_MODES = frozenset({"auto", "manual"}) _SCHEMA_MODE_KEY = "PGW_TELEMETRY_SCHEMA_MODE" # 遥测正文字符上限(issue #12);二态键,未设 = 不截断 _TEXT_CAP_KEY = "PGW_TELEMETRY_TEXT_CAP" +# 遥测池的资源占用与写入预算(issue #15);缺省只写在这里,recorder 侧是必填参数 +_POOL_MAX_KEY = "PGW_TELEMETRY_PG_POOL_MAX" +_WRITE_TIMEOUT_KEY = "PGW_TELEMETRY_PG_WRITE_TIMEOUT_S" +# 4 条 ≈ 32 行/秒(实测跨内网 RTT 123ms),够单 client 数十并发;闲时占 0 条 +_DEFAULT_PG_POOL_MAX = 4 +# 实测稳态写入 123ms、首次含建连 513ms;5s 宽松且**有界** +_DEFAULT_PG_WRITE_TIMEOUT_S = 5.0 _REDIS_DEPENDENT_BACKENDS = ("limiter_backend", "breaker_backend", "cache_backend") # 背压默认(M1 仅 poll 生效;CHS _BACKOFF_S=0.05 同源) _DEFAULT_STALL_WINDOW_S = 300.0 @@ -152,6 +159,13 @@ class GatewaySettings: # 既有下游正依赖这一行为。值域(> 0)由 `_validate_telemetry` 把关,直接构造、 # `dataclasses.replace` 与 env 三条路一并覆盖 telemetry_text_cap: int | None + # 遥测池对外声明的资源占用上限与整次写入的硬预算(issue #15)。库内每一处外部 + # 资源都按需建连,唯独遥测池此前预占 10 条(asyncpg 默认 `min_size`),共享实例 + # 余量紧张时先倒下的必然是它。这两个字段是库对自己占用的**显式表态**: + # 稳态并发上限 = `pool_max`,闲时 0 条;单次写入(准备+取连接+执行)≤ 预算。 + # 值域由 `_validate_telemetry` 把关,直接构造、`dataclasses.replace` 与 env 三条路一致 + telemetry_pg_pool_max: int + telemetry_pg_write_timeout_s: float redis_url: str | None pricing_path: str | None structured_max_retries: int @@ -248,6 +262,7 @@ class GatewaySettings: f"telemetry_text_cap({_TEXT_CAP_KEY})必须 > 0: {self.telemetry_text_cap};" "不截断请不设该键(None),0 只会让每条正文退化成一个省略标记" ) + self._validate_telemetry_pool() if self.telemetry_backend == "none" and self.telemetry_auto_migrate: object.__setattr__(self, "telemetry_auto_migrate", False) if self.telemetry_backend == "sqlite" and not self.telemetry_sqlite_path: @@ -266,6 +281,24 @@ class GatewaySettings: ) object.__setattr__(self, "telemetry_pg_dsn", stripped) + def _validate_telemetry_pool(self) -> None: + """遥测池两个标量的值域(issue #15);与 backend 无关,三条装配路一并覆盖。 + + 不按 `telemetry_backend == "postgres"` 才校验: 值域错就是错,提前拦住 + 比等到有人把 backend 切成 postgres 时才炸更接近"缺失关键配置直接报错"。 + 报错文本同时点字段名与 env 键名(两类调用方各看得懂自己那套)。 + """ + if self.telemetry_pg_pool_max < 1: + raise ValueError( + f"telemetry_pg_pool_max({_POOL_MAX_KEY})必须 >= 1: " + f"{self.telemetry_pg_pool_max};0 条上限等于永远取不到连接,遥测会全灭" + ) + if self.telemetry_pg_write_timeout_s <= 0: + raise ValueError( + f"telemetry_pg_write_timeout_s({_WRITE_TIMEOUT_KEY})必须 > 0: " + f"{self.telemetry_pg_write_timeout_s};预算 0 会让每一行当场超预算被丢弃" + ) + def _validate_lease(self) -> None: """调用超时须 ≤ permit 租约 TTL,防租约先于请求过期使并发超出配额。""" slowest = max(s.timeout_s for s in self.sources) @@ -486,6 +519,8 @@ def _load_pgw(env: Mapping[str, str]) -> dict[str, object]: "telemetry_pg_dsn": _load_pg_dsn(env) if telemetry_backend == "postgres" else None, "telemetry_auto_migrate": auto_migrate, "telemetry_text_cap": _load_text_cap(env), + "telemetry_pg_pool_max": _load_pool_max(env), + "telemetry_pg_write_timeout_s": _load_write_timeout(env), "redis_url": redis_url, "pricing_path": env.get("PGW_PRICING_PATH") or None, "structured_max_retries": _load_structured_retries(env), @@ -543,6 +578,43 @@ def _load_text_cap(env: Mapping[str, str]) -> int | None: return int(_cast(found[1], "int", found[0])) +def _load_pool_max(env: Mapping[str, str]) -> int: + """读 `PGW_TELEMETRY_PG_POOL_MAX`(issue #15);未设即缺省 4。 + + 与 `_load_text_cap` 同为二态键,只是"未设"落到一个具体缺省而非 None: + 池上限没有"不设上限"这一档——不表态就是继承第三方默认值,而那正是本 issue + 的病灶。值域(>= 1)留给构造期守卫,它同时覆盖直接构造与 `dataclasses.replace`。 + + Args: + env: 已合并的环境映射。 + + Returns: + 遥测池允许的最大连接数。 + """ + found = _first(env, _POOL_MAX_KEY) + if found is None: + return _DEFAULT_PG_POOL_MAX + return int(_cast(found[1], "int", found[0])) + + +def _load_write_timeout(env: Mapping[str, str]) -> float: + """读 `PGW_TELEMETRY_PG_WRITE_TIMEOUT_S`(issue #15);未设即缺省 5.0 秒。 + + 这个值同时是 connect、acquire 与整次写入的上界: 遥测是业务路径上的内联 + await,"不设预算"不是一个允许存在的档位(铁律"丢一条 < 拖垮调用")。 + + Args: + env: 已合并的环境映射。 + + Returns: + 单次遥测写入的硬预算(秒)。 + """ + found = _first(env, _WRITE_TIMEOUT_KEY) + if found is None: + return _DEFAULT_PG_WRITE_TIMEOUT_S + return float(_cast(found[1], "float", found[0])) + + def _strip_dsn_driver(dsn: str) -> str: """剥 SQLAlchemy 风格的 `+driver` 后缀(asyncpg 不认);已干净的原样返回。""" scheme, sep, rest = dsn.partition("://") diff --git a/src/polygateway/telemetry/postgres.py b/src/polygateway/telemetry/postgres.py index a94b1ac..cf8ed53 100644 --- a/src/polygateway/telemetry/postgres.py +++ b/src/polygateway/telemetry/postgres.py @@ -17,6 +17,7 @@ asyncpg 池自恢复;避免浸泡开头一次抖动导致后续全程失遥测) from __future__ import annotations import asyncio +import time from typing import TYPE_CHECKING from loguru import logger @@ -31,6 +32,8 @@ from polygateway.telemetry.schema import ( from polygateway.telemetry.status import TelemetryStatusTracker if TYPE_CHECKING: + from collections.abc import Callable + import asyncpg from polygateway.types import TelemetryStatus @@ -38,6 +41,12 @@ if TYPE_CHECKING: # 探测表是否存在;不需要任何权限,且与 INSERT 走同一套 search_path 解析 _TABLE_EXISTS = "SELECT to_regclass('llm_calls')" +# 归还连接的独立上限(issue #15)。**不**复用写入预算: 写入预算已经花在 +# acquire+execute 上,归还再给它一个同样大的额度,等于允许业务路径上的一次遥测 +# 写入吃掉 2 倍预算。归还是本地动作(reset 一次往返),1 秒足够;超时即断开, +# asyncpg 会在下次 acquire 时补一条新连接 +_RELEASE_TIMEOUT_S = 1.0 + # 探测现有列;尊重 search_path(to_regclass 按当前 search_path 解析) _EXISTING_COLUMNS = ( "SELECT attname FROM pg_attribute " @@ -48,7 +57,16 @@ _EXISTING_COLUMNS = ( class PostgresRecorder: """TelemetryRecorder 端口的 Postgres 实现;asyncpg 原生异步,无线程桥接。""" - def __init__(self, dsn: str, *, pool: asyncpg.Pool | None = None, auto_migrate: bool) -> None: + def __init__( + self, + dsn: str, + *, + pool: asyncpg.Pool | None = None, + auto_migrate: bool, + pool_max: int, + write_timeout_s: float, + now: Callable[[], float] = time.monotonic, + ) -> None: """记下装配参数(不连库);列与 INSERT 语句在首次准备期定型。 Args: @@ -59,6 +77,10 @@ class PostgresRecorder: 锁,会排在长事务后阻塞该表其后所有查询,而遥测是业务路径上的内联 await。keyword-only **必填**: 缺省规则只写在 config 一处,不与本类 签名漂移(设计 D-c)。 + pool_max: 自建池的连接数上限(issue #15)。稳态吞吐 ≈ `pool_max / RTT`。 + 与 `auto_migrate` 同一纪律: 必填,缺省只写在 config 一处。 + write_timeout_s: 单次写入的硬预算,同时用作 connect 与 acquire 的上限。 + now: 单调时钟,注入给降级 tracker(测试可推进冷却与节流窗口)。 """ try: import asyncpg # noqa: F401 - 仅探测 extra 是否安装 @@ -70,6 +92,8 @@ class PostgresRecorder: self._pool: asyncpg.Pool | None = pool self._external_pool = pool is not None self._auto_migrate = auto_migrate + self._pool_max = pool_max + self._write_timeout_s = write_timeout_s # 先按全量列定型: 准备期探测失败时保守沿用全量(今天的行为) self._columns: tuple[str, ...] = COLUMNS self._insert = insert_sql("postgres", COLUMNS) @@ -77,7 +101,7 @@ class PostgresRecorder: self._failed = False # 结构性降级标志: 置位后所有写入短路 # 降级的可编程出口与节流日志;`_failed` 与它并存是 issue #15 的过渡态, # 判据改造(冷却自愈)落地时状态收归 tracker 一处 - self._status = TelemetryStatusTracker(backend="postgres") + self._status = TelemetryStatusTracker(backend="postgres", now=now) self._init_lock = asyncio.Lock() @property @@ -102,13 +126,27 @@ class PostgresRecorder: return await self._prepare_schema(pool) async def _open_pool(self) -> asyncpg.Pool | None: - """建池;失败即永久降级(唯一一处「无条件判死」)。""" + """建池;失败即永久降级(唯一一处「无条件判死」)。 + + **池的资源占用由本库显式声明**(issue #15): `min_size=0` 的语义是"不预 + 连接"(asyncpg `pool.py:457` 为 0 时只造 holder 对象,一条连接都不连), + 建池因此从"要么拿到 10 条、要么失败"的重资源动作变成零成本、不触库的 + 动作;连接失败自然落到 acquire 那条本来就正确的"丢一行、池自恢复"路径。 + `max_size` 是库对自己占用的表态——继承第三方默认值等于不表态(P4),而 + 那正是共享实例余量紧张时先倒下的原因。 + """ if self._pool is not None: return self._pool try: import asyncpg - self._pool = await asyncpg.create_pool(self._dsn, timeout=10) + self._pool = await asyncpg.create_pool( + self._dsn, + min_size=0, + max_size=self._pool_max, + timeout=self._write_timeout_s, + command_timeout=self._write_timeout_s, + ) except asyncio.CancelledError: raise except Exception as exc: @@ -120,10 +158,17 @@ class PostgresRecorder: return self._pool async def _prepare_schema(self, pool: asyncpg.Pool) -> asyncpg.Pool | None: - """备好表并交回可用的池;瞬时失败只跳过本次,确定写不进去才判死。""" + """备好表并交回可用的池;瞬时失败只跳过本次,确定写不进去才判死。 + + 取连接走显式 acquire/release(理由见 `_release`): 准备期同样跑在调用方的 + 写入预算里,`async with` 那条路的归还会把真实上界撑到 ≈2× 预算。 + """ try: - async with pool.acquire() as conn: + conn = await pool.acquire(timeout=self._write_timeout_s) + try: columns = await self._prepare_table(conn) + finally: + await self._release(pool, conn) except asyncio.CancelledError: raise except Exception as exc: @@ -239,7 +284,33 @@ class PostgresRecorder: logger.warning("Postgres 遥测补列失败(写入将逐行降级): {}", exc) async def record_llm_call(self, **fields: object) -> None: - """写一行遥测;单条失败逐条 warning 丢弃(两级降级之二),绝不冒泡。 + """写一行遥测;整次写入受硬预算约束,失败逐条丢弃(两级降级之二),绝不冒泡。 + + **硬预算**(issue #15): 准备 + 取连接 + 执行合计不得超过 `write_timeout_s`。 + 这把"遥测绝不拖垮业务"从"靠各处 timeout 参数凑"变成一条可陈述、可测试的 + 保证——此前 `pool.acquire()` 无超时(asyncpg 缺省 `timeout=None` = 无限等), + 池满时会无限期挂在业务路径上。 + + 外部取消照常穿透: `asyncio.timeout` 只把**自己**触发的 cancel 转成 + TimeoutError,故 `CancelledError` 分支必须排在最前且原样 re-raise(铁律)。 + """ + try: + async with asyncio.timeout(self._write_timeout_s): + await self._write_row(fields) + except asyncio.CancelledError: + raise + except TimeoutError: + logger.warning( + "Postgres 遥测写入超预算 {}s(丢弃该行);后端慢不得拖垮业务调用", + self._write_timeout_s, + ) + self._status.record_drop("写入超预算") + except Exception as exc: + # 遥测铁律: 丢一条 < 拖垮调用;仅记 warning(非 pass),池自恢复 + logger.warning("Postgres 遥测写入失败(丢弃该行): {}", exc) + + async def _write_row(self, fields: dict[str, object]) -> None: + """预算内的写入本体: 准备 → 取连接 → 执行 → 归还。 取值按 `self._columns`(manual 档可能已被裁剪),与 `self._insert` 的 占位符同序——两者必须一起改,分开改就是把值写进错位的列。 @@ -250,14 +321,42 @@ class PostgresRecorder: self._status.record_drop("遥测已降级") return row = tuple(fields[col] for col in self._columns) + conn = await pool.acquire(timeout=self._write_timeout_s) try: - async with pool.acquire() as conn: - await conn.execute(self._insert, *row) + await conn.execute(self._insert, *row) + finally: + await self._release(pool, conn) + + async def _release(self, pool: asyncpg.Pool, conn: object) -> None: + """归还连接;归还路径独立有界,失败即断开(下次 acquire 会补一条新的)。 + + **不用 `async with pool.acquire()`**(设计 §3.1,已核实): asyncpg 的 + `Pool.release` 是 `await asyncio.shield(ch.release(timeout))`,且那个 + timeout 默认复用 acquire 时记录的 `ch._timeout`(`pool.py:886-889, + 930-937`)。写入预算到期时 cancel 在 execute 处抛出,异常传播中执行 + `__aexit__`,此时没有新的 cancel 投递——那次 shielded release 会**正常 + 等到完成**,业务路径的真实上界因此变成 ≈2 × 预算。显式归还才能给它一个 + 独立的小上限,承诺才精确成立: 主写入尝试 ≤ 预算,归还路径独立有界。 + """ + try: + await pool.release(conn, timeout=_RELEASE_TIMEOUT_S) except asyncio.CancelledError: raise except Exception as exc: - # 遥测铁律: 丢一条 < 拖垮调用;仅记 warning(非 pass),池自恢复 - logger.warning("Postgres 遥测写入失败(丢弃该行): {}", exc) + # 含 TimeoutError: 归还超时与归还出错的处置相同——断开而不是留一条 + # 状态不明的连接在池里(asyncpg 的 reset 失败路径也是这么做的) + logger.warning("Postgres 遥测连接归还失败(强制断开): {}", exc) + self._terminate(conn) + + @staticmethod + def _terminate(conn: object) -> None: + """强制断开一条连接;断开本身再失败也只记 warning(遥测绝不冒泡)。""" + try: + conn.terminate() # type: ignore[attr-defined] + except asyncio.CancelledError: + raise + except Exception as exc: + logger.warning("Postgres 遥测连接断开失败(交给池自行回收): {}", exc) async def aclose(self) -> None: """幂等关闭自建池;注入的池归注入方管理。""" diff --git a/tests/integration/test_postgres_telemetry.py b/tests/integration/test_postgres_telemetry.py index e301c2d..49c8551 100644 --- a/tests/integration/test_postgres_telemetry.py +++ b/tests/integration/test_postgres_telemetry.py @@ -130,6 +130,27 @@ async def _record_minimal( return fields +# 集成用例统一的池上限与写入预算(issue #15;两者是 recorder 的必填 keyword-only)。 +# 池上限取 config 的生产缺省(4),让本文件跑的就是下游真实会跑的那个形状。 +# +# 预算却**远比生产的 5s 宽**,这不是抄错: `test_concurrent_writes_all_land` 一次 +# 发 50 行,50 行共享 4 条连接,实测跨内网 RTT 123ms 下整批约 3.2s——而那 50 个 +# `record_llm_call` 的预算是**同时**起算的,批越慢离预算越近。这个实例被多项目 +# 共用,别人的一次负载尖峰就能让批耗时翻几倍,于是"丢行"变成掷硬币(pool_max=2 +# 时实测批耗时 5.3s/15s 预算,已经在全套件里红过一次)。给它 60s 是把余量拉到 +# 近 20 倍,让这个用例只在真出 bug 时红(CLAUDE.md §4.6: 重跑一次就绿的测试是 +# 信号污染源)。突发排队本身超预算即丢行是设计上的既定取舍(设计 §6),不在此改。 +_POOL_MAX = 4 +_WRITE_TIMEOUT_S = 60.0 + + +def _recorder(dsn: str, *, auto_migrate: bool) -> PostgresRecorder: + """本文件唯一的 recorder 构造点: 池参数只写一遍,免得 16 处各抄一份。""" + return PostgresRecorder( + dsn, auto_migrate=auto_migrate, pool_max=_POOL_MAX, write_timeout_s=_WRITE_TIMEOUT_S + ) + + async def _fetch(dsn: str, sql: str, *args): import asyncpg @@ -205,7 +226,7 @@ class TestObservabilityColumns: """issue #3: 两列写入可回读,且已存在的 18 列旧表会被自动补列。""" async def test_values_round_trip(self, dsn): - recorder = PostgresRecorder(dsn, auto_migrate=True) + recorder = _recorder(dsn, auto_migrate=True) try: await _record_minimal(recorder, call_id=_cid("hit"), cached_prompt_tokens=64) await _record_minimal(recorder, call_id=_cid("zero"), cached_prompt_tokens=0) @@ -233,7 +254,7 @@ class TestObservabilityColumns: async def test_legacy_table_is_upgraded_in_place(self, legacy_schema): """18 列旧表不补列的话,每行写入都会被逐行 warning 丢弃(遥测静默全失)。""" schema_dsn, schema = legacy_schema - recorder = PostgresRecorder(schema_dsn, auto_migrate=True) + recorder = _recorder(schema_dsn, auto_migrate=True) try: await _record_minimal( recorder, call_id=_cid("legacy"), cached_prompt_tokens=7, model_reported="m-real" @@ -258,7 +279,7 @@ class TestObservabilityColumns: class TestSchema: async def test_schema_has_frozen_columns_in_order(self, dsn): - recorder = PostgresRecorder(dsn, auto_migrate=True) + recorder = _recorder(dsn, auto_migrate=True) try: await _record_minimal(recorder) rows = await _fetch( @@ -271,7 +292,7 @@ class TestSchema: await recorder.aclose() async def test_call_id_idempotent(self, dsn): - recorder = PostgresRecorder(dsn, auto_migrate=True) + recorder = _recorder(dsn, auto_migrate=True) try: await _record_minimal(recorder, call_id=_cid("dup")) await _record_minimal(recorder, call_id=_cid("dup"), response="second") @@ -283,7 +304,7 @@ class TestSchema: await recorder.aclose() async def test_concurrent_writes_all_land(self, dsn): - recorder = PostgresRecorder(dsn, auto_migrate=True) + recorder = _recorder(dsn, auto_migrate=True) try: await asyncio.gather( *(_record_minimal(recorder, call_id=_cid(f"c{i}")) for i in range(50)) @@ -301,14 +322,14 @@ class TestSchema: class TestDegradation: async def test_unreachable_server_degrades_silently(self): """结构性失败(建池不通)→ warning 一次后永久降级,业务零感知。""" - recorder = PostgresRecorder("postgresql://u:p@127.0.0.1:1/x", auto_migrate=True) + recorder = _recorder("postgresql://u:p@127.0.0.1:1/x", auto_migrate=True) await _record_minimal(recorder) # 不抛 await _record_minimal(recorder, call_id=_cid("c2")) # 已降级短路,同样不抛 await recorder.aclose() async def test_row_failure_does_not_poison_later_rows(self, dsn): """运行时单条写失败(NUL 字节文本被 PG 拒)→ 丢该行,后续行照常落库。""" - recorder = PostgresRecorder(dsn, auto_migrate=True) + recorder = _recorder(dsn, auto_migrate=True) try: await _record_minimal(recorder, call_id=_cid("bad"), response="nul\x00byte") await _record_minimal(recorder, call_id=_cid("good")) @@ -322,7 +343,7 @@ class TestDegradation: await recorder.aclose() async def test_aclose_idempotent(self, dsn): - recorder = PostgresRecorder(dsn, auto_migrate=True) + recorder = _recorder(dsn, auto_migrate=True) await _record_minimal(recorder) await recorder.aclose() await recorder.aclose() @@ -394,7 +415,7 @@ class TestLeastPrivilegeDeployment: async def test_records_land_without_schema_create_privilege(self, least_privilege_dsn): """修复前: 建表被拒 → _failed → 整个进程一条不落(下游 150 次调用全丢)。""" low_dsn, schema = least_privilege_dsn - recorder = PostgresRecorder(low_dsn, auto_migrate=True) + recorder = _recorder(low_dsn, auto_migrate=True) try: await _record_minimal(recorder, call_id=_cid("lp1")) await _record_minimal(recorder, call_id=_cid("lp2"), cost=1.5) @@ -563,7 +584,7 @@ class TestCallerDimensionsAcceptance: async def test_fresh_schema_round_trips_the_dimensions(self, fresh_schema): """新建库: 列齐全,且维度值原样读回——只验列存在会漏掉写错列位的错。""" fresh_dsn, schema = fresh_schema - recorder = PostgresRecorder(fresh_dsn, auto_migrate=True) + recorder = _recorder(fresh_dsn, auto_migrate=True) try: await _record_minimal( recorder, call_id=_cid("dim"), tenant_id="tenant-a", meta='{"batch": "b7"}' @@ -597,7 +618,7 @@ class TestCallerDimensionsAcceptance: 审计出来,历史欠账是可见、可量化、可补录的。 """ schema_dsn, schema = pre_tenant_schema - recorder = PostgresRecorder(schema_dsn, auto_migrate=True) + recorder = _recorder(schema_dsn, auto_migrate=True) try: await _record_minimal( recorder, call_id=_cid("new"), tenant_id="tenant-a", meta='{"k": 1}' @@ -649,7 +670,7 @@ class TestCallerDimensionsAcceptance: 置 `_failed` 会让整个进程从此一条遥测都不写(比逐行丢弃严重得多), 且一旦 DBA 补上列也不会自愈——必须等重启。 """ - recorder = PostgresRecorder(least_privilege_pre_tenant_dsn, auto_migrate=True) + recorder = _recorder(least_privilege_pre_tenant_dsn, auto_migrate=True) try: await _record_minimal(recorder, call_id=_cid("lpp1")) # 不得抛 assert recorder._failed is False @@ -745,7 +766,7 @@ class TestConflictTargetFreeInsert: 断言"无写入失败 warning"是为了区分"冲突被忽略"与"整条被 PG 拒收"。 """ fresh_dsn, _ = fresh_schema - recorder = PostgresRecorder(fresh_dsn, auto_migrate=True) + recorder = _recorder(fresh_dsn, auto_migrate=True) try: await _record_minimal(recorder, call_id=_cid("nodup")) await _record_minimal(recorder, call_id=_cid("nodup"), response="second") @@ -766,7 +787,7 @@ class TestConflictTargetFreeInsert: 遥测全线写不进去却一声不吭,只能靠"读不回来"暴露。 """ part_dsn, _ = partitioned_schema - recorder = PostgresRecorder(part_dsn, auto_migrate=True) + recorder = _recorder(part_dsn, auto_migrate=True) try: await _record_minimal(recorder, call_id=_cid("part"), tenant_id="tenant-p") assert [m for m in captured_warnings if "写入失败" in m] == [] @@ -797,7 +818,7 @@ class TestManualSchemaModeAcceptance: 同一张表、同一份负载,只有 `auto_migrate` 不同,列数就必须是 23 与 25 之别。 """ schema_dsn, schema = pre_tenant_schema - recorder = PostgresRecorder(schema_dsn, auto_migrate=False) + recorder = _recorder(schema_dsn, auto_migrate=False) try: recorded = await _record_minimal( recorder, call_id=_cid("man"), tenant_id="tenant-a", meta='{"k": 1}' @@ -837,7 +858,7 @@ class TestManualSchemaModeAcceptance: 消灭的噪声。manual 档下 ALTER 压根不发,取而代之的是一条点名缺列并附可直接 执行的 ALTER 的提示,而遥测照常落库。 """ - recorder = PostgresRecorder(least_privilege_pre_tenant_dsn, auto_migrate=False) + recorder = _recorder(least_privilege_pre_tenant_dsn, auto_migrate=False) try: recorded = await _record_minimal( recorder, call_id=_cid("manlp1"), tenant_id="tenant-b", meta='{"k": 2}' diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 82aeb8d..e29bc4d 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -425,6 +425,73 @@ class TestTelemetryTextCap: GatewaySettings.from_env("LLM", env=_env(PGW_TELEMETRY_TEXT_CAP="2k")) +class TestTelemetryPoolKeys: + """`PGW_TELEMETRY_PG_POOL_MAX` / `PGW_TELEMETRY_PG_WRITE_TIMEOUT_S`(issue #15)。 + + 两键都带 `PG` 前缀,与 `PGW_TELEMETRY_PG_DSN` 一致: SQLite 侧没有池、也没有 + 等价的写入预算旋钮,这个不对称是已知且有理由的。缺省值(4 / 5.0)只写在 + config 一处——recorder 的两个同名参数是必填 keyword-only,不许各带一份缺省。 + """ + + def _pg_env(self, **overrides): + return _env( + PGW_TELEMETRY_BACKEND="postgres", + PGW_TELEMETRY_PG_DSN="postgresql://u:p@h:5432/polygateway", + **overrides, + ) + + def test_unset_keys_fall_back_to_the_documented_defaults(self): + s = GatewaySettings.from_env("LLM", env=self._pg_env()) + assert s.telemetry_pg_pool_max == 4 + assert s.telemetry_pg_write_timeout_s == 5.0 + + def test_values_parsed_from_env(self): + s = GatewaySettings.from_env( + "LLM", + env=self._pg_env(PGW_TELEMETRY_PG_POOL_MAX="8", PGW_TELEMETRY_PG_WRITE_TIMEOUT_S="1.5"), + ) + assert s.telemetry_pg_pool_max == 8 + assert s.telemetry_pg_write_timeout_s == 1.5 + + @pytest.mark.parametrize("raw", ["0", "-1"]) + def test_non_positive_pool_max_rejected_naming_the_env_key(self, raw): + """池上限 0 = 永远拿不到连接(遥测全灭),负数无意义。""" + with pytest.raises(ValueError, match="PGW_TELEMETRY_PG_POOL_MAX"): + GatewaySettings.from_env("LLM", env=self._pg_env(PGW_TELEMETRY_PG_POOL_MAX=raw)) + + @pytest.mark.parametrize("raw", ["0", "-1"]) + def test_non_positive_write_timeout_rejected_naming_the_env_key(self, raw): + """预算 0 = 每一行都当场超预算;不设预算不是这个键的写法。""" + with pytest.raises(ValueError, match="PGW_TELEMETRY_PG_WRITE_TIMEOUT_S"): + GatewaySettings.from_env("LLM", env=self._pg_env(PGW_TELEMETRY_PG_WRITE_TIMEOUT_S=raw)) + + def test_non_numeric_rejected_naming_the_env_key(self): + with pytest.raises(ValueError, match="PGW_TELEMETRY_PG_POOL_MAX"): + GatewaySettings.from_env("LLM", env=self._pg_env(PGW_TELEMETRY_PG_POOL_MAX="many")) + + @pytest.mark.parametrize( + ("field", "value"), + [("telemetry_pg_pool_max", 0), ("telemetry_pg_write_timeout_s", 0.0)], + ) + def test_direct_construction_and_replace_are_validated_too(self, field, value): + """env 路只覆盖 from_env;直接构造与 replace 是同等官方的装配路(与 text_cap 同款)。""" + base = GatewaySettings.from_env("LLM", env=_env()) + with pytest.raises(ValueError, match=field): + dataclasses.replace(base, **{field: value}) + + def test_values_reach_the_recorder(self): + """配置到 recorder 之间不得断链——两个键唯一的作用就是抵达那里。""" + from polygateway.client import _build_telemetry + + settings = GatewaySettings.from_env( + "LLM", + env=self._pg_env(PGW_TELEMETRY_PG_POOL_MAX="7", PGW_TELEMETRY_PG_WRITE_TIMEOUT_S="2.5"), + ) + recorder = _build_telemetry(settings) + assert recorder._pool_max == 7 + assert recorder._write_timeout_s == 2.5 + + class TestOcrSettings: """M3 OcrSettings(设计 §3.4): 复用 GatewaySettings,无 OCR 专用键。""" diff --git a/tests/unit/test_telemetry.py b/tests/unit/test_telemetry.py index 38a687d..b594633 100644 --- a/tests/unit/test_telemetry.py +++ b/tests/unit/test_telemetry.py @@ -706,6 +706,13 @@ class TestSQLiteSchemaMode: SQLiteRecorder(tmp_path / "t.db") # type: ignore[call-arg] +# 假池用例的池上限与写入预算: 两者都是必填 keyword-only(缺省只写在 config 一处), +# 本文件统一取这一份,免得每个 helper 各写一个数字 +_TEST_POOL_MAX = 2 +_TEST_WRITE_TIMEOUT_S = 5.0 +_PG_DSN = "postgresql://u:p@h:5432/polygateway" + + class _FakePgConn: """记录执行过的语句;可让 ALTER/CREATE/探测抛错以模拟权限不足与抖动。 @@ -720,15 +727,20 @@ class _FakePgConn: fail_alter: bool = False, fail_create: bool = False, probe_errors: int = 0, + hang_insert: bool = False, ): self.existing = existing self.fail_alter = fail_alter self.fail_create = fail_create self.probe_errors = probe_errors + # 只挂 INSERT: 准备期照常完成,挂住的才是业务路径上那次内联 await + self.hang_insert = hang_insert self.statements: list[str] = [] async def execute(self, sql, *args): self.statements.append(sql) + if sql.startswith("INSERT INTO") and self.hang_insert: + await asyncio.sleep(3600) if sql.startswith("ALTER TABLE") and self.fail_alter: raise RuntimeError("must be owner of table llm_calls") if sql.lstrip().startswith("CREATE TABLE"): @@ -749,20 +761,33 @@ class _FakePgConn: class _FakePgPool: - def __init__(self, conn): + """假池: 记 acquire/release 的配对次数与实参 timeout(issue #15 T3)。 + + 形状跟着被测代码走: recorder 改用**显式** `acquire(timeout=)` / + `release(conn, timeout=)`,不再用 `async with pool.acquire()`(那条路 + 的 shielded release 会把写入的真实上界撑成 ≈2× 预算,设计 §3.1), + 故这里也不再提供上下文管理器。 + """ + + def __init__(self, conn, *, hang_acquire: bool = False): self._conn = conn + self.hang_acquire = hang_acquire + self.acquired = 0 + self.released = 0 + self.acquire_timeouts: list[object] = [] + self.release_timeouts: list[object] = [] - def acquire(self): - conn = self._conn + async def acquire(self, *, timeout=None): + self.acquire_timeouts.append(timeout) + if self.hang_acquire: + await asyncio.sleep(3600) + self.acquired += 1 + return self._conn - class _Ctx: - async def __aenter__(self): - return conn - - async def __aexit__(self, *exc): - return False - - return _Ctx() + async def release(self, conn, *, timeout=None): + assert conn is self._conn + self.release_timeouts.append(timeout) + self.released += 1 class TestPostgresBackfillDiscipline: @@ -785,7 +810,11 @@ class TestPostgresBackfillDiscipline: from polygateway.telemetry.postgres import PostgresRecorder return PostgresRecorder( - "postgresql://u:p@h:5432/polygateway", pool=_FakePgPool(conn), auto_migrate=True + "postgresql://u:p@h:5432/polygateway", + pool=_FakePgPool(conn), + auto_migrate=True, + pool_max=_TEST_POOL_MAX, + write_timeout_s=_TEST_WRITE_TIMEOUT_S, ) async def test_alter_failure_does_not_disable_the_recorder(self): @@ -841,7 +870,11 @@ class TestPostgresTableProbe: from polygateway.telemetry.postgres import PostgresRecorder return PostgresRecorder( - "postgresql://u:p@h:5432/polygateway", pool=_FakePgPool(conn), auto_migrate=True + "postgresql://u:p@h:5432/polygateway", + pool=_FakePgPool(conn), + auto_migrate=True, + pool_max=_TEST_POOL_MAX, + write_timeout_s=_TEST_WRITE_TIMEOUT_S, ) def _created(self, conn): @@ -903,7 +936,11 @@ class TestPostgresSchemaMode: from polygateway.telemetry.postgres import PostgresRecorder return PostgresRecorder( - "postgresql://u:p@h:5432/polygateway", pool=_FakePgPool(conn), auto_migrate=auto_migrate + "postgresql://u:p@h:5432/polygateway", + pool=_FakePgPool(conn), + auto_migrate=auto_migrate, + pool_max=_TEST_POOL_MAX, + write_timeout_s=_TEST_WRITE_TIMEOUT_S, ) async def test_manual_mode_trims_the_insert_instead_of_altering(self, captured_warnings): @@ -1852,7 +1889,11 @@ class TestPostgresStatusVisibility: from polygateway.telemetry.postgres import PostgresRecorder return PostgresRecorder( - "postgresql://u:p@h:5432/polygateway", pool=_FakePgPool(conn), auto_migrate=True + "postgresql://u:p@h:5432/polygateway", + pool=_FakePgPool(conn), + auto_migrate=True, + pool_max=_TEST_POOL_MAX, + write_timeout_s=_TEST_WRITE_TIMEOUT_S, ) async def test_unusable_table_shows_up_in_the_status(self, captured_warnings): @@ -1869,3 +1910,107 @@ class TestPostgresStatusVisibility: await _record_minimal(recorder) status = recorder.telemetry_status assert status.degraded is False and status.dropped_rows == 0 + + +class TestPostgresPoolResourceSemantics: + """issue #15 A 组: 库必须自己声明池的资源占用,并给写入一个硬预算。 + + 建池这条路在本 issue 之前**零测试覆盖**(全部 PG 用例都经 `pool=` 注入, + 走的是外部池分支),`min_size=10` 因此潜伏至今: 4 个 client × 10 = 40 条 + 常驻连接专用于写遥测,共享实例余量不足时先倒下的必然是它。 + """ + + def _recorder(self, pool, **overrides): + from polygateway.telemetry.postgres import PostgresRecorder + + kwargs: dict[str, object] = { + "auto_migrate": True, + "pool_max": _TEST_POOL_MAX, + "write_timeout_s": _TEST_WRITE_TIMEOUT_S, + } + kwargs.update(overrides) + return PostgresRecorder(_PG_DSN, pool=pool, **kwargs) + + async def test_pool_is_created_without_preconnecting(self, monkeypatch): + """**主回归钉子**: `min_size=0` 且 `max_size` 取配置值。 + + `min_size` 的语义是"预连接"而非"下限"(asyncpg `pool.py:457` 为 0 时 + 一条连接都不连),故它是"建池要么全有要么全无"这个脆点的唯一来源。 + 继承第三方默认值等于库对自己的资源占用不表态(P4),本条防的就是回归。 + """ + import asyncpg + + from polygateway.telemetry.postgres import PostgresRecorder + + captured: dict[str, object] = {} + pool = _FakePgPool(_FakePgConn(list(_EXPECTED_COLUMNS))) + + async def fake_create_pool(dsn, **kwargs): + captured["dsn"] = dsn + captured.update(kwargs) + return pool + + monkeypatch.setattr(asyncpg, "create_pool", fake_create_pool) + recorder = PostgresRecorder(_PG_DSN, auto_migrate=True, pool_max=3, write_timeout_s=2.5) + await _record_minimal(recorder) + + assert captured["min_size"] == 0 + assert captured["max_size"] == 3 + # connect 与单条语句都在同一份写入预算内,不留继承来的 10s 默认值 + assert captured["timeout"] == 2.5 + assert captured["command_timeout"] == 2.5 + + async def test_acquire_gets_an_explicit_timeout(self): + """`pool.acquire()` 无参 = 无限等(asyncpg 缺省 `timeout=None`)。 + + 池满时那是挂在业务路径上的无限期 await,`max_size` 收到个位数后必现。 + """ + pool = _FakePgPool(_FakePgConn(list(_EXPECTED_COLUMNS))) + await _record_minimal(self._recorder(pool)) + assert pool.acquire_timeouts # 准备期与写入期各一次 + assert all(t == _TEST_WRITE_TIMEOUT_S for t in pool.acquire_timeouts) + + async def test_write_budget_drops_the_row_instead_of_blocking_the_caller( + self, captured_warnings + ): + """整次写入有硬预算: 后端挂住时丢一行,绝不把业务调用拖在那里。""" + pool = _FakePgPool(_FakePgConn(list(_EXPECTED_COLUMNS), hang_insert=True)) + recorder = self._recorder(pool, write_timeout_s=0.05) + loop = asyncio.get_running_loop() + started = loop.time() + # 挂死就当场红,而不是把整个套件拖到 CI 超时 + await asyncio.wait_for(_record_minimal(recorder), timeout=5) + assert loop.time() - started < 1.0 + assert any("预算" in m for m in captured_warnings) + assert recorder.telemetry_status.dropped_rows == 1 + + async def test_release_is_paired_even_when_the_budget_fires(self): + """预算取消发生在 execute 上,连接照样要还回去——否则池被慢查询吃干。""" + pool = _FakePgPool(_FakePgConn(list(_EXPECTED_COLUMNS), hang_insert=True)) + recorder = self._recorder(pool, write_timeout_s=0.05) + await asyncio.wait_for(_record_minimal(recorder), timeout=5) + assert pool.acquired == 2 # 准备期一次 + 写入一次 + assert pool.released == pool.acquired + # 归还有独立的小上限: 复用写入预算就等于允许再等一个预算(设计 §3.1) + assert all(t is not None and t < _TEST_WRITE_TIMEOUT_S for t in pool.release_timeouts) + + async def test_acquire_timeout_drops_the_row_without_leaking(self, captured_warnings): + """取连接本身挂住时同样丢行;没拿到的连接不许伪造一次 release。""" + pool = _FakePgPool(_FakePgConn(list(_EXPECTED_COLUMNS)), hang_acquire=True) + recorder = self._recorder(pool, write_timeout_s=0.05) + await asyncio.wait_for(_record_minimal(recorder), timeout=5) + assert pool.acquired == 0 and pool.released == 0 + assert captured_warnings + + async def test_external_cancellation_is_not_swallowed_as_a_timeout(self): + """铁律"取消可穿透": `asyncio.timeout` 只把**自己**触发的 cancel 转成 + TimeoutError,外部取消必须照常以 CancelledError 冒出去。 + """ + pool = _FakePgPool(_FakePgConn(list(_EXPECTED_COLUMNS), hang_insert=True)) + recorder = self._recorder(pool, write_timeout_s=30.0) + task = asyncio.create_task(_record_minimal(recorder)) + await asyncio.sleep(0.05) # 让它跑到挂住的那次 INSERT + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert pool.released == pool.acquired # 取消路径上也不许泄漏连接