Prevent stuck cancellation from blocking connection teardown

Use connection command timeout budget when waiting for cancellation and
abort the connection if cancellation cannot finish in time.

Fixes: #1326
Closes: #1349

Author: elprans

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):

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:

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
 

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: ...