Add a workaround for bpo-37658

`asyncio.wait_for()` currently has a bug where it raises a
`CancelledError` even when the wrapped awaitable has completed.
The upstream fix is in python/cpython#21894.  This adds a workaround
until the aforementioned PR is merged, backported and released.

Co-authored-by: Adam Liddell <git@aliddell.com>
Fixes: #467
Fixes: #547
Related: #468
Supersedes: #548

Author: elprans

asyncpg/compat.py

@@ -90,3 +90,19 @@ async def wait_closed(stream):
             # On Windows wait_closed() sometimes propagates
             # ConnectionResetError which is totally unnecessary.
             pass
+
+
+# Workaround for https://bugs.python.org/issue37658
+async def wait_for(fut, timeout):
+    if timeout is None:
+        return await fut
+
+    fut = asyncio.ensure_future(fut)
+
+    try:
+        return await asyncio.wait_for(fut, timeout)
+    except asyncio.CancelledError:
+        if fut.done():
+            return fut.result()
+        else:
+            raise

asyncpg/connect_utils.py

@@ -636,18 +636,13 @@ async def _connect_addr(
 
     connector = asyncio.ensure_future(connector)
     before = time.monotonic()
-    try:
-        tr, pr = await asyncio.wait_for(
-            connector, timeout=timeout)
-    except asyncio.CancelledError:
-        connector.add_done_callback(_close_leaked_connection)
-        raise
+    tr, pr = await compat.wait_for(connector, timeout=timeout)
     timeout -= time.monotonic() - before
 
     try:
         if timeout <= 0:
             raise asyncio.TimeoutError
-        await asyncio.wait_for(connected, timeout=timeout)
+        await compat.wait_for(connected, timeout=timeout)
     except (Exception, asyncio.CancelledError):
         tr.close()
         raise
@@ -745,12 +740,3 @@ def _create_future(loop):
         return asyncio.Future(loop=loop)
     else:
         return create_future()
-
-
-def _close_leaked_connection(fut):
-    try:
-        tr, pr = fut.result()
-        if tr:
-            tr.close()
-    except asyncio.CancelledError:
-        pass  # hide the exception

asyncpg/pool.py

@@ -12,6 +12,7 @@
 import time
 import warnings
 
+from . import compat
 from . import connection
 from . import connect_utils
 from . import exceptions
@@ -198,7 +199,7 @@ async def release(self, timeout):
                 # If the connection is in cancellation state,
                 # wait for the cancellation
                 started = time.monotonic()
-                await asyncio.wait_for(
+                await compat.wait_for(
                     self._con._protocol._wait_for_cancellation(),
                     budget)
                 if budget is not None:
@@ -623,7 +624,7 @@ async def _acquire_impl():
         if timeout is None:
             return await _acquire_impl()
         else:
-            return await asyncio.wait_for(
+            return await compat.wait_for(
                 _acquire_impl(), timeout=timeout)
 
     async def release(self, connection, *, timeout=None):

tests/test_pool.py

@@ -379,6 +379,26 @@ async def worker():
             self.cluster.trust_local_connections()
             self.cluster.reload()
 
+    async def test_pool_handles_task_cancel_in_acquire_with_timeout(self):
+        # See https://github.com/MagicStack/asyncpg/issues/547
+        pool = await self.create_pool(database='postgres',
+                                      min_size=1, max_size=1)
+
+        async def worker():
+            async with pool.acquire(timeout=100):
+                pass
+
+        # Schedule task
+        task = self.loop.create_task(worker())
+        # Yield to task, but cancel almost immediately
+        await asyncio.sleep(0.00000000001)
+        # Cancel the worker.
+        task.cancel()
+        # Wait to make sure the cleanup has completed.
+        await asyncio.sleep(0.4)
+        # Check that the connection has been returned to the pool.
+        self.assertEqual(pool._queue.qsize(), 1)
+
     async def test_pool_handles_task_cancel_in_release(self):
         # Use SlowResetConnectionPool to simulate
         # the Task.cancel() and __aexit__ race.