asyncpg: shield connection close in terminate to avoid connection leak by jraby · Pull Request #12076 · sqlalchemy/sqlalchemy

jraby · 2024-11-10T12:31:40Z

I hit a connection leak when cancelling a 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 solution proposed here might not be the right place to fix this issue, but it does fix my reproducer :)

The following code leaks a connection on every loop iteration:

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())

Description

If terminate is called from an asyncio.CancelledError exception handler, under an anyio task group, await will bail with asyncio.CancelledError right away, leading to a connection leak.

Shielding the close() call allows the connection to be properly closed in those cases.

Tentatively removed the asyncio.CancelledError except clause since close() doesn't raise it itself so it shouldn't be possible to get it.

Fixes #12077

CaselIT · 2024-11-10T12:51:57Z

Hi,

The pr template asks you to open an issue first. Please do that

zzzeek · 2024-11-10T13:36:51Z

+                await_(asyncio.shield(self._connection.close(timeout=2)))
            except (
                asyncio.TimeoutError,
-                asyncio.CancelledError,


why remove CancelledError here? is it 100% impossible for CancelledError to be raised from sheild()?

jraby · 2024-11-10T14:37:29Z

From what I understand, it is possible only if the shielded function cancels itself from within. But I'm no asyncio expert.

…

On Sun, Nov 10, 2024, 8:37 a.m. Michael Bayer ***@***.***> wrote: ***@***.**** commented on this pull request. ------------------------------ In lib/sqlalchemy/dialects/postgresql/asyncpg.py <#12076 (comment)> : > except ( asyncio.TimeoutError, - asyncio.CancelledError, why remove CancelledError here? is it 100% impossible for CancelledError to be raised from sheild()? — Reply to this email directly, view it on GitHub <#12076 (review)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AASBOTOXEWYJBEZPTS6IFPLZ75OQTAVCNFSM6AAAAABRQFJDJ6VHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZDIMRVG43DINBQGE> . You are receiving this because you authored the thread.Message ID: ***@***.***>

zzzeek · 2024-11-10T14:48:31Z

From what I understand, it is possible only if the shielded function cancels itself from within. But I'm no asyncio expert.
…
On Sun, Nov 10, 2024, 8:37 a.m. Michael Bayer @.> wrote: @.* commented on this pull request. ------------------------------ In lib/sqlalchemy/dialects/postgresql/asyncpg.py <#12076 (comment)> : > except ( asyncio.TimeoutError, - asyncio.CancelledError, why remove CancelledError here? is it 100% impossible for CancelledError to be raised from sheild()? — Reply to this email directly, view it on GitHub <#12076 (review)>, or unsubscribe https://github.com/notifications/unsubscribe-auth/AASBOTOXEWYJBEZPTS6IFPLZ75OQTAVCNFSM6AAAAABRQFJDJ6VHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZDIMRVG43DINBQGE . You are receiving this because you authored the thread.Message ID: @.***>

I would leave it, if only to make sure it's still there if the code above changes. I'm not sure it's apparent but this block of code has been a struggle for years to get asyncpg to close connections

jraby · 2024-11-10T16:01:28Z

Fair. I'll add it back later today (not at the terminal for a while) On Sun, Nov 10, 2024, 9:48 a.m. Michael Bayer ***@***.***> wrote:

…

From what I understand, it is possible only if the shielded function cancels itself from within. But I'm no asyncio expert. … <#m_1772178836461687846_> On Sun, Nov 10, 2024, 8:37 a.m. Michael Bayer *@*.*> wrote: @.** commented on this pull request. ------------------------------ In lib/sqlalchemy/dialects/postgresql/asyncpg.py <#12076 (comment) <#12076 (comment)>> : > except ( asyncio.TimeoutError, - asyncio.CancelledError, why remove CancelledError here? is it 100% impossible for CancelledError to be raised from sheild()? — Reply to this email directly, view it on GitHub <#12076 (review) <#12076 (review)>>, or unsubscribe https://github.com/notifications/unsubscribe-auth/AASBOTOXEWYJBEZPTS6IFPLZ75OQTAVCNFSM6AAAAABRQFJDJ6VHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZDIMRVG43DINBQGE . You are receiving this because you authored the thread.Message ID: *@* .***> I would leave it, if only to make sure it's still there if the code above changes. I'm not sure it's apparent but this block of code has been a struggle for years to get asyncpg to close connections — Reply to this email directly, view it on GitHub <#12076 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AASBOTO2GOOTQR7M5XQAO4LZ75W5JAVCNFSM6AAAAABRQFJDJ6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDINRWG43DIMBSGI> . You are receiving this because you authored the thread.Message ID: ***@***.***>

CaselIT · 2024-11-10T17:11:24Z

I would need to re-look at the old PRs that changed this bit. If I'm not mistaken there was a reason why shield was not added here.
Also if I'm not mistaken the session and connection method do use shield. why it's not working here?
Also does this effect other async drivers too?

In any case it's very annoying that asyncio behaves one way while anyio behaves in another way...

jraby · 2024-11-11T00:44:34Z

Indeed, having asyncio and anyio behave differently in the error path is surprising.

I see what you mean by connection and sessions both using shield, I'll run some tests to see why it doesn't seem to be working as intended.

btw I suggested the fix in this PR as a way to avoid the leak observed in the reproducer script, but I have very little exposure to async python and sqlalchemy in general, so feel free to ignore the proposed fix.
I think the reproducer script might have more value than the fix itself since it shows a straightforward way to get the current implementation to leak a connection.

jraby · 2024-11-11T01:17:49Z

ok the key here seems to be that the reproducer uses session.begin().
removing the transaction avoids the problem.

here's a traceback from where SA starts to invalidate the session/connection (and ends up trying to abort the connection:

  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/engine/util.py", line 162, in __exit__
    self.rollback()
  File "<string>", line 2, in rollback
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/orm/state_changes.py", line 139, in _go
    ret_value = fn(self, *arg, **kw)
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/orm/session.py", line 1360, in rollback
    t[1].rollback()
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 2610, in rollback
    self._do_rollback()
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 2729, in _do_rollback
    self._close_impl(try_deactivate=True)
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 2712, in _close_impl
    self._connection_rollback_impl()
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 2704, in _connection_rollback_impl
    self.connection._rollback_impl()
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1129, in _rollback_impl
    self._handle_dbapi_exception(e, None, None, None, None)    # <<<<<<<<< the exception here is a Cancelled by cancel scope
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 2368, in _handle_dbapi_exception
    self.invalidate(e)
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 751, in invalidate
    pool_proxied_connection.invalidate(exception)
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/pool/base.py", line 1477, in invalidate
    self._connection_record.invalidate(e=e, soft=soft)

And the traceback that _handle_dbapi_exception sees:

Traceback (most recent call last):
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1127, in _rollback_impl
    self.engine.dialect.do_rollback(self.connection)
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/engine/default.py", line 699, in do_rollback
    dbapi_connection.rollback()
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/dialects/postgresql/asyncpg.py", line 864, in rollback
    self.await_(self._rollback_and_discard())
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/util/_concurrency_py3k.py", line 132, in await_only
    return current.parent.switch(awaitable)  # type: ignore[no-any-return,attr-defined] # noqa: E501
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/util/_concurrency_py3k.py", line 196, in greenlet_spawn
    value = await result
            ^^^^^^^^^^^^
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/sqlalchemy/dialects/postgresql/asyncpg.py", line 845, in _rollback_and_discard
    await self._transaction.rollback()
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/asyncpg/transaction.py", line 227, in rollback
    await self.__rollback()
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/asyncpg/transaction.py", line 206, in __rollback
    await self._connection.execute(query)
  File "/home/jean/.cache/pypoetry/virtualenvs/starlette-sa-conn-leak-kv5kKC85-py3.12/lib/python3.12/site-packages/asyncpg/connection.py", line 349, in execute
    result = await self._protocol.query(query, timeout)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

It is possible that the transaction __exit__ func of the transaction should also be shielded against CancelledError ?

If terminate is called from an asyncio.CancelledError exception handler, under an anyio task group, await will bail with asyncio.CancelledError right away, leading to a connection leak. Shielding the close() call allows the connection to be properly closed in those cases.

sqla-tester

OK, this is sqla-tester setting up my work on behalf of zzzeek to try to get revision 6ae261e of this pull request into gerrit so we can run tests and reviews and stuff

sqla-tester · 2024-11-15T13:55:57Z

New Gerrit review created for change 6ae261e: https://gerrit.sqlalchemy.org/c/sqlalchemy/sqlalchemy/+/5594

CaselIT · 2024-11-15T15:30:29Z

is there any way of testing this?

jraby · 2024-11-15T15:44:38Z

the only way I could reproduce and test this reliably is with the code in the description and observing that the socket is closed with netstat + a select in pg_stat_activity.
It is hard to test from within SA since the connection is unusable from its perspective, but the underlying socket is still live.

In other words, I'm not sure how to add a test for this.

sqla-tester · 2024-11-18T17:49:30Z

Michael Bayer (zzzeek) wrote:

this is fine can we work out a changelog message and a 2.0 port

View this in Gerrit at https://gerrit.sqlalchemy.org/c/sqlalchemy/sqlalchemy/+/5594

sqla-tester · 2024-11-18T20:07:01Z

Federico Caselli (CaselIT) wrote:

I don't particularly like changing stuff randomly until hopefully things stop breaking, but I guess in asyncio that's the expectation...

Doesn't seems to break any test, so hopefully it's fine

View this in Gerrit at https://gerrit.sqlalchemy.org/c/sqlalchemy/sqlalchemy/+/5594

Redoubts · 2025-01-14T20:31:49Z

Would you be willing to backport this change to the 1.4 branch? I'm noticing a similar issue there, and if I replicate the change you're making here (sqlalchemy/dialects/postgresql/asyncpg.py:749) then the issue is resolved.

CaselIT · 2025-01-14T20:50:42Z

It's not a given that 1.4 will have other releases. It's probably best if you patch the dialect with the changed method

sqla-tester · 2025-01-15T16:51:43Z

Michael Bayer (zzzeek) wrote:

OK changelog + 2.0 backport?

View this in Gerrit at https://gerrit.sqlalchemy.org/c/sqlalchemy/sqlalchemy/+/5594

sqla-tester · 2025-01-19T20:22:39Z

Michael Bayer (zzzeek) wrote:

mypy-recheck

View this in Gerrit at https://gerrit.sqlalchemy.org/c/sqlalchemy/sqlalchemy/+/5667

sqla-tester · 2025-01-21T00:35:04Z

Gerrit review https://gerrit.sqlalchemy.org/c/sqlalchemy/sqlalchemy/+/5594 has been merged. Congratulations! :)

sqla-tester · 2025-01-21T00:35:09Z

Gerrit review https://gerrit.sqlalchemy.org/c/sqlalchemy/sqlalchemy/+/5667 has been merged. Congratulations! :)

Added an additional ``shield()`` call within the connection terminate process of the asyncpg driver, to mitigate an issue where terminate would be prevented from completing under the anyio concurrency library. Fixes #12077 Closes: #12076 Pull-request: #12076 Pull-request-sha: 6ae261e Change-Id: Iaec63d02b620201dc60ffdad76ba9d61d427cac1 (cherry picked from commit 2f6ca65)

jraby changed the title ~~asyncpy: shield connection close in terminate to avoid connection leak~~ asyncpg: shield connection close in terminate to avoid connection leak Nov 10, 2024

jraby mentioned this pull request Nov 10, 2024

anyio taskgroup cancellation leading to asyncpg connection leak #12077

Closed

zzzeek reviewed Nov 10, 2024

View reviewed changes

jraby mentioned this pull request Nov 14, 2024

Task cancellation with SA connection pool + asyncpg + direct_tls can lead to connection leak #12099

Closed

jraby force-pushed the task-cancellation-anyio-tg-asyncpg-conn-leak branch from c7e6729 to 6ae261e Compare November 15, 2024 13:36

zzzeek requested a review from sqla-tester November 15, 2024 13:55

sqla-tester reviewed Nov 15, 2024

View reviewed changes

sqlalchemy-bot closed this in 2f6ca65 Jan 21, 2025

coderabbitai Bot mentioned this pull request Jul 22, 2026

Fix/bug Bongseop-Kim/essesion#21

Merged

coderabbitai Bot mentioned this pull request Jul 28, 2026

fix: #527 the 250ms resolver budget must bound resolution, not cold DB connect alfred-os/AlfredOS#531

Merged

fidelix mentioned this pull request Aug 18, 2026

[BUG] POST /threads/{id}/stream/events leaks pooled asyncpg connections on SSE cancel aegra/aegra#530

Closed

Conversation

jraby commented Nov 10, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Uh oh!

CaselIT commented Nov 10, 2024

Uh oh!

zzzeek Nov 10, 2024

Choose a reason for hiding this comment

Uh oh!

jraby Nov 15, 2024

Choose a reason for hiding this comment

Uh oh!

jraby commented Nov 10, 2024 via email

Uh oh!

zzzeek commented Nov 10, 2024

Uh oh!

jraby commented Nov 10, 2024 via email

Uh oh!

CaselIT commented Nov 10, 2024

Uh oh!

jraby commented Nov 11, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jraby commented Nov 11, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sqla-tester left a comment

Choose a reason for hiding this comment

Uh oh!

sqla-tester commented Nov 15, 2024

Uh oh!

CaselIT commented Nov 15, 2024

Uh oh!

jraby commented Nov 15, 2024

Uh oh!

sqla-tester commented Nov 18, 2024

Uh oh!

sqla-tester commented Nov 18, 2024

Uh oh!

Redoubts commented Jan 14, 2025

Uh oh!

CaselIT commented Jan 14, 2025

Uh oh!

sqla-tester commented Jan 15, 2025

Uh oh!

sqla-tester commented Jan 19, 2025

Uh oh!

sqla-tester commented Jan 21, 2025

Uh oh!

sqla-tester commented Jan 21, 2025

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

jraby commented Nov 10, 2024 •

edited

Loading

jraby commented Nov 11, 2024 •

edited

Loading

jraby commented Nov 11, 2024 •

edited

Loading