Async wait_for shim raise cancellation right away

Deferring the cancellation makes it much harder to reason about and might make
the Python driver behave in a surprising way when cancelled.

Author: robsdedude

neo4j/_async/work/session.py

@@ -90,7 +90,6 @@ def __init__(self, pool, session_config):
         super().__init__(pool, session_config)
         assert isinstance(session_config, SessionConfig)
         self._bookmarks = self._prepare_bookmarks(session_config.bookmarks)
-        self._cancelled = False
 
     async def __aenter__(self):
         return self
@@ -140,7 +139,6 @@ def _collect_bookmark(self, bookmark):
             self._bookmarks = bookmark,
 
     def _handle_cancellation(self, message="General"):
-        self._cancelled = True
         self._transaction = None
         self._auto_result = None
         connection = self._connection
@@ -178,7 +176,6 @@ async def close(self):
         This will release any borrowed resources, such as connections, and will
         roll back any outstanding transactions.
         """
-        # if self._closed or self._cancelled:
         if self._closed:
             return
         if self._connection:

neo4j/_async_compat/shims/__init__.py

@@ -93,9 +93,8 @@ async def wait_for(fut, timeout):
                     # We got cancelled, but we are already done. Therefore,
                     # we defer the cancellation until next time the task yields
                     # to the event loop.
-                    asyncio.current_task().cancel()
+                    raise
                     # [/PATCH]
-                    return fut.result()
                 else:
                     fut.remove_done_callback(cb)
                     # We must ensure that the task is not running

neo4j/_sync/work/session.py

@@ -90,7 +90,6 @@ def __init__(self, pool, session_config):
         super().__init__(pool, session_config)
         assert isinstance(session_config, SessionConfig)
         self._bookmarks = self._prepare_bookmarks(session_config.bookmarks)
-        self._cancelled = False
 
     def __enter__(self):
         return self
@@ -140,7 +139,6 @@ def _collect_bookmark(self, bookmark):
             self._bookmarks = bookmark,
 
     def _handle_cancellation(self, message="General"):
-        self._cancelled = True
         self._transaction = None
         self._auto_result = None
         connection = self._connection
@@ -178,7 +176,6 @@ def close(self):
         This will release any borrowed resources, such as connections, and will
         roll back any outstanding transactions.
         """
-        # if self._closed or self._cancelled:
         if self._closed:
             return
         if self._connection:

tests/integration/mixed/test_async_cancellation.py

@@ -48,7 +48,7 @@ async def work(tx, i=1):
             assert isinstance(summary, neo4j.ResultSummary)
             assert len(records) == 1
             assert list(records[0]) == [i]
-        except asyncio.CancelledError as e:
+        except asyncio.CancelledError:
             work_cancelled = True
             raise
 

tests/unit/mixed/async_compat/test_concurrency.py

@@ -176,9 +176,14 @@ async def test_async_r_lock_acquire_cancellation(waits):
     async def acquire_task():
         while True:
             count = lock._count
-            await lock.acquire(timeout=0.1)
-            assert lock._count == count + 1
             try:
+                await lock.acquire(timeout=0.1)
+                assert lock._count == count + 1
+            except asyncio.CancelledError:
+                assert lock._count == count
+                raise
+            try:
+                # we're also ok with a deferred cancellation
                 await asyncio.sleep(0)
             except asyncio.CancelledError:
                 raise