Fix #216 (#223)

* Fix #216

* Make test_cancelled_connection_is_usable_asap more readable

Pass cur and fut explicitly as arguments, so it is easier to see where they come from

* Test that cancelled connection is not usable until cancellation is done

Author: themylogin

aiopg/connection.py

@@ -113,6 +113,8 @@ def __init__(self, dsn, loop, timeout, waiter, echo, **kwargs):
         self._timeout = timeout
         self._waiter = waiter
         self._writing = False
+        self._cancelling = False
+        self._cancellation_waiter = None
         self._echo = echo
         self._notifies = asyncio.Queue(loop=loop)
         self._weakref = weakref.ref(self)
@@ -199,8 +201,14 @@ def _fatal_error(self, message):
 
     def _create_waiter(self, func_name):
         if self._waiter is not None:
-            raise RuntimeError('%s() called while another coroutine is '
-                               'already waiting for incoming data' % func_name)
+            if self._cancelling:
+                if not self._waiter.done():
+                    raise RuntimeError('%s() called while connection is '
+                                       'being cancelled' % func_name)
+            else:
+                raise RuntimeError('%s() called while another coroutine is '
+                                   'already waiting for incoming '
+                                   'data' % func_name)
         self._waiter = create_future(self._loop)
         return self._waiter
 
@@ -212,6 +220,8 @@ def _poll(self, waiter, timeout):
         @asyncio.coroutine
         def cancel():
             self._waiter = create_future(self._loop)
+            self._cancelling = True
+            self._cancellation_waiter = self._waiter
             self._conn.cancel()
             if not self._conn.isexecuting():
                 return
@@ -231,7 +241,13 @@ def cancel():
         except psycopg2.extensions.QueryCanceledError:
             raise asyncio.CancelledError
         finally:
-            self._waiter = None
+            if self._cancelling:
+                self._cancelling = False
+                if self._waiter is self._cancellation_waiter:
+                    self._waiter = None
+                self._cancellation_waiter = None
+            else:
+                self._waiter = None
 
     def _isexecuting(self):
         return self._conn.isexecuting()

tests/test_connection.py

@@ -3,6 +3,7 @@
 import gc
 import psycopg2
 import psycopg2.extras
+import psycopg2.extensions
 import pytest
 import socket
 import time
@@ -287,6 +288,65 @@ def inner():
         yield from task
 
 
+@asyncio.coroutine
+def test_cancelled_connection_is_usable_asap(connect, loop):
+    @asyncio.coroutine
+    def inner(future, cursor):
+        future.set_result(None)
+        yield from cursor.execute("SELECT pg_sleep(10)")
+
+    fut = asyncio.Future(loop=loop)
+    conn = yield from connect()
+    cur = yield from conn.cursor()
+    task = ensure_future(inner(fut, cur), loop=loop)
+    yield from fut
+    yield from asyncio.sleep(0.1, loop=loop)
+
+    task.cancel()
+
+    for tick in range(100):
+        yield from asyncio.sleep(0, loop=loop)
+        status = conn._conn.get_transaction_status()
+        if status == psycopg2.extensions.TRANSACTION_STATUS_IDLE:
+            cur = yield from conn.cursor()
+            yield from cur.execute("SELECT 1")
+            ret = yield from cur.fetchone()
+            assert (1,) == ret
+            break
+    else:
+        assert False, "Cancelled connection transaction status never got idle"
+
+
+@asyncio.coroutine
+def test_cancelled_connection_is_not_usable_until_cancellation(connect, loop):
+    @asyncio.coroutine
+    def inner(future, cursor):
+        future.set_result(None)
+        yield from cursor.execute("SELECT pg_sleep(10)")
+
+    fut = asyncio.Future(loop=loop)
+    conn = yield from connect()
+    cur = yield from conn.cursor()
+    task = ensure_future(inner(fut, cur), loop=loop)
+    yield from fut
+    yield from asyncio.sleep(0.1, loop=loop)
+
+    task.cancel()
+
+    for i in range(100):
+        yield from asyncio.sleep(0)
+        if conn._cancelling:
+            break
+    else:
+        assert False, "Connection did not start cancelling"
+
+    cur = yield from conn.cursor()
+    with pytest.raises(RuntimeError) as e:
+        yield from cur.execute('SELECT 1')
+    assert str(e.value) == ('cursor.execute() called while connection '
+                            'is being cancelled')
+
+
 @asyncio.coroutine
 def test_close2(connect, loop):
     conn = yield from connect()