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