diff --git a/asyncpg/connect_utils.py b/asyncpg/connect_utils.py index c3a8fe8a..61eae0c0 100644 --- a/asyncpg/connect_utils.py +++ b/asyncpg/connect_utils.py @@ -1310,29 +1310,31 @@ def connection_lost(self, exc): if not self.on_disconnect.done(): self.on_disconnect.set_result(True) - if isinstance(addr, str): - tr, pr = await loop.create_unix_connection(CancelProto, addr) - else: - if params.ssl and params.sslmode != SSLMode.allow: - tr, pr = await _create_ssl_connection( - CancelProto, - *addr, - loop=loop, - ssl_context=params.ssl, - ssl_is_advisory=params.sslmode == SSLMode.prefer) + tr = None + try: + if isinstance(addr, str): + tr, pr = await loop.create_unix_connection(CancelProto, addr) else: - tr, pr = await loop.create_connection( - CancelProto, *addr) - _set_nodelay(_get_socket(tr)) + if params.ssl and params.sslmode != SSLMode.allow: + tr, pr = await _create_ssl_connection( + CancelProto, + *addr, + loop=loop, + ssl_context=params.ssl, + ssl_is_advisory=params.sslmode == SSLMode.prefer) + else: + tr, pr = await loop.create_connection( + CancelProto, *addr) + _set_nodelay(_get_socket(tr)) - # Pack a CancelRequest message - msg = struct.pack('!llll', 16, 80877102, backend_pid, backend_secret) + # Pack a CancelRequest message + msg = struct.pack('!llll', 16, 80877102, backend_pid, backend_secret) - try: tr.write(msg) await pr.on_disconnect finally: - tr.close() + if tr is not None: + tr.close() def _get_socket(transport): diff --git a/asyncpg/connection.py b/asyncpg/connection.py index c1363b6f..1939b158 100644 --- a/asyncpg/connection.py +++ b/asyncpg/connection.py @@ -1520,7 +1520,7 @@ async def close(self, *, timeout=None): def terminate(self): """Terminate the connection without waiting for pending data.""" - if not self.is_closed(): + if not self._aborted and self._protocol is not None: self._abort() self._cleanup() @@ -1573,8 +1573,9 @@ async def reset(self, *, timeout=None): def _abort(self): # Put the connection into the aborted state. self._aborted = True - self._protocol.abort() - self._protocol = None + if self._protocol is not None: + self._protocol.abort() + self._protocol = None def _cleanup(self): self._call_termination_listeners() @@ -1595,8 +1596,9 @@ def _cleanup(self): def _clean_tasks(self): # Wrap-up any remaining tasks associated with this connection. if self._cancellations: + current = asyncio.current_task(self._loop) for fut in self._cancellations: - if not fut.done(): + if fut is not current and not fut.done(): fut.cancel() self._cancellations.clear() @@ -1649,37 +1651,42 @@ async def _cleanup_stmts(self): # so we ignore the timeout. await self._protocol.close_statement(stmt, protocol.NO_TIMEOUT) - async def _cancel(self, waiter): + async def _cancel(self, waiter, cancel_waiter=None): try: - # Open new connection to the server - await connect_utils._cancel( - loop=self._loop, addr=self._addr, params=self._params, - backend_pid=self._protocol.backend_pid, - backend_secret=self._protocol.backend_secret) - except ConnectionResetError as ex: - # On some systems Postgres will reset the connection - # after processing the cancellation command. - if not waiter.done(): - waiter.set_exception(ex) + async with compat.timeout(self._config.command_timeout): + try: + await connect_utils._cancel( + loop=self._loop, addr=self._addr, params=self._params, + backend_pid=self._protocol.backend_pid, + backend_secret=self._protocol.backend_secret) + except ConnectionResetError: + # Some servers reset the auxiliary connection after + # receiving the CancelRequest. The original connection + # still has to acknowledge the cancelled query. + pass + + if not waiter.done(): + waiter.set_result(None) + if cancel_waiter is not None: + await asyncio.shield(cancel_waiter) except asyncio.CancelledError: - # There are two scenarios in which the cancellation - # itself will be cancelled: 1) the connection is being closed, - # 2) the event loop is being shut down. - # In either case we do not care about the propagation of - # the CancelledError, and don't want the loop to warn about - # an unretrieved exception. + # Teardown can cancel this background task. Its waiters are + # completed in finally, without leaking CancelledError. pass - except (Exception, asyncio.CancelledError) as ex: - if not waiter.done(): - waiter.set_exception(ex) + except Exception: + if not self._aborted: + # A failed CancelRequest leaves the original connection's + # protocol state uncertain. It cannot be reused safely. + self.terminate() finally: self._cancellations.discard( asyncio.current_task(self._loop)) if not waiter.done(): waiter.set_result(None) - def _cancel_current_command(self, waiter): - self._cancellations.add(self._loop.create_task(self._cancel(waiter))) + def _cancel_current_command(self, waiter, cancel_waiter=None): + self._cancellations.add(self._loop.create_task( + self._cancel(waiter, cancel_waiter))) def _process_log_message(self, fields, last_query): if not self._log_listeners: diff --git a/asyncpg/pool.py b/asyncpg/pool.py index c7d624dc..81c2ad3b 100644 --- a/asyncpg/pool.py +++ b/asyncpg/pool.py @@ -11,7 +11,6 @@ import functools import inspect import logging -import time from types import TracebackType from typing import Any, Optional, Type import warnings @@ -226,12 +225,17 @@ async def release(self, timeout: Optional[float]) -> None: if self._con._protocol._is_cancelling(): # If the connection is in cancellation state, # wait for the cancellation - started = time.monotonic() - await compat.wait_for( - self._con._protocol._wait_for_cancellation(), + budget = await self._con._protocol._wait_for_cancellation( budget) - if budget is not None: - budget -= time.monotonic() - started + + # The background cancellation may have timed out and terminated + # the connection while we were waiting. In that case cleanup + # has already returned the holder to the pool. + if self._con is None: + return + if self._con.is_closed(): + self._con.terminate() + return if self._pool._reset is not None: async with compat.timeout(budget): @@ -246,7 +250,8 @@ async def release(self, timeout: Optional[float]) -> None: try: # An exception in `reset` is most likely caused by # an IO error, so terminate the connection. - self._con.terminate() + if self._con is not None: + self._con.terminate() finally: raise ex diff --git a/asyncpg/protocol/protocol.pyi b/asyncpg/protocol/protocol.pyi index 3bed28df..a1ad67d7 100644 --- a/asyncpg/protocol/protocol.pyi +++ b/asyncpg/protocol/protocol.pyi @@ -176,7 +176,9 @@ class BaseProtocol(CoreProtocol, Generic[_Record]): async def close(self, timeout: _TimeoutType) -> None: ... def _get_timeout(self, timeout: _TimeoutType) -> float | None: ... def _is_cancelling(self) -> bool: ... - async def _wait_for_cancellation(self) -> None: ... + async def _wait_for_cancellation( + self, timeout: float | None = None + ) -> float | None: ... async def close_statement( self, state: PreparedStatementState[_OtherRecord], timeout: _TimeoutType ) -> Any: ... diff --git a/asyncpg/protocol/protocol.pyx b/asyncpg/protocol/protocol.pyx index 4e1494e7..7497dcec 100644 --- a/asyncpg/protocol/protocol.pyx +++ b/asyncpg/protocol/protocol.pyx @@ -141,14 +141,10 @@ cdef class BaseProtocol(CoreProtocol): PreparedStatementState state=None, ignore_custom_codec=False, record_class): - if self.cancel_waiter is not None: - await self.cancel_waiter - if self.cancel_sent_waiter is not None: - await self.cancel_sent_waiter - self.cancel_sent_waiter = None + timeout = self._get_timeout_impl(timeout) + timeout = await self._wait_for_cancellation(timeout) self._check_state() - timeout = self._get_timeout_impl(timeout) waiter = self._new_waiter(timeout) try: @@ -173,14 +169,10 @@ cdef class BaseProtocol(CoreProtocol): return_extra: bool, timeout, ): - if self.cancel_waiter is not None: - await self.cancel_waiter - if self.cancel_sent_waiter is not None: - await self.cancel_sent_waiter - self.cancel_sent_waiter = None + timeout = self._get_timeout_impl(timeout) + timeout = await self._wait_for_cancellation(timeout) self._check_state() - timeout = self._get_timeout_impl(timeout) args_buf = state._encode_bind_msg(args) waiter = self._new_waiter(timeout) @@ -212,14 +204,10 @@ cdef class BaseProtocol(CoreProtocol): timeout, return_rows: bool, ): - if self.cancel_waiter is not None: - await self.cancel_waiter - if self.cancel_sent_waiter is not None: - await self.cancel_sent_waiter - self.cancel_sent_waiter = None + timeout = self._get_timeout_impl(timeout) + timeout = await self._wait_for_cancellation(timeout) self._check_state() - timeout = self._get_timeout_impl(timeout) timer = Timer(timeout) # Make sure the argument sequence is encoded lazily with @@ -268,14 +256,10 @@ cdef class BaseProtocol(CoreProtocol): async def bind(self, PreparedStatementState state, args, str portal_name, timeout): - if self.cancel_waiter is not None: - await self.cancel_waiter - if self.cancel_sent_waiter is not None: - await self.cancel_sent_waiter - self.cancel_sent_waiter = None + timeout = self._get_timeout_impl(timeout) + timeout = await self._wait_for_cancellation(timeout) self._check_state() - timeout = self._get_timeout_impl(timeout) args_buf = state._encode_bind_msg(args) waiter = self._new_waiter(timeout) @@ -300,14 +284,10 @@ cdef class BaseProtocol(CoreProtocol): str portal_name, int limit, return_extra, timeout): - if self.cancel_waiter is not None: - await self.cancel_waiter - if self.cancel_sent_waiter is not None: - await self.cancel_sent_waiter - self.cancel_sent_waiter = None + timeout = self._get_timeout_impl(timeout) + timeout = await self._wait_for_cancellation(timeout) self._check_state() - timeout = self._get_timeout_impl(timeout) waiter = self._new_waiter(timeout) try: @@ -327,14 +307,10 @@ cdef class BaseProtocol(CoreProtocol): async def close_portal(self, str portal_name, timeout): - if self.cancel_waiter is not None: - await self.cancel_waiter - if self.cancel_sent_waiter is not None: - await self.cancel_sent_waiter - self.cancel_sent_waiter = None + timeout = self._get_timeout_impl(timeout) + timeout = await self._wait_for_cancellation(timeout) self._check_state() - timeout = self._get_timeout_impl(timeout) waiter = self._new_waiter(timeout) try: @@ -348,17 +324,10 @@ cdef class BaseProtocol(CoreProtocol): return await waiter async def query(self, query, timeout): - if self.cancel_waiter is not None: - await self.cancel_waiter - if self.cancel_sent_waiter is not None: - await self.cancel_sent_waiter - self.cancel_sent_waiter = None + timeout = self._get_timeout(timeout) + timeout = await self._wait_for_cancellation(timeout) self._check_state() - # query() needs to call _get_timeout instead of _get_timeout_impl - # for consistent validation, as it is called differently from - # prepare/bind/execute methods. - timeout = self._get_timeout(timeout) waiter = self._new_waiter(timeout) try: @@ -372,15 +341,11 @@ cdef class BaseProtocol(CoreProtocol): return await waiter async def copy_out(self, copy_stmt, sink, timeout): - if self.cancel_waiter is not None: - await self.cancel_waiter - if self.cancel_sent_waiter is not None: - await self.cancel_sent_waiter - self.cancel_sent_waiter = None + timeout = self._get_timeout_impl(timeout) + timeout = await self._wait_for_cancellation(timeout) self._check_state() - timeout = self._get_timeout_impl(timeout) timer = Timer(timeout) # The copy operation is guarded by a single timeout @@ -431,15 +396,11 @@ cdef class BaseProtocol(CoreProtocol): ssize_t num_cols Codec codec - if self.cancel_waiter is not None: - await self.cancel_waiter - if self.cancel_sent_waiter is not None: - await self.cancel_sent_waiter - self.cancel_sent_waiter = None + timeout = self._get_timeout_impl(timeout) + timeout = await self._wait_for_cancellation(timeout) self._check_state() - timeout = self._get_timeout_impl(timeout) timer = Timer(timeout) waiter = self._new_waiter(timer.get_remaining_budget()) @@ -563,11 +524,8 @@ cdef class BaseProtocol(CoreProtocol): return status_msg async def close_statement(self, PreparedStatementState state, timeout): - if self.cancel_waiter is not None: - await self.cancel_waiter - if self.cancel_sent_waiter is not None: - await self.cancel_sent_waiter - self.cancel_sent_waiter = None + timeout = self._get_timeout_impl(timeout) + timeout = await self._wait_for_cancellation(timeout) self._check_state() @@ -576,7 +534,6 @@ cdef class BaseProtocol(CoreProtocol): 'cannot close prepared statement; refs == {} != 0'.format( state.refs)) - timeout = self._get_timeout_impl(timeout) waiter = self._new_waiter(timeout) try: self._close(state.name, False) # network op @@ -624,7 +581,7 @@ cdef class BaseProtocol(CoreProtocol): # Transport loss completes the cancellation futures too. # There will be no further disconnect notification to await. - if self.con_status != CONNECTION_OK: + if self.con_status != CONNECTION_OK or self.transport is None: return assert self.waiter is None @@ -659,7 +616,8 @@ cdef class BaseProtocol(CoreProtocol): if con is not None: # if 'con' is None it means that the connection object has been # garbage collected and that the transport will soon be aborted. - con._cancel_current_command(self.cancel_sent_waiter) + con._cancel_current_command( + self.cancel_sent_waiter, self.cancel_waiter) else: self.loop.call_exception_handler({ 'message': 'asyncpg.Protocol has no reference to its ' @@ -770,12 +728,33 @@ cdef class BaseProtocol(CoreProtocol): self.cancel_sent_waiter is not None ) - async def _wait_for_cancellation(self): - 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 + async def _wait_for_cancellation(self, timeout=None): + if not self._is_cancelling(): + return timeout + + started = self.loop.time() + try: + async with compat.timeout(timeout): + if self.cancel_sent_waiter is not None: + await asyncio.shield(self.cancel_sent_waiter) + self.cancel_sent_waiter = None + if self.cancel_waiter is not None: + await asyncio.shield(self.cancel_waiter) + except asyncio.TimeoutError: + # The old query may still be running. A new query cannot use + # this connection unless the cancellation is acknowledged. + con = self.get_connection() + if con is not None: + con.terminate() + else: + self.abort() + raise + + if timeout is not None: + timeout -= self.loop.time() - started + if timeout <= 0: + raise asyncio.TimeoutError() + return timeout cdef _coreproto_error(self): try: diff --git a/tests/test_adversity.py b/tests/test_adversity.py index a6e03feb..1d061ee6 100644 --- a/tests/test_adversity.py +++ b/tests/test_adversity.py @@ -19,6 +19,28 @@ platform.system() == 'Windows', 'not compatible with ProactorEventLoop which is default in Python 3.8+') class TestConnectionLoss(tb.ProxiedClusterTestCase): + @tb.with_timeout(30.0) + async def test_cancel_request_times_out_during_network_loss(self): + con = await self.connect(command_timeout=0.2) + try: + self.proxy.trigger_connectivity_loss() + loss_started = asyncio.run_coroutine_threadsafe( + self.proxy.connectivity_loss.wait(), self.proxy.loop) + await asyncio.wait_for(asyncio.wrap_future(loss_started), 1) + + with self.assertRaises(asyncio.TimeoutError): + await con.execute('SELECT 1') + self.assertTrue(con._cancellations) + + async def wait_until_closed(): + while not con.is_closed(): + await asyncio.sleep(0.005) + + await asyncio.wait_for(wait_until_closed(), 1) + finally: + self.proxy.restore_connectivity() + con.terminate() + @tb.with_timeout(30.0) async def test_connection_close_timeout(self): con = await self.connect() diff --git a/tests/test_pool.py b/tests/test_pool.py index 0bc35ac0..f8000a59 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -14,9 +14,11 @@ import textwrap import time import unittest +from unittest import mock import asyncpg from asyncpg import _testbase as tb +from asyncpg import connect_utils from asyncpg import connection as pg_connection from asyncpg import pool as pg_pool from asyncpg import cluster as pg_cluster @@ -36,9 +38,9 @@ async def reset(self, *, timeout=None): class SlowCancelConnection(pg_connection.Connection): """Connection class to simulate races with Connection._cancel().""" - async def _cancel(self, waiter): + async def _cancel(self, waiter, cancel_waiter=None): await asyncio.sleep(0.2) - return await super()._cancel(waiter) + return await super()._cancel(waiter, cancel_waiter) class TestPool(tb.ConnectedTestCase): @@ -471,6 +473,47 @@ async def worker(): # Check that the connection has been returned to the pool. self.assertEqual(pool._queue.qsize(), 1) + async def test_pool_release_timeout_during_cancellation(self): + pool = await self.create_pool(database='postgres', + min_size=1, max_size=1) + con = await pool.acquire() + cancel_started = asyncio.Event() + + async def cancel(**kwargs): + cancel_started.set() + await self.loop.create_future() + + with mock.patch.object(connect_utils, '_cancel', cancel): + with self.assertRaises(asyncio.TimeoutError): + await con.execute('SELECT pg_sleep(10)', timeout=0.01) + await cancel_started.wait() + + with self.assertRaises(asyncio.TimeoutError): + await pool.release(con, timeout=0.05) + + async with pool.acquire(timeout=1) as replacement: + self.assertEqual(await replacement.fetchval('SELECT 1'), 1) + + async def test_pool_release_after_background_cancel_timeout(self): + pool = await self.create_pool(database='postgres', + min_size=1, max_size=1, + command_timeout=0.1) + con = await pool.acquire() + cancel_started = asyncio.Event() + + async def cancel(**kwargs): + cancel_started.set() + await self.loop.create_future() + + with mock.patch.object(connect_utils, '_cancel', cancel): + with self.assertRaises(asyncio.TimeoutError): + await con.execute('SELECT pg_sleep(10)', timeout=0.01) + await cancel_started.wait() + await asyncio.wait_for(pool.release(con), 1) + + async with pool.acquire(timeout=1) as replacement: + self.assertEqual(await replacement.fetchval('SELECT 1'), 1) + async def test_pool_no_acquire_deadlock(self): async with self.create_pool(database='postgres', min_size=1, max_size=1, diff --git a/tests/test_timeout.py b/tests/test_timeout.py index 4b176ea8..0bc8de8f 100644 --- a/tests/test_timeout.py +++ b/tests/test_timeout.py @@ -199,9 +199,12 @@ async def test_close_times_out_pending_cancel(self): 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() + with self.assertRunUnder(MAX_RUNTIME): + try: + await con.close() + except asyncio.TimeoutError: + # The close and background cancel deadlines can race. + pass self.assertTrue(con._transport.is_closing()) async def test_cancel_close_aborts_transport(self): @@ -293,3 +296,83 @@ async def cancel(**kwargs): con._transport.abort() task.cancel() await asyncio.gather(task, return_exceptions=True) + + async def test_next_command_timeout_includes_pending_cancel(self): + for cancel_sent in (False, True): + for method in ('execute', 'fetchval'): + with self.subTest(cancel_sent=cancel_sent, method=method): + async with self.pending_cancel( + cancel_sent=cancel_sent) as con: + with self.assertRaises(asyncio.TimeoutError), \ + self.assertRunUnder(MAX_RUNTIME): + await getattr(con, method)( + 'select 1', timeout=0.05) + self.assertTrue(con.is_closed()) + + async def test_cancel_attempt_has_one_deadline(self): + for cancel_sent in (False, True): + with self.subTest(cancel_sent=cancel_sent): + con = await self.connect(command_timeout=0.05) + started = asyncio.Event() + + async def cancel(**kwargs): + started.set() + if not cancel_sent: + await self.loop.create_future() + + try: + with mock.patch.object(connect_utils, '_cancel', cancel): + with self.assertRaises(asyncio.TimeoutError): + await con.execute( + 'select pg_sleep(10)', timeout=0.01) + await started.wait() + + async def wait_until_closed(): + while not con.is_closed(): + await asyncio.sleep(0.005) + + await asyncio.wait_for( + wait_until_closed(), MAX_RUNTIME) + self.assertTrue(con._transport.is_closing()) + finally: + con.terminate() + + async def test_cancel_attempt_expires_during_close(self): + con = await self.connect(command_timeout=0.05) + started = asyncio.Event() + + async def cancel(**kwargs): + started.set() + await self.loop.create_future() + + try: + with mock.patch.object(connect_utils, '_cancel', cancel): + with self.assertRaises(asyncio.TimeoutError): + await con.execute('select pg_sleep(10)', timeout=0.01) + await started.wait() + with self.assertRunUnder(0.3): + await con.close(timeout=0.1) + self.assertTrue(con.is_closed()) + finally: + con.terminate() + + async def test_failed_cancel_during_close_without_timeout(self): + con = await self.connect() + cancel_started = asyncio.Event() + + async def cancel(**kwargs): + cancel_started.set() + raise ConnectionRefusedError('cancel port unreachable') + + query = self.loop.create_task(con.execute('select pg_sleep(5)')) + try: + await asyncio.sleep(0.05) + self.assertFalse(query.done()) + with mock.patch.object(connect_utils, '_cancel', cancel): + await asyncio.wait_for(con.close(), MAX_RUNTIME) + self.assertTrue(cancel_started.is_set()) + self.assertTrue(con.is_closed()) + finally: + con.terminate() + query.cancel() + await asyncio.gather(query, return_exceptions=True)