Async connections not returned to pool if task is cancelled

Describe the bug

The AsyncConnection context manager does not reliably close the connection when exiting due to a cancellation. This results in resources leaking, which is clearly surfaced by the warning during garbage collection. Garbage collection can't perform the cleanup itself, but tells us that something was missed earlier.

To Reproduce

import anyio
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.sql import text

async def main():
    engine = create_async_engine(
        "postgresql+asyncpg://username:password@localhost/db_name",
        echo=True,
        connect_args={"timeout": 5}
        )

    async def sleep_in_db():
        print(engine.pool.status()) # 0 connections
        async with engine.connect() as connection:
            print(engine.pool.status()) # 1 connections
            while True:
                await connection.execute(text("select now();"))
                await anyio.sleep(1)

    async with anyio.create_task_group() as db_tasks:
        db_tasks.start_soon(sleep_in_db)

        await anyio.sleep(2.5)

        db_tasks.cancel_scope.cancel() # Will raise a CancelledError under the hood


    print(engine.pool.status()) # 1 connection????

    await engine.dispose()

anyio.run(main)

Error

The garbage collector is trying to clean up connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7f1cf4760660>>. This feature is unsupported on async dbapi, since no IO can be performed at this stage to reset the connection. Please close out all connections when they are no longer used, calling ``close()`` or using a context manager to manage their lifetime.
sys:1: SAWarning: The garbage collector is trying to clean up connection <AdaptedConnection <asyncpg.connection.Connection object at 0x7f1cf4760660>>. This feature is unsupported on async dbapi, since no IO can be performed at this stage to reset the connection. Please close out all connections when they are no longer used, calling ``close()`` or using a context manager to manage their lifetime.

Versions

OS: Ubuntu on WSL2 in Windows 10
Python: 3.10.4
SQLAlchemy: 1.4.37
Database: Postgres 14
DBAPI (eg: psycopg, cx_oracle, mysqlclient): asyncpg 0.25.0

Additional context

I think the problem is that the way cancellation works is that the event loop (asyncio in this case) raises CancelledError when you do await, so if any await is needed while actually processing a cancellation (ex. to cleanup), that needs to be "shielded" from the cancellation we're trying to process. We can see in AsyncConnection that on __aexit__() it's trying to await self.close(), but it's unable to (self.close() is never run) because of the cancellation: https://github.com/sqlalchemy/sqlalchemy/blob/main/lib/sqlalchemy/ext/asyncio/engine.py#L700

    async def __aexit__(self, type_: Any, value: Any, traceback: Any) -> None:
        await self.close()

Here's a good description of this concept in the trio docs: https://trio.readthedocs.io/en/stable/reference-core.html?highlight=shield#cancellation-semantics (it takes a bit to get to the relevant part, but read the whole block).

I don't know asyncio well enough to code up a robust solution but I believe these are the relevant docs: https://docs.python.org/3/library/asyncio-task.html#shielding-from-cancellation

Using anyio to achieve trio-like structured concurrency on top of asyncio, I am able to patch this __aexit__() method and then the connection is properly closed. I'm not sure if you'd consider using anyio in sqlalchemy. Could be used in just this one place, throughout the codebase in order to add compatibility with trio, or piecemeal to get the advantages of structured concurrency without a big refactor. Here's the anyio docs for shielding: https://anyio.readthedocs.io/en/stable/cancellation.html?highlight=shield#shielding

import anyio

class PatchedAsyncConnection(AsyncConnection):
    async def __aexit__(self, type_, value, traceback):
        with anyio.open_cancel_scope(shield=True):
            await self.close()

hi -

I dont really see where this "shield" could be on the SQLAlhcemy side. if you check out a connection from the pool, everything you do with that connection is fine, but as long as you are checked out, that entire operation needs to be shielded to prevent this issue. that is above, everything that happens in sleep_in_db() after engine.connect(). if anything is cancelled in that block, there's no longer a context in which to close the connection.

as far as anyio, looking at https://anyio.readthedocs.io/en/stable/cancellation.html#shielding I dont see this really being changed much, they just give you a context manager to do the same thing which is fine.

but if we just shield in our __aexit__() method, that's not going to work IIUC, because if the task is cancelled before we get to __aexit__(), then we're already cancelled and aexit isnt' called from within the event loop, right? or no?

@zzzeek

Hi @zzzeek. Thanks for the response.

The shield should be in __aexit__(), where it calls self.close(). See my "patch" above. I think that the event loop is smart enough to call __aexit__() even when there is a cancellation in progress and the problem is that it's the subsequent call to self.close() that doesn't happen. I'm not sure how to elaborate on this further other than to say that my code above reproduces the issue and the fix above solves it. You can try running this yourself, adding print statements, etc...

If desired for functionality reasons, one could theoretically add a shield to the application code (ex. the entire sleep_in_db() function). But in this case, we don't want that: that's the code we intend to cancel and it's just the connection cleanup that shouldn't also be cancelled. Like, what we want is the connection.execute(text("select now();")) or anyio.sleep(1) to be cancelled. Maybe there's a database query we'd like to try, but if it takes too long so we abort. Or maybe, as in this example code, we want to run a database query periodically indefinitely but cancel it based on some external action. Originally I'd coded up a call to postgres' pg_sleep() (hence the function name of sleep_in_db()), but since changed it to actually just run a query on a timer. Hopefully the basic idea still comes across.

As the anyio docs put it:

There are cases where you want to shield your task from cancellation, at least temporarily. The most important such use case is performing shutdown procedures on asynchronous resources.

The issue here seems only when using anyio.
If I create the same example using asyncio there is no issue since the aexits get a chance of cleaning up correctly

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

async def main():
    engine = create_async_engine(
        "postgresql+asyncpg://scott:tiger@localhost/test",
        echo=True,
        connect_args={"timeout": 5}
        )

    async def sleep_in_db():
        print(engine.pool.status()) # 0 connections
        async with engine.connect() as connection:
            print(engine.pool.status()) # 1 connections
            while True:
                await connection.execute(text("select now();"))
                await asyncio.sleep(1)

    t = asyncio.create_task(sleep_in_db())
    await asyncio.sleep(2.5)
    t.cancel()


    await asyncio.sleep(1)
    print(engine.pool.status()) # 1 connection????
    await engine.dispose()

asyncio.run(main())

also from my understating the leaked connection here is because asyncpg does not use shield.
If I wrap the asyncpg connection in a proxy that uses shield on the close the connection is closed from the db side. There is still the garbage connection message, but that's a false positive, caused by the not execution of the other cleanup logic sqlalchemy

If I read the problem correctly is how anyio works, that prevents any normal use of aexit.
Not sure if the burden should be on other libraries to make changes to accommodate anyio strange behaviour

By strange behaviour I mean this.

Using asyncio

import asyncio

async def go():
    print("go")

    async def task():
        try:
            print("before sleep 1")
            await asyncio.sleep(1)
            print("after sleep 1, sleep 2")
        finally:
            print("finally")
            await asyncio.sleep(1)
            print("end sleep 2")

    t = asyncio.create_task(task())
    print("sleep go")
    await asyncio.sleep(0.1)
    t.cancel()
    await asyncio.sleep(2)
    print("end")

asyncio.run(go())

here finally has a change of properly clean up after itself

Using anyio instead

import anyio

async def go():
    print("go")

    async def task():
        try:
            print("before sleep 1")
            await anyio.sleep(1)
            print("after sleep 1, sleep 2")
        finally:
            print("finally")
            await anyio.sleep(1)
            print("end sleep 2")

    async with anyio.create_task_group() as tasks:
        tasks.start_soon(task)
        print("sleep go")
        await anyio.sleep(0.1)
        tasks.cancel_scope.cancel()
    await anyio.sleep(2)
    print("end")

anyio.run(go)

any await in finally does not run, "end sleep 2" is never executed. This seems to be quite broken on their part.

does asyncpg include any context manager functionalty and if so what do they do here? they are after all authors of asyncio itself

they do have context manager functionality and no they do not seem to use shield.

but from a quick look I can't say if the close method is safe from a cancel prospective

Regarding the issue here, I think the issue here could also happen (with a less probability) also when using asyncio, since it could happen that cancel is called on a task that's executing the close operation. I'll try to experiment on that.

If that's the case then I guess we may need to investigate where and how to use shield of equivalent functionality

This should do it I think?

with anyio.CancelScope() as scope:
    async with engine.connect() as engine:
        try:
             ...
        finally:
            scope.sheild = True

but generally an engine should try do something like:

try:
    await self.drain()  # or flush or etc
finally:
    self._sock.close() # this should close the OS socket and will result in a ragged eof in TLS if drain is cancelled

If you have a connection pool you could return the connection using a background task scoped to that pool

Federico Caselli has proposed a fix for this issue in the main branch:

Tried adding shield where close cleanup is done https://gerrit.sqlalchemy.org/c/sqlalchemy/sqlalchemy/+/3933

You probably don't want to shield this, you should instead synchronously invalidate the unwrapped socket (close the underlying file descriptor) on a BaseException

Federico Caselli has proposed a fix for this issue in the main branch:

Tried adding shield where close cleanup is done https://gerrit.sqlalchemy.org/c/sqlalchemy/sqlalchemy/+/3933

In general closing a sqlalchemy connection will not close the driver connection, but return it to the pool, so it would not be appropriate to close the driver connection (even if possible, since sqlalchemy uses drivers not socket directly and not all may expose a sync terminate method similar to asyncpg).

the pool doesnt have a background thread that's culling old connections. so if you have a pool that isn't being used, the connection stays there until someone goes to get a new connection from it. if the pool is only lightly used, then there still can be connections in there that may be any level of "old", they could be days old if the application has not seen any periods of high concurrency.

with the fix to just terminate them, there shouldn't be any problem.

also your LB here is not being friendly. there's ways to close connections that is more quickly detectable

with the fix to just terminate them, there shouldn't be any problem.

Are you referring to #8419? I really hope so. I will be glued to my chair and wait to test it out :)

also your LB here is not being friendly. there's ways to close connections that is more quickly detectable

I agree. Im sure Sanic has something to do with the whole process as well. I am still investigating those issues through my APM software. It's keeping me busy :)

Even if the connection is currently used, it should just die and get refreshed gracefully

that's not really possible to refresh a connection gracefully since if it's in a transaction you would lose everything since the last commit

considering the number of api calls, that's strange the connection would be left in the pool for more than the hour. maybe try setting a shorter pool recycle? like 30min for 1h of credential validity?

considering the number of api calls, that's strange the connection would be left in the pool for more than the hour. maybe try setting a shorter pool recycle? like 30min for 1h of credential validity?

if the pool is configured for 50 connections, and their web container is configured for only 45 concurrent requests, then no matter what load you have, you will be subject to potentially 5 connections that might not get used for a long time.

considering the number of api calls, that's strange the connection would be left in the pool for more than the hour. maybe try setting a shorter pool recycle? like 30min for 1h of credential validity?

if the pool is configured for 50 connections, and their web container is configured for only 45 concurrent requests, then no matter what load you have, you will be subject to potentially 5 connections that might not get used for a long time.

sure, but usually you would have just 45 connections in the pool, since at no point 45+ would be checked out.
still that may be a reason

this is invalid asyncio, see https://docs.python.org/3.11/library/asyncio-task.html?highlight=important%20weak%20references%20disappearing%20mid-execution#asyncio.shield

how would you improve the existing code?

Keep a reference to a group of tasks on the connection pool

eg:

task = asyncio.create_task(conn.close())
pool._tasks.add(task)
task.add_done_callback(pool._tasks.discard)
await asyncio.shield(task)

Then your connection pool can wait for self._tasks

Even better is using an anyio TaskGroup in your pool
Even better still is making an upstream PR that makes sure the connection is always closed when asyncpg close terminates

Thanks for the example, I'll open a new issue to improve this.

fix(test): adjust performance test timeouts for CI environment

feat(frontend): add variant detail page with comprehensive visualiza…

fix(infra): pool-leak + Scrape.do 429 storm

Navigation Menu

Repository navigation

Async connections not returned to pool if task is cancelled #8145

Describe the bug

To Reproduce

Error

Versions

Additional context

44 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Async connections not returned to pool if task is cancelled #8145

Description

Describe the bug

To Reproduce

Error

Versions

Additional context

Activity

zzzeek commented on Jun 16, 2022

camlee commented on Jun 16, 2022

CaselIT commented on Jun 17, 2022

CaselIT commented on Jun 17, 2022

zzzeek commented on Jun 17, 2022

CaselIT commented on Jun 17, 2022

graingert commented on Jun 17, 2022

graingert commented on Jun 17, 2022

sqla-tester commented on Jun 17, 2022

graingert commented on Jun 17, 2022

CaselIT commented on Jun 17, 2022

44 remaining items

zzzeek commented on Aug 31, 2022

zzzeek commented on Aug 31, 2022

zikphil commented on Aug 31, 2022

CaselIT commented on Aug 31, 2022

zzzeek commented on Aug 31, 2022

CaselIT commented on Aug 31, 2022

graingert commented on Sep 13, 2022

CaselIT commented on Sep 13, 2022

graingert commented on Sep 13, 2022

CaselIT commented on Sep 13, 2022

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions