fix(db): invalidate pool connection when cancelled close fails

Under high-concurrency MySQL load, uvicorn cancels request tasks when
clients disconnect.  If cancellation lands mid-query, session.close()
tries to ROLLBACK on a connection that aiomysql has already marked as
closed — raising InterfaceError("Cancelled during execution") and
leaving the connection checked-out until GC, which the pool then
warns about as a 'non-checked-in connection'.

The old fallback tried sync.rollback() + sync.close(), but those still
go through the async driver and fail the same way on a dead connection.
Replace them with session.sync_session.invalidate(), which just flips
the pool's internal record — no I/O, so it can't be cancelled — and
tells the pool to drop the connection immediately instead of waiting
for garbage collection.
This commit is contained in:
2026-04-17 21:04:04 -04:00
parent e967aaabfb
commit 1446f6da94

View File

@@ -40,28 +40,26 @@ _log = get_logger("db.pool")
async def _force_close(session: AsyncSession) -> None: async def _force_close(session: AsyncSession) -> None:
"""Close a session, forcing connection invalidation if clean close fails. """Close a session, forcing connection invalidation if clean close fails.
Shielded from cancellation and catches every exception class including Under cancellation, ``session.close()`` may try to issue a ROLLBACK on
CancelledError. If session.close() fails (corrupted connection), we a connection that was interrupted mid-query — aiomysql then raises
invalidate the underlying connection so the pool discards it entirely ``InterfaceError("Cancelled during execution")`` and the connection is
rather than leaving it checked-out forever. left checked-out, reported by the pool as ``non-checked-in connection``
on GC. When clean close fails, invalidate the session's connections
directly (no I/O, just flips the pool record) so the pool discards
them immediately instead of waiting for garbage collection.
""" """
try: try:
await asyncio.shield(session.close()) await asyncio.shield(session.close())
return
except BaseException: except BaseException:
# close() failed — connection is likely corrupted. pass
# Try to invalidate the raw connection so the pool drops it. try:
try: # invalidate() is sync and does no network I/O — safe inside a
bind = session.get_bind() # cancelled task. Tells the pool to drop the underlying DBAPI
if hasattr(bind, "dispose"): # connection rather than return it for reuse.
pass # don't dispose the whole engine session.sync_session.invalidate()
# The sync_session holds the connection record; invalidating except BaseException:
# it tells the pool to discard rather than reuse. _log.debug("force-close: invalidate failed", exc_info=True)
sync = session.sync_session
if sync.is_active:
sync.rollback()
sync.close()
except BaseException:
_log.debug("force-close: fallback cleanup failed", exc_info=True)
@asynccontextmanager @asynccontextmanager