Bound close(timeout) around a pending cancel

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)
 

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;

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)