本文详解如何通过自定义 __await__ 方法安全包装异步上下文管理器（如 asyncpg.pool.PoolAcquireContext），确保日志计数器仅在连接真正获取/释放后更新，避免因 await 被取消或异常导致状态不一致。
本文详解如何通过自定义 `__await__` 方法安全包装异步上下文管理器（如 `asyncpg.pool.poolacquirecontext`），确保日志计数器仅在连接真正获取/释放后更新，避免因 `await` 被取消或异常导致状态不一致。
在异步数据库连接池监控场景中，精确跟踪连接的实际获取（acquire）与释放（release）时机至关重要。常见误区是直接在 __await__ 中提前调用 acquire()，但这会导致计数器在协程尚未完成、甚至被取消时就被错误递增——因为 __await__ 仅返回一个可等待对象（generator），并不保证其后续一定会被 await 执行完毕。
正确的做法是：将日志逻辑嵌入协程执行流本身，利用 Python 协程生成器的 yield from 机制，使 acquire() 和 release() 分别在协程开始和结束（包括异常退出）时精准触发。
以下是关键实现：
class LoggingPoolAcquireContext:
def __init__(self, pool_acquire_context: asyncpg.pool.PoolAcquireContext, cnt_logger: CntPoolLogger):
self.pool_acquire_context = pool_acquire_context
self.cnt_logger = cnt_logger
def __await__(self):
# ✅ 正确：acquire 在 await 开始时执行（即连接即将被获取）
self.cnt_logger.acquire()
try:
# ✅ yield from 将控制权交还给底层协程，并透传其执行结果
# ✅ 协程完成后（无论成功或异常），finally 块确保 release 执行
return (yield from self.pool_acquire_context.__await__())
except BaseException:
# ⚠️ 注意：若 acquire 成功但后续 await 失败（如网络中断），连接可能未被实际使用
# 此处仍需 release，否则连接泄漏；业务层应结合连接健康检查进一步处理
self.cnt_logger.release()
raise
finally:
# ✅ 确保 release 总在协程生命周期结束时调用
# （注意：此处不能无条件 release！见下方说明）
pass
# ✅ 更健壮的实现：仅在 __aenter__ 成功后才 release，避免重复释放
async def __aenter__(self):
conn = await self.pool_acquire_context.__aenter__()
self.cnt_logger.acquire() # 实际获取连接后计数
return conn
async def __aexit__(self, exc_type, exc_val, exc_tb):
try:
await self.pool_acquire_context.__aexit__(exc_type, exc_val, exc_tb)
finally:
self.cnt_logger.release() # 确保连接归还后计数⚠️ 重要注意事项：
- __await__ 中的 finally 块不能直接调用 release() —— 因为 __await__ 只负责协程迭代，不等价于资源生命周期管理；它可能被多次 await 或中途丢弃（如被 asyncio.wait_for 取消），导致 release() 被误触发。
- 推荐优先使用 __aenter__/__aexit__ 组合（即 async with 语义），因其语义明确、生命周期可控，且与 asyncpg 原生行为完全兼容。
- 若必须支持裸 await（如 conn = await pool.acquire()），则 __await__ 应仅包装 __aenter__ 的调用，并确保 release() 由配套的 __aexit__ 或显式清理逻辑承担，而非 __await__ 的 finally。
- CntPoolLogger 的 __del__ 不可靠（无法保证及时调用），应改用上下文管理器或显式 close() + asyncio.atexit 清理日志文件。
✅ 总结：
包装协程的 __await__ 并非简单“返回另一个 __await__”，而是要理解其本质是生成器协议。通过 yield from 委托执行，并配合 try/except/finally 控制副作用时机，才能实现原子性日志记录。但在生产环境中，更推荐基于 async with 的上下文管理方案，兼顾可读性、健壮性与标准兼容性。
