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..be8fa7e7 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,40 +612,41 @@ cdef class BaseProtocol(CoreProtocol): return self.closing = True - - 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 - - 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) - self._terminate() + close_waiter = None 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 + timeout = self._get_timeout_impl(timeout) + # Cancellation and the final disconnect share one deadline. + async with compat.timeout(timeout): + await self._drain_cancels() + + # Transport loss completes the cancellation futures too. + # There will be no further disconnect notification to await. + if self.con_status != CONNECTION_OK: + return + + 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 - self.transport.abort() + 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: + transport = self.transport + self.transport = None + transport.abort() def _request_cancel(self): self.cancel_waiter = self.create_future() @@ -686,6 +695,22 @@ 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 +720,7 @@ 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..4b176ea8 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 @@ -152,3 +155,141 @@ 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.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): + 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): + 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() + + 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(terminate_sent.is_set()) + 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)