Repository navigation
Task cancellation with SA connection pool + asyncpg + direct_tls can lead to connection leak #12099
Description
Describe the bug
Hi, while tracking down #12076 / #12077 we found that we were hitting another connection leak under similar circumstances.
The setup is as follows:
- SA connection pool + asyncpg with a
direct_tls=trueconnection. - an asyncio task takes a session from the pool, starts a transaction and locks some rows with
SELECT ... FOR UPDATE, then does some async work, while holding the lock + transaction - the asyncio task is cancelled.
At this point, it is possible that the connection holding the lock will leak (and still hold the lock).
The direct_tls stuff is what is used to connect to GCP CloudSQL instances with the gcp cloud sql python connector, where we were having the issue in the first place.
Optional link from https://docs.sqlalchemy.org which documents the behavior that is expected
No response
SQLAlchemy Version in Use
2.0.36
DBAPI (i.e. the database driver)
asyncpg
Database Vendor and Major Version
postgresql 16
Python Version
3.12
Operating system
linux
To Reproduce
The setup is a little bit involved, so I created a repo with a docker compose config to make it easier to reproduce:
The idea is to:
- install the deps
- docker compose up to start postgres + nginx (ssl connection termination)
- run direct_tls_leak.py
- watch the number of connections grow beyond the configured pool size.
Here's a copy of the script.
import asyncio
import asyncpg
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.sql import text
# direct_tls seems to be the key to seeing the issue
SQLALCHEMY_DATABASE_URL = "postgresql+asyncpg://postgres:postgres@localhost:5443/postgres?direct_tls=true"
# pool_size must be >1
engine = create_async_engine(SQLALCHEMY_DATABASE_URL, pool_size=2, max_overflow=0, echo_pool="debug")
async_session = async_sessionmaker(engine, expire_on_commit=False)
async def querier():
async with async_session() as session, session.begin():
# locking rows is important for reproducing the issue
res = await session.execute(text("select * from test for update;"))
print([row for row in res])
# simulate more asyncio work while tx is open and lock is held
# task will be cancelled during this
await asyncio.sleep(10)
async def main():
async with async_session() as session, session.begin():
res = await session.execute(text("create table if not exists test (a text);"))
print(res)
res = await session.execute(text("insert into test values ('patate');"))
print(res)
while True:
ts = []
for i in range(3): # 1 more than pool size
ts.append(asyncio.create_task(querier()))
async def timeouter():
await asyncio.sleep(1)
for t in ts:
t.cancel()
asyncio.create_task(timeouter())
try:
await asyncio.gather(*ts)
except asyncio.CancelledError:
print("cancelled")
if __name__ == '__main__':
asyncio.run(main())Error
2024-11-14 11:44:58,563 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Created new connection <AdaptedConnection <asyncpg.connection.Connection object at 0x1224d1040>>
2024-11-14 11:44:58,564 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Created new connection <AdaptedConnection <asyncpg.connection.Connection object at 0x1224d1130>>
2024-11-14 11:44:58,566 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Connection <AdaptedConnection <asyncpg.connection.Connection object at 0x1224d1040>> checked out from pool
2024-11-14 11:44:58,567 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Connection <AdaptedConnection <asyncpg.connection.Connection object at 0x1224d1130>> checked out from pool
2024-11-14 11:44:59,521 INFO sqlalchemy.pool.impl.AsyncAdaptedQueuePool Invalidate connection <AdaptedConnection <asyncpg.connection.Connection object at 0x1224d1040>> (reason: CancelledError:)
2024-11-14 11:44:59,521 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Hard-closing connection <AdaptedConnection <asyncpg.connection.Connection object at 0x1224d1040>>
2024-11-14 11:44:59,522 INFO sqlalchemy.pool.impl.AsyncAdaptedQueuePool Invalidate connection <AdaptedConnection <asyncpg.connection.Connection object at 0x1224d1130>> (reason: CancelledError:)
2024-11-14 11:44:59,522 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Hard-closing connection <AdaptedConnection <asyncpg.connection.Connection object at 0x1224d1130>>
the hard-closed connections are not closed.
Additional context
I am not sure where the issue lies exactly, is it in asyncpg + tls handling or in the error handling for sessions holding a transaction + lock ?
I'm also not sure what would be the best behavior here, when the lock is released / transaction is rolled back, should the session be returned to the pool or hard closed like it is trying to do currently?
Activity
zzzeek commentedon Nov 14, 2024
there's nothing fundamentally different here from #12077 AFAICT. the log shows SQLAlchemy is doing the right thing, hard-closing the connection where we call asyncpg terminate(). asyncpg refuses. seems like people are reporting similar issues over there such as MagicStack/asyncpg#1156 .
if your patch at #12076 works, please make the changes requested (restore CancellledError) and we can merge that
jraby commentedon Nov 14, 2024
I forgot to mention, using
asyncio.shieldas per #12076 does not fix this case, that's why I opened a new bug report.zzzeek commentedon Nov 14, 2024
What would you have SQLAlchemy do ? We call the command asyncpg gives us. we try close(), we try terminate(). I dont see other options
zzzeek commentedon Nov 14, 2024
perhaps try psycopg instead?
zzzeek commentedon Nov 14, 2024
or if all these issues are local to anyio, clearly you'd want to just stop using that 3rd party library
jraby commentedon Nov 14, 2024