Strengthening against external cancellation

tests/test_traio.py

@@ -162,6 +162,34 @@ async def run_me():
         raise Exception('should have raised')
 
 
+@pytest.mark.asyncio
+async def test_external_cancel_nasty():
+    """Raise an external cancellation with task which fails cancelling"""
+    async def nasty():
+        try:
+            await asyncio.sleep(10)
+        except asyncio.CancelledError:
+            # Something bad happens
+            raise ValueError('boom')
+
+    async def run_me():
+        n = Nursery()
+        async with n:
+            n.start_soon(nasty())
+
+        assert n.state == State.CANCELLED
+
+    task = asyncio.ensure_future(run_me())
+
+    try:
+        await asyncio.wait_for(task, 0.1)
+    except asyncio.TimeoutError:
+        with pytest.raises(asyncio.CancelledError):
+            await task
+    else:
+        raise Exception('should have raised')
+
+
 @pytest.mark.asyncio
 async def test_internal_cancel():
     """Test an internal cancellation"""
@@ -186,17 +214,6 @@ async def test_cancel_not_started():
         n.cancel()
 
 
-@pytest.mark.asyncio
-async def test_cancel_already_finished():
-    """Cancel a finished nursery"""
-    n = Nursery()
-    async with n:
-        pass
-
-    with pytest.raises(AssertionError):
-        n.cancel()
-
-
 @pytest.mark.asyncio
 async def test_master():
     """Test if a master correctly cancels pending tasks"""
@@ -286,3 +303,20 @@ async def trivial():
     after = time.time()
 
     assert 0.1 < (after - before) < 0.5
+
+
+@pytest.mark.asyncio
+async def test_cancelling_going_bad():
+    """Test cancelling a pending task, but things go wrong..."""
+    async def nasty():
+        try:
+            await asyncio.sleep(10)
+        except asyncio.CancelledError:
+            # Something bad happens
+            raise ValueError('boom')
+
+    with pytest.raises(TimeoutError):
+        async with Nursery(timeout=0.5) as n:
+            n.start_soon(nasty())
+
+    await asyncio.sleep(0.1)

traio/nursery.py

@@ -129,7 +129,7 @@ def cancel(self, exception: Exception = None):
         and raise given Exception if needed.
         :param exception: Exception to be raised
         """
-        assert self.state == State.STARTED, 'can only cancel a running nursery'
+        assert self.state != State.INIT, 'cannot cancel before even starting'
 
         for task in self._pending_tasks:
             if not task.done():
@@ -181,7 +181,13 @@ async def join(self):
             # We may still have pending tasks if the Nursery is cancelled
             for task in self._pending_tasks:
                 if not task.done():
-                    await task.ensure_cancelled()
+                    try:
+                        await task.ensure_cancelled()
+                    except Exception as ex:  # pylint: disable=broad-except
+                        # Too late for raising... and we need to move on cleaning other tasks!
+                        self.logger.error(
+                            'task `%s` failed to cancelled with exception: %s %s',
+                            task, ex.__class__.__name__, ex)
 
     def start_soon(
             self, awaitable: Awaitable, *,