anyio taskgroup cancellation leading to asyncpg connection leak #12077

Describe the bug

I hit a connection leak when cancelling an http handler task in a fastapi + starlette + BaseHTTPMiddleware + SA + asyncpg setup and managed to reproduce it with the code below.

From my understanding, when cancelling an anyio taskgroup which contains a task that checked out a connection from the pool, SA tries to close the connection, but most (all?) await calls within that code path bail with asyncio.CancelledError, leading to SA and asyncpg thinking that the connection was closed / released, but the socket is still alive and idle from postgres's point of view.

This is related to MagicStack/asyncpg#1156 which doesn't currently have a resolution.

The following code leaks a connection on every loop operation, as observed by running select * from pg_stat_activity where datname='postgres'; in the db while the script is running.

#12076 is an attempt a fixing this behavior by shielding connection termination from asyncio.CancelledError

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

postgres 17

Python Version

3.12

Operating system

linux

To Reproduce

import asyncio
import asyncpg
import anyio

from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
from sqlalchemy.sql import text

SQLALCHEMY_DATABASE_URL = "postgresql+asyncpg://postgres:postgres@localhost/postgres"
engine = create_async_engine(SQLALCHEMY_DATABASE_URL, pool_size=1, max_overflow=0, echo_pool="debug")
async_session = async_sessionmaker(engine, expire_on_commit=False)

async def anyio_tg_wrapper(func, *args, **kwargs):
    async with anyio.create_task_group() as tg:
        tg.start_soon(func, *args, **kwargs)

async def tg_wrapper(f):
    async with asyncio.TaskGroup() as tg:
        tg.create_task(f)


async def querier():
    async with async_session() as session, session.begin():
        res = await session.execute(text("select 1"))
        print([row for row in res])
        
        # simulate other async io work while holding the session
        await asyncio.sleep(10)

async def main():
    while True:
        t = asyncio.create_task(anyio_tg_wrapper(querier))
        
        # asyncio task group does not lead to a leak
        #t = asyncio.create_task(tg_wrapper(querier()))

        async def timeouter():
            await asyncio.sleep(1)
            t.cancel()

        asyncio.create_task(timeouter())

        try:
            await t
        except asyncio.CancelledError:
            print("cancelled")

if __name__ == '__main__':
    asyncio.run(main())

Error

This is the output from the script above.

2024-11-10 08:07:01,830 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Created new connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7e46b7088e60>>
2024-11-10 08:07:01,831 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7e46b7088e60>> checked out from pool
[(1,)]
2024-11-10 08:07:02,252 INFO sqlalchemy.pool.impl.AsyncAdaptedQueuePool Invalidate connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7e46b7088e60>> (reason: CancelledError:Cancelled by cancel scope 7e46b705a690)
2024-11-10 08:07:02,253 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Hard-closing connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7e46b7088e60>>

at this point, the "hard-closed" connection, is not in the pool anymore. But still alive.

When using an asyncio task group instead of an anyio task group, the session is rolled back and put back in the pool, which is probably a better outcome than closing the connection in this case.

2024-11-10 08:10:56,558 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Created new connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7e75c812caa0>>
2024-11-10 08:10:56,564 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7e75c812caa0>> checked out from pool
[(1,)]
2024-11-10 08:10:57,526 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7e75c812caa0>> being returned to pool
2024-11-10 08:10:57,526 DEBUG sqlalchemy.pool.impl.AsyncAdaptedQueuePool Connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7e75c812caa0>> rollback-on-return

Additional context

No response