From 9c7d60ed1f68d597598ca31684a42eb31fe60a31 Mon Sep 17 00:00:00 2001 From: Gyanu Mayank Date: Fri, 18 Sep 2026 23:04:42 +0530 Subject: [PATCH 1/3] Bound close(timeout) around a pending cancel close() awaited cancel_waiter with no timeout, and connection_lost only resolved waiter. A frozen server left close() hanging even after the socket dropped. Fixes #1356 --- asyncpg/protocol/protocol.pxd | 1 + asyncpg/protocol/protocol.pyx | 61 +++++++++++++++++++++++++---------- tests/test_timeout.py | 17 ++++++++++ 3 files changed, 62 insertions(+), 17 deletions(-) diff --git a/asyncpg/protocol/protocol.pxd b/asyncpg/protocol/protocol.pxd index cd221fbb..39b914b3 100644 --- a/asyncpg/protocol/protocol.pxd +++ b/asyncpg/protocol/protocol.pxd @@ -70,6 +70,7 @@ cdef class BaseProtocol(CoreProtocol): cdef _on_result__copy_in(self, object waiter) cdef _handle_waiter_on_connection_lost(self, cause) + cdef _complete_cancel_waiters(self) cdef _dispatch_result(self) diff --git a/asyncpg/protocol/protocol.pyx b/asyncpg/protocol/protocol.pyx index acce4e9f..486f7de0 100644 --- a/asyncpg/protocol/protocol.pyx +++ b/asyncpg/protocol/protocol.pyx @@ -591,7 +591,15 @@ cdef class BaseProtocol(CoreProtocol): return not self.closing and self.con_status == CONNECTION_OK def abort(self): + # Always finish pending cancel waiters. close() sets closing=True + # before awaiting them, so a later abort() must still unblock + # those futures and drop the transport. + self._complete_cancel_waiters() if self.closing: + if self.transport is not None: + transport = self.transport + self.transport = None + transport.abort() return self.closing = True self._handle_waiter_on_connection_lost(None) @@ -604,26 +612,25 @@ cdef class BaseProtocol(CoreProtocol): return self.closing = True + timeout = self._get_timeout_impl(timeout) - if self.cancel_sent_waiter is not None: - await self.cancel_sent_waiter - self.cancel_sent_waiter = None - - if self.cancel_waiter is not None: - await self.cancel_waiter - - if self.waiter is not None: - # If there is a query running, cancel it - self._request_cancel() - await self.cancel_sent_waiter - self.cancel_sent_waiter = None - if self.cancel_waiter is not None: - await self.cancel_waiter + try: + if timeout is None: + await self._drain_cancels() + else: + await asyncio.wait_for(self._drain_cancels(), timeout) + except asyncio.TimeoutError: + # Server never acknowledged the cancel. Abort rather than + # hanging on cancel_waiter until the transport dies. + self._complete_cancel_waiters() + if self.transport is not None: + transport = self.transport + self.transport = None + transport.abort() + return assert self.waiter is None - timeout = self._get_timeout_impl(timeout) - # Ask the server to terminate the connection and wait for it # to drop. self.waiter = self._new_waiter(timeout) @@ -637,7 +644,9 @@ cdef class BaseProtocol(CoreProtocol): pass finally: self.waiter = None - self.transport.abort() + if self.transport is not None: + self.transport.abort() + self.transport = None def _request_cancel(self): self.cancel_waiter = self.create_future() @@ -686,6 +695,21 @@ cdef class BaseProtocol(CoreProtocol): def _create_future_fallback(self): return asyncio.Future(loop=self.loop) + cdef _complete_cancel_waiters(self): + if self.cancel_sent_waiter is not None and not self.cancel_sent_waiter.done(): + self.cancel_sent_waiter.set_result(None) + self.cancel_sent_waiter = None + if self.cancel_waiter is not None and not self.cancel_waiter.done(): + self.cancel_waiter.set_result(None) + self.cancel_waiter = None + + async def _drain_cancels(self): + await self._wait_for_cancellation() + if self.waiter is not None: + # If there is a query running, cancel it + self._request_cancel() + await self._wait_for_cancellation() + cdef _handle_waiter_on_connection_lost(self, cause): if self.waiter is not None and not self.waiter.done(): exc = apg_exc.ConnectionDoesNotExistError( @@ -695,6 +719,8 @@ cdef class BaseProtocol(CoreProtocol): exc.__cause__ = cause self.waiter.set_exception(exc) self.waiter = None + self._complete_cancel_waiters() + cdef _set_server_parameter(self, name, val): self.settings.add_setting(name, val) @@ -940,6 +966,7 @@ cdef class BaseProtocol(CoreProtocol): else: self.waiter.set_exception(exc) self.waiter = None + self._complete_cancel_waiters() else: # The connection was lost because it was # terminated or due to another error; diff --git a/tests/test_timeout.py b/tests/test_timeout.py index 152a504a..433324de 100644 --- a/tests/test_timeout.py +++ b/tests/test_timeout.py @@ -152,3 +152,20 @@ async def test_timeout_covers_prepare_01(self): with self.assertRaises(asyncio.TimeoutError): meth = getattr(self.con, methname) await meth('select pg_sleep($1)', 0.2) + + +class TestCloseTimeoutPendingCancel(tb.ConnectedTestCase): + + async def test_close_times_out_pending_cancel(self): + with self.assertRaises(asyncio.TimeoutError): + await self.con.fetch('select pg_sleep(10)', timeout=0.05) + + proto = self.con._protocol + proto.cancel_waiter = self.loop.create_future() + proto.cancel_sent_waiter = None + + with self.assertRunUnder(1): + await self.con.close(timeout=0.15) + + self.assertTrue(self.con.is_closed()) + From 0da418812304bcaea3cc0041327d48711c6f3302 Mon Sep 17 00:00:00 2001 From: Elvis Pranskevichus Date: Mon, 21 Sep 2026 10:39:01 -0700 Subject: [PATCH 2/3] Fix close timeout budgeting and cancellation cleanup --- asyncpg/protocol/protocol.pyx | 60 +++++++-------- tests/test_timeout.py | 134 ++++++++++++++++++++++++++++++---- 2 files changed, 151 insertions(+), 43 deletions(-) diff --git a/asyncpg/protocol/protocol.pyx b/asyncpg/protocol/protocol.pyx index 486f7de0..be8fa7e7 100644 --- a/asyncpg/protocol/protocol.pyx +++ b/asyncpg/protocol/protocol.pyx @@ -612,41 +612,41 @@ cdef class BaseProtocol(CoreProtocol): return self.closing = True - timeout = self._get_timeout_impl(timeout) - + close_waiter = None try: - if timeout is None: + timeout = self._get_timeout_impl(timeout) + # Cancellation and the final disconnect share one deadline. + async with compat.timeout(timeout): await self._drain_cancels() - else: - await asyncio.wait_for(self._drain_cancels(), timeout) - except asyncio.TimeoutError: - # Server never acknowledged the cancel. Abort rather than - # hanging on cancel_waiter until the transport dies. - self._complete_cancel_waiters() - if self.transport is not None: - transport = self.transport - self.transport = None - transport.abort() - return - assert self.waiter is None + # Transport loss completes the cancellation futures too. + # There will be no further disconnect notification to await. + if self.con_status != CONNECTION_OK: + return - # Ask the server to terminate the connection and wait for it - # to drop. - self.waiter = self._new_waiter(timeout) - self._terminate() - try: - await self.waiter - except ConnectionResetError: - # There appears to be a difference in behaviour of asyncio - # in Windows, where, instead of calling protocol.connection_lost() - # a ConnectionResetError will be thrown into the task. - pass + assert self.waiter is None + + # The timeout context also covers errors while sending + # Terminate; do not start a separate query timeout timer. + close_waiter = self._new_waiter(None) + self._terminate() + try: + await close_waiter + except ConnectionResetError: + # On Windows the transport may raise this instead of + # calling protocol.connection_lost(). + pass finally: - self.waiter = None + if close_waiter is not None and not close_waiter.done(): + close_waiter.cancel() + if self.timeout_handle is not None: + self.timeout_handle.cancel() + self.timeout_handle = None + self._handle_waiter_on_connection_lost(None) if self.transport is not None: - self.transport.abort() + transport = self.transport self.transport = None + transport.abort() def _request_cancel(self): self.cancel_waiter = self.create_future() @@ -696,7 +696,8 @@ cdef class BaseProtocol(CoreProtocol): return asyncio.Future(loop=self.loop) cdef _complete_cancel_waiters(self): - if self.cancel_sent_waiter is not None and not self.cancel_sent_waiter.done(): + if (self.cancel_sent_waiter is not None and + not self.cancel_sent_waiter.done()): self.cancel_sent_waiter.set_result(None) self.cancel_sent_waiter = None if self.cancel_waiter is not None and not self.cancel_waiter.done(): @@ -721,7 +722,6 @@ cdef class BaseProtocol(CoreProtocol): self.waiter = None self._complete_cancel_waiters() - cdef _set_server_parameter(self, name, val): self.settings.add_setting(name, val) diff --git a/tests/test_timeout.py b/tests/test_timeout.py index 433324de..549bb2d3 100644 --- a/tests/test_timeout.py +++ b/tests/test_timeout.py @@ -6,8 +6,11 @@ import asyncio +import contextlib +from unittest import mock import asyncpg +from asyncpg import connect_utils from asyncpg import connection as pg_connection from asyncpg import _testbase as tb @@ -154,18 +157,123 @@ async def test_timeout_covers_prepare_01(self): await meth('select pg_sleep($1)', 0.2) -class TestCloseTimeoutPendingCancel(tb.ConnectedTestCase): +class TestCloseTimeoutPendingCancel(tb.ClusterTestCase): + + @contextlib.asynccontextmanager + async def pending_cancel(self, *, cancel_sent, command_timeout=None): + con = await self.connect(command_timeout=command_timeout) + cancel_started = asyncio.Event() + + async def cancel(**kwargs): + cancel_started.set() + if not cancel_sent: + await self.loop.create_future() + + # Exercise the actual Cython futures: either the cancel connection + # stalls, or it closes without the query receiving ReadyForQuery. + with mock.patch.object(connect_utils, '_cancel', cancel): + cancellations = [] + try: + with self.assertRaises(asyncio.TimeoutError): + await con.execute('select pg_sleep(10)', timeout=0.01) + await cancel_started.wait() + cancellations = list(con._cancellations) + self.assertTrue(con._protocol._is_cancelling()) + yield con + finally: + con.terminate() + con._transport.abort() + await asyncio.gather(*cancellations, return_exceptions=True) async def test_close_times_out_pending_cancel(self): - with self.assertRaises(asyncio.TimeoutError): - await self.con.fetch('select pg_sleep(10)', timeout=0.05) - - proto = self.con._protocol - proto.cancel_waiter = self.loop.create_future() - proto.cancel_sent_waiter = None - - with self.assertRunUnder(1): - await self.con.close(timeout=0.15) - - self.assertTrue(self.con.is_closed()) - + for cancel_sent in (False, True): + with self.subTest(cancel_sent=cancel_sent): + async with self.pending_cancel(cancel_sent=cancel_sent) as con: + proto = con._protocol + with self.assertRaises(asyncio.TimeoutError), \ + self.assertRunUnder(MAX_RUNTIME): + await con.close(timeout=0.05) + self.assertTrue(con._transport.is_closing()) + self.assertFalse(proto._is_cancelling()) + + async def test_close_uses_command_timeout(self): + async with self.pending_cancel( + cancel_sent=False, command_timeout=0.05) as con: + with self.assertRaises(asyncio.TimeoutError), \ + self.assertRunUnder(MAX_RUNTIME): + await con.close() + self.assertTrue(con._transport.is_closing()) + + async def test_cancel_close_aborts_transport(self): + for cancel_sent in (False, True): + with self.subTest(cancel_sent=cancel_sent): + async with self.pending_cancel(cancel_sent=cancel_sent) as con: + task = self.loop.create_task(con.close()) + await asyncio.sleep(0) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + self.assertTrue(con._transport.is_closing()) + + async def test_connection_lost_during_close(self): + for cancel_sent in (False, True): + for timeout in (None, 0.05): + with self.subTest(cancel_sent=cancel_sent, timeout=timeout): + async with self.pending_cancel( + cancel_sent=cancel_sent) as con: + task = self.loop.create_task( + con.close(timeout=timeout)) + await asyncio.sleep(0) + con._transport.abort() + await asyncio.wait_for(task, MAX_RUNTIME) + self.assertTrue(con.is_closed()) + self.assertFalse(con._protocol._is_cancelling()) + # A lost transport must not leave a new shutdown + # timer that attempts another cancel after cleanup. + await asyncio.sleep(0.1) + self.assertFalse(con._cancellations) + + async def test_close_has_one_timeout_budget(self): + con = await self.connect() + proto = con._protocol + drain_cancels = proto._drain_cancels + + async def delayed_drain(): + await asyncio.sleep(0.3) + await drain_cancels() + # Hold back the final disconnect notification after cancellation + # completes, so both phases consume part of the timeout budget. + con._transport.pause_reading() + + try: + with self.assertRaises(asyncio.TimeoutError): + await con.execute('select pg_sleep(10)', timeout=0.01) + with mock.patch.object(proto, '_drain_cancels', delayed_drain): + with self.assertRaises(asyncio.TimeoutError), \ + self.assertRunUnder(0.7): + await con.close(timeout=0.5) + self.assertTrue(con._transport.is_closing()) + finally: + con.terminate() + con._transport.abort() + + async def test_close_timeout_resolves_running_query(self): + con = await self.connect() + + async def cancel(**kwargs): + await self.loop.create_future() + + with mock.patch.object(connect_utils, '_cancel', cancel): + task = self.loop.create_task(con.execute('select pg_sleep(10)')) + try: + await asyncio.sleep(0) + with self.assertRaises(asyncio.TimeoutError): + await con.close(timeout=0.05) + with self.assertRaises(asyncpg.ConnectionDoesNotExistError): + await asyncio.wait_for(task, MAX_RUNTIME) + self.assertTrue(con._transport.is_closing()) + finally: + con.terminate() + con._transport.abort() + task.cancel() + await asyncio.gather(task, return_exceptions=True) From 24e3bf29d7bd486cdb92753293de6faa35d1116a Mon Sep 17 00:00:00 2001 From: Elvis Pranskevichus Date: Mon, 21 Sep 2026 11:19:31 -0700 Subject: [PATCH 3/3] Make close deadline test portable to Windows --- tests/test_timeout.py | 24 ++++++++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/tests/test_timeout.py b/tests/test_timeout.py index 549bb2d3..4b176ea8 100644 --- a/tests/test_timeout.py +++ b/tests/test_timeout.py @@ -234,16 +234,31 @@ async def test_connection_lost_during_close(self): self.assertFalse(con._cancellations) async def test_close_has_one_timeout_budget(self): - con = await self.connect() + terminate_sent = asyncio.Event() + + class NoTerminateProtocol(connect_utils.protocol.Protocol): + def connection_made(self, transport): + def write(data): + if bytes(data) == b'X\x00\x00\x00\x04': + terminate_sent.set() + else: + transport.write(data) + + # Keep the server connected after Terminate. Pausing reads + # does not suppress disconnects on Windows' Proactor loop. + wrapped = mock.Mock(wraps=transport) + wrapped.write.side_effect = write + super().connection_made(wrapped) + + with mock.patch.object(connect_utils.protocol, 'Protocol', + NoTerminateProtocol): + con = await self.connect() proto = con._protocol drain_cancels = proto._drain_cancels async def delayed_drain(): await asyncio.sleep(0.3) await drain_cancels() - # Hold back the final disconnect notification after cancellation - # completes, so both phases consume part of the timeout budget. - con._transport.pause_reading() try: with self.assertRaises(asyncio.TimeoutError): @@ -252,6 +267,7 @@ async def delayed_drain(): with self.assertRaises(asyncio.TimeoutError), \ self.assertRunUnder(0.7): await con.close(timeout=0.5) + self.assertTrue(terminate_sent.is_set()) self.assertTrue(con._transport.is_closing()) finally: con.terminate()