Terminate function does not work when a stream is cancelled prematurely when in an anyio task group #1156

Activity

1. elprans commented on Jul 17, 2024

   

   elprans

   on Jul 17, 2024

   Member

   More actions

   What happens is that anyio cancels await conn.close() before it has a chance to send the termination message to Postgres. Connection does get closed from the client side, but Postgres never receives the memo. I need to think if shielding Protocol.close() makes sense, but you can use the idle_in_transaction_session_timeout server setting as a workaround for now to kill inactive sessions.
2. JordanZimmitti commented on Jul 31, 2024

   

   JordanZimmitti

   on Jul 31, 2024

   Author

   More actions

   @elprans Thanks for the insight! I'll experiment with idle_in_transaction_session_timeout in the meantime though I will add that the connections can run up quickly on us. Keep me posted on any updates or if I can help provide anymore info/context on the situation.

   Thanks for looking into this issue!
3. JordanZimmitti commented on Oct 21, 2024

   

   JordanZimmitti

   on Oct 21, 2024

   Author

   More actions

   @elprans Hi!

   I am hoping for a small update on this, as it's still an ongoing issue and forcing me not to update the version of sqlalchemy I'm using

   Thanks!
4. jraby commented on Nov 14, 2024

   

   jraby

   on Nov 14, 2024

   More actions

   I think I hit the same issue as this and put together a self contained reproducer in https://github.com/brittlesoft/repro-starlette-sa-conn-leak

   There's also a similar issue that seem to only happen with direct_tls=true connections as shown in the direct_tls_leak.py script from the repo above.

   I'm not sure if it is the same issue or not so I wanted to post here before opening a possibly duplicate issue.
5. circlingthesun commented on Dec 2, 2024

   

   circlingthesun

   on Dec 2, 2024

   More actions

   @elprans have you given shielding Protocol.close() any more thought?
6. circlingthesun commented on Dec 3, 2024

   

   circlingthesun

   on Dec 3, 2024

   More actions

   I don't believe the idle_in_transaction_session_timeout helps if a connections isn't in the middle of a transaction.
7. circlingthesun commented on Dec 3, 2024

   

   circlingthesun

   on Dec 3, 2024

   More actions

   Monkey patching the close function appears to work in the example copied from sqlalchemy/sqlalchemy#11128 (reply in thread)

   import asyncio
   
   import anyio
   import asyncpg
   from asyncpg import Connection
   
   
   def monkey_patch():
       original_close = Connection.close
   
       # shield the close method from being cancelled
       async def close(self, *, timeout=None):
           await asyncio.shield(original_close(self, timeout=timeout))
   
       Connection.close = close
   
   
   # Uncomment for bug
   monkey_patch()
   
   async def partitions(conn, cursor, size):
       try:
           while True:
               partition = await cursor.fetch(size)
               if partition:
                   yield partition
               else:
                   break
       except BaseException as e:
           print(f"Got exception {e}; attempting to close")
           try:
               await conn.close(timeout=2)
           except (
               asyncio.CancelledError,
               asyncio.TimeoutError,
               OSError,
           ) as e2:
               print(
                   f"Got a second exception {e2} while trying to close, trying to terminate"
               )
               conn.terminate()
               print(
                   "Conn terminate succeeded!   is the connection closed?  no?  then that's a bug in asyncpg"
               )
   
   
   async def run():
       async def go_raw_anyio():
           conn = await asyncpg.connect(
               user="doadmin",
               password="bookem",
               host="localhost",
               database="defaultdb",
               server_settings={"application_name": "bug"},
           )
           _transaction = conn.transaction()
           await _transaction.start()
   
           cursor = await conn.cursor(
               "SELECT anon_1 FROM "
               "generate_series(1::INTEGER, 500::INTEGER) AS anon_1"
           )
   
           async for _ in partitions(conn, cursor, 10):
               tg.cancel_scope.cancel()
   
       for _ in range(5):
           async with anyio.create_task_group() as tg:
               tg.start_soon(go_raw_anyio)
   
       await asyncio.sleep(5)
   
       print("Check the number of connections in the database:")
       print(
           "SELECT * FROM pg_stat_activity WHERE datname = 'defaultdb' and application_name = 'bug';"
       )
       input("Press Enter to continue...")
   
   
   if __name__ == "__main__":
       asyncio.run(run())

   👍React with 👍1Reacted by Savio Mak