Add a test for Pool.release() interaction with Connection.cancel()

Author: elprans

asyncpg/_testbase/__init__.py

@@ -243,9 +243,11 @@ def get_connection_spec(cls, kwargs={}):
                 conn_spec['user'] = 'postgres'
         return conn_spec
 
-    def create_pool(self, pool_class=pg_pool.Pool, **kwargs):
+    def create_pool(self, pool_class=pg_pool.Pool,
+                    connection_class=pg_connection.Connection, **kwargs):
         conn_spec = self.get_connection_spec(kwargs)
-        return create_pool(loop=self.loop, pool_class=pool_class, **conn_spec)
+        return create_pool(loop=self.loop, pool_class=pool_class,
+                           connection_class=connection_class, **conn_spec)
 
     @classmethod
     def connect(cls, **kwargs):

asyncpg/connection.py

@@ -1052,47 +1052,47 @@ async def _cleanup_stmts(self):
             # so we ignore the timeout.
             await self._protocol.close_statement(stmt, protocol.NO_TIMEOUT)
 
+    async def _cancel(self, waiter):
+        r = w = None
+
+        try:
+            # Open new connection to the server
+            r, w = await connect_utils._open_connection(
+                loop=self._loop, addr=self._addr, params=self._params)
+
+            # Pack CancelRequest message
+            msg = struct.pack('!llll', 16, 80877102,
+                              self._protocol.backend_pid,
+                              self._protocol.backend_secret)
+
+            w.write(msg)
+            await r.read()  # Wait until EOF
+        except ConnectionResetError as ex:
+            # On some systems Postgres will reset the connection
+            # after processing the cancellation command.
+            if r is None and not waiter.done():
+                waiter.set_exception(ex)
+        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.
+            pass
+        except Exception as ex:
+            if not waiter.done():
+                waiter.set_exception(ex)
+        finally:
+            self._cancellations.discard(
+                asyncio.Task.current_task(self._loop))
+            if not waiter.done():
+                waiter.set_result(None)
+            if w is not None:
+                w.close()
+
     def _cancel_current_command(self, waiter):
-        async def cancel():
-            r = w = None
-
-            try:
-                # Open new connection to the server
-                r, w = await connect_utils._open_connection(
-                    loop=self._loop, addr=self._addr, params=self._params)
-
-                # Pack CancelRequest message
-                msg = struct.pack('!llll', 16, 80877102,
-                                  self._protocol.backend_pid,
-                                  self._protocol.backend_secret)
-
-                w.write(msg)
-                await r.read()  # Wait until EOF
-            except ConnectionResetError as ex:
-                # On some systems Postgres will reset the connection
-                # after processing the cancellation command.
-                if r is None and not waiter.done():
-                    waiter.set_exception(ex)
-            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.
-                pass
-            except Exception as ex:
-                if not waiter.done():
-                    waiter.set_exception(ex)
-            finally:
-                self._cancellations.discard(
-                    asyncio.Task.current_task(self._loop))
-                if not waiter.done():
-                    waiter.set_result(None)
-                if w is not None:
-                    w.close()
-
-        self._cancellations.add(self._loop.create_task(cancel()))
+        self._cancellations.add(self._loop.create_task(self._cancel(waiter)))
 
     def _process_log_message(self, fields, last_query):
         if not self._log_listeners:

asyncpg/pool.py

@@ -200,7 +200,7 @@ async def release(self, timeout):
                     started = time.monotonic()
                     await asyncio.wait_for(
                         self._con._protocol._wait_for_cancellation(),
-                        timeout, loop=self._pool._loop)
+                        budget, loop=self._pool._loop)
                     if budget is not None:
                         budget -= time.monotonic() - started
 

tests/test_pool.py

@@ -36,10 +36,11 @@ async def reset(self):
         return await super().reset()
 
 
-class SlowResetConnectionPool(pg_pool.Pool):
-    async def _connect(self, *args, **kwargs):
-        return await pg_connection.connect(
-            *args, connection_class=SlowResetConnection, **kwargs)
+class SlowCancelConnection(pg_connection.Connection):
+    """Connection class to simulate races with Connection._cancel()."""
+    async def _cancel(self, waiter):
+        await asyncio.sleep(0.2, loop=self._loop)
+        return await super()._cancel(waiter)
 
 
 class TestPool(tb.ConnectedTestCase):
@@ -351,12 +352,12 @@ async def worker():
             self.cluster.trust_local_connections()
             self.cluster.reload()
 
-    async def test_pool_handles_cancel_in_release(self):
+    async def test_pool_handles_task_cancel_in_release(self):
         # Use SlowResetConnectionPool to simulate
         # the Task.cancel() and __aexit__ race.
         pool = await self.create_pool(database='postgres',
                                       min_size=1, max_size=1,
-                                      pool_class=SlowResetConnectionPool)
+                                      connection_class=SlowResetConnection)
 
         async def worker():
             async with pool.acquire():
@@ -372,6 +373,27 @@ async def worker():
         # Check that the connection has been returned to the pool.
         self.assertEqual(pool._queue.qsize(), 1)
 
+    async def test_pool_handles_query_cancel_in_release(self):
+        # Use SlowResetConnectionPool to simulate
+        # the Task.cancel() and __aexit__ race.
+        pool = await self.create_pool(database='postgres',
+                                      min_size=1, max_size=1,
+                                      connection_class=SlowCancelConnection)
+
+        async def worker():
+            async with pool.acquire() as con:
+                await con.execute('SELECT pg_sleep(10)')
+
+        task = self.loop.create_task(worker())
+        # Let the worker() run.
+        await asyncio.sleep(0.1, loop=self.loop)
+        # Cancel the worker.
+        task.cancel()
+        # Wait to make sure the cleanup has completed.
+        await asyncio.sleep(0.5, loop=self.loop)
+        # Check that the connection has been returned to the pool.
+        self.assertEqual(pool._queue.qsize(), 1)
+
     async def test_pool_no_acquire_deadlock(self):
         async with self.create_pool(database='postgres',
                                     min_size=1, max_size=1,