Nursery: don't act as a cancellation point in __aexit__

catern · catern · commit a3d9d2536a4f · 2020-09-01T11:48:44.000-04:00
Previously, a Nursery could raise Cancelled even if all its children had completed successfully. This is undesirable, as discussed in #1457, so now a Nursery will shield itself from cancels in __aexit__. A Nursery with children can still be cancelled fine: if any of its children raise Cancelled, then the whole Nursery will raise Cancelled. If none of the Nursery's children raise Cancelled, then the Nursery will never raise Cancelled. As another consequence, a Nursery which never had any child tasks will never raise Cancelled. As yet another consequence, since an internal Nursery is used in the implementation of Nursery.start(), if start() is cancelled, but the task it's starting does not raise Cancelled, start() won't raise Cancelled either.
diff --git a/trio/_core/_run.py b/trio/_core/_run.py
@@ -916,21 +916,25 @@ async def _nested_child_finished(self, nested_child_exc):
         self._check_nursery_closed()
 
         if not self._closed:
-            # If we get cancelled (or have an exception injected, like
-            # KeyboardInterrupt), then save that, but still wait until our
-            # children finish.
+            # If we have an exception injected, like KeyboardInterrupt,
+            # then save that, but still wait until our children finish.
             def aborted(raise_cancel):
-                self._add_exc(capture(raise_cancel).error)
+                exn = capture(raise_cancel).error
+                # We ignore Cancelled, since we should get it again
+                # through our children, if they were actually cancelled
+                # rather than successfully completing.
+                if not isinstance(exn, Cancelled):
+                    self._add_exc(exn)
                 return Abort.FAILED
 
             self._parent_waiting_in_aexit = True
             await wait_task_rescheduled(aborted)
         else:
-            # Nothing to wait for, so just execute a checkpoint -- but we
-            # still need to mix any exception (e.g. from an external
-            # cancellation) in with the rest of our exceptions.
+            # Nothing to wait for, so execute a schedule point, but don't
+            # allow us to be cancelled, just like the other branch.  We
+            # still need to catch and store non-Cancelled exceptions.
             try:
-                await checkpoint()
+                await cancel_shielded_checkpoint()
             except BaseException as exc:
                 self._add_exc(exc)
 
diff --git a/trio/_core/tests/test_run.py b/trio/_core/tests/test_run.py
@@ -32,6 +32,7 @@
     Sequencer,
     assert_checkpoints,
 )
+from ...testing._checkpoints import assert_yields_or_not
 
 
 # slightly different from _timeouts.sleep_forever because it returns the value
@@ -433,12 +434,12 @@ async def crasher():
                 # nursery block exited, all cancellations inside the
                 # nursery block continue propagating to reach the
                 # outer scope.
-                assert len(multi_exc.exceptions) == 5
+                assert len(multi_exc.exceptions) == 4
                 summary = {}
                 for exc in multi_exc.exceptions:
                     summary.setdefault(type(exc), 0)
                     summary[type(exc)] += 1
-                assert summary == {_core.Cancelled: 4, KeyError: 1}
+                assert summary == {_core.Cancelled: 3, KeyError: 1}
                 raise
     except AssertionError:  # pragma: no cover
         raise
@@ -1605,20 +1606,15 @@ async def test_trivial_yields():
         await _core.checkpoint_if_cancelled()
         await _core.cancel_shielded_checkpoint()
 
-    with assert_checkpoints():
+    with assert_yields_or_not(expect_cancel_point=False, expect_schedule_point=True):
         async with _core.open_nursery():
             pass
 
     with _core.CancelScope() as cancel_scope:
         cancel_scope.cancel()
-        with pytest.raises(_core.MultiError) as excinfo:
+        with pytest.raises(KeyError):
             async with _core.open_nursery():
                 raise KeyError
-        assert len(excinfo.value.exceptions) == 2
-        assert {type(e) for e in excinfo.value.exceptions} == {
-            KeyError,
-            _core.Cancelled,
-        }
 
 
 async def test_nursery_start(autojump_clock):
@@ -1685,13 +1681,24 @@ async def nothing(task_status=_core.TASK_STATUS_IGNORED):
     # is ignored; start() raises Cancelled.
     async def just_started(task_status=_core.TASK_STATUS_IGNORED):
         task_status.started("hi")
+        await _core.checkpoint()
 
     async with _core.open_nursery() as nursery:
         with _core.CancelScope() as cs:
             cs.cancel()
             with pytest.raises(_core.Cancelled):
                 await nursery.start(just_started)
 
+    # but if the task does not execute any checkpoints, and exits, then start()
+    # doesn't raise Cancelled, since the task completed successfully.
+    async def started_with_no_checkpoint(task_status=_core.TASK_STATUS_IGNORED):
+        task_status.started("hi")
+
+    async with _core.open_nursery() as nursery:
+        with _core.CancelScope() as cs:
+            cs.cancel()
+            await nursery.start(started_with_no_checkpoint)
+
     # and if after the no-op started(), the child crashes, the error comes out
     # of start()
     async def raise_keyerror_after_started(task_status=_core.TASK_STATUS_IGNORED,):
@@ -1701,12 +1708,8 @@ async def raise_keyerror_after_started(task_status=_core.TASK_STATUS_IGNORED,):
     async with _core.open_nursery() as nursery:
         with _core.CancelScope() as cs:
             cs.cancel()
-            with pytest.raises(_core.MultiError) as excinfo:
+            with pytest.raises(KeyError):
                 await nursery.start(raise_keyerror_after_started)
-    assert {type(e) for e in excinfo.value.exceptions} == {
-        _core.Cancelled,
-        KeyError,
-    }
 
     # trying to start in a closed nursery raises an error immediately
     async with _core.open_nursery() as closed_nursery:
