ignore_after: patch Task.cancel to separate timeout vs external cancel

fixes kyuupichan#44

aiorpcx/curio.py

@@ -324,9 +324,16 @@ class UncaughtTimeoutError(Exception):
 
 def _set_new_deadline(task, deadline):
     def timeout_task():
-        # Unfortunately task.cancel is all we can do with asyncio
-        task.cancel()
-        task._timed_out = deadline
+        task._orig_cancel()
+        task._timed_out = None if getattr(task, "_externally_cancelled", False) else deadline
+    def mycancel(*args, **kwargs):
+        task._orig_cancel(*args, **kwargs)
+        task._externally_cancelled = True
+        task._timed_out = None
+    # We monkey-patch Task.cancel to distinguish timeouts and external cancellation:
+    if not hasattr(task, "_orig_cancel"):
+        task._orig_cancel = task.cancel
+        task.cancel = mycancel
     task._deadline_handle = task._loop.call_at(deadline, timeout_task)
 
 

tests/test_curio.py

@@ -1112,6 +1112,38 @@ async def test_ignore_at_context():
         assert False
 
 
+# See https://github.com/kyuupichan/aiorpcX/issues/44
+@pytest.mark.asyncio
+async def test_ignore_after_cancellation_race():
+    failed_evt = Event()
+
+    async def kill_group():
+        await sleep(0.01)
+        raise Exception("kill_group")
+
+    async def second_timeout():
+        await sleep(0.2)
+        failed_evt.set()  # mark unit test as failed
+
+    async def infinite_loop():
+        while True:
+            async with ignore_after(0):
+                await failed_evt.wait()
+                break
+
+    async with TaskGroup() as group:
+        await group.spawn(infinite_loop)
+        await group.spawn(kill_group)
+        # kill_group is supposed to kill this group soon. The unit test is testing this.
+        # If that fails, to avoid infinite-looping and to fail the test:
+        task_second_timeout = await spawn(second_timeout)  # top-level spawn
+    exc = group.exception
+    assert exc and exc.args and exc.args[0] == "kill_group"
+    assert not failed_evt.is_set()
+    # clean-up
+    task_second_timeout.cancel()
+
+
 #
 # Task group tests snitched from curio
 #