iomgaa/PolyGateway
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases 9 Wiki Activity

feat: make the telemetry pool declare what it costs

Browse Source 
The pool was the only external resource in the library that pre-allocated:
asyncpg's default min_size=10 turned pool creation into an all-or-nothing
action, so on a shared instance running low on connection budget the first
thing to fall over was the one component that must not fail silently
(4 clients x 10 = 40 idle connections just to write telemetry).

min_size=0 means "do not pre-connect" - asyncpg only builds holders - so
pool creation becomes free and never touches the database; connection
failures then land on acquire, the path that already drops one row and lets
the pool recover. max_size and the write budget become the library's
explicit statement about its own footprint, configurable through two new
keys whose defaults live in config alone (the recorder parameters are
required keyword-only, same discipline as auto_migrate).

The whole write - prepare, acquire, execute - now runs inside one
asyncio.timeout: acquire used to have no timeout at all, so a full pool
would hang forever on the caller's path. Release is explicit rather than
`async with`, because asyncpg shields release and reuses the acquire
timeout, which would let a single telemetry write consume twice the budget.
This commit is contained in:
iomgaa
2026-08-24 09:32:35 -04:00
parent f958138e83
commit 84c2cc11a4
6 changed files with 450 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/polygateway/telemetry/postgres.py
+110 -11
View File
@@ -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:
"""幂等关闭自建池;注入的池归注入方管理。"""