Merge pull request #1061 from oremanj/notice-cancel-misnesting

Attempt to recover from cancel scope nesting violations

Author: njsmith

newsfragments/882.feature.rst

@@ -0,0 +1,5 @@
+Trio now gives a reasonable traceback and error message in most cases
+when its invariants surrounding cancel scope nesting have been
+violated. (One common source of such violations is an async generator
+that yields within a cancel scope.) The previous behavior was an
+inscrutable chain of TrioInternalErrors.

trio/_core/_run.py

@@ -173,7 +173,8 @@ class CancelStatus:
     # The CancelStatus whose cancellations can propagate to us; we
     # become effectively cancelled when they do, unless scope.shield
     # is True.  May be None (for the outermost CancelStatus in a call
-    # to trio.run(), or briefly during TaskStatus.started()).
+    # to trio.run(), briefly during TaskStatus.started(), or during
+    # recovery from mis-nesting of cancel scopes).
     _parent = attr.ib(default=None, repr=False)
 
     # All of the CancelStatuses that have this CancelStatus as their parent.
@@ -185,6 +186,12 @@ class CancelStatus:
     # Invariant: all(task._cancel_status is self for task in self._tasks)
     _tasks = attr.ib(factory=set, init=False, repr=False)
 
+    # Set to True on still-active cancel statuses that are children
+    # of a cancel status that's been closed. This is used to permit
+    # recovery from mis-nested cancel scopes (well, at least enough
+    # recovery to show a useful traceback).
+    abandoned_by_misnesting = attr.ib(default=False, init=False, repr=False)
+
     def __attrs_post_init__(self):
         if self._parent is not None:
             self._parent._children.add(self)
@@ -213,9 +220,43 @@ def children(self):
     def tasks(self):
         return frozenset(self._tasks)
 
+    def encloses(self, other):
+        """Returns true if this cancel status is a direct or indirect
+        parent of cancel status *other*, or if *other* is *self*.
+        """
+        while other is not None:
+            if other is self:
+                return True
+            other = other.parent
+        return False
+
     def close(self):
-        assert not self._tasks and not self._children
         self.parent = None  # now we're not a child of self.parent anymore
+        if self._tasks or self._children:
+            # Cancel scopes weren't exited in opposite order of being
+            # entered. CancelScope._close() deals with raising an error
+            # if appropriate; our job is to leave things in a reasonable
+            # state for unwinding our dangling children. We choose to leave
+            # this part of the CancelStatus tree unlinked from everyone
+            # else, cancelled, and marked so that exiting a CancelScope
+            # within the abandoned subtree doesn't affect the active
+            # CancelStatus. Note that it's possible for us to get here
+            # without CancelScope._close() raising an error, if a
+            # nursery's cancel scope is closed within the nursery's
+            # nested child and no other cancel scopes are involved,
+            # but in that case task_exited() will deal with raising
+            # the error.
+            self._mark_abandoned()
+
+            # Since our CancelScope is about to forget about us, and we
+            # have no parent anymore, there's nothing left to call
+            # recalculate(). So, we can stay cancelled by setting
+            # effectively_cancelled and updating our children.
+            self.effectively_cancelled = True
+            for task in self._tasks:
+                task._attempt_delivery_of_any_pending_cancel()
+            for child in self._children:
+                child.recalculate()
 
     @property
     def parent_cancellation_is_visible_to_us(self):
@@ -237,6 +278,11 @@ def recalculate(self):
             for child in self._children:
                 child.recalculate()
 
+    def _mark_abandoned(self):
+        self.abandoned_by_misnesting = True
+        for child in self._children:
+            child._mark_abandoned()
+
     def effective_deadline(self):
         if self.effectively_cancelled:
             return -inf
@@ -245,6 +291,30 @@ def effective_deadline(self):
         return min(self._scope.deadline, self._parent.effective_deadline())
 
 
+MISNESTING_ADVICE = """
+This is probably a bug in your code, that has caused Trio's internal state to
+become corrupted. We'll do our best to recover, but from now on there are
+no guarantees.
+
+Typically this is caused by one of the following:
+  - yielding within a generator or async generator that's opened a cancel
+    scope or nursery (unless the generator is a @contextmanager or
+    @asynccontextmanager); see https://github.com/python-trio/trio/issues/638
+  - manually calling __enter__ or __exit__ on a trio.CancelScope, or
+    __aenter__ or __aexit__ on the object returned by trio.open_nursery();
+    doing so correctly is difficult and you should use @[async]contextmanager
+    instead, or maybe [Async]ExitStack
+  - using [Async]ExitStack to interleave the entries/exits of cancel scopes
+    and/or nurseries in a way that couldn't be achieved by some nesting of
+    'with' and 'async with' blocks
+  - using the low-level coroutine object protocol to execute some parts of
+    an async function in a different cancel scope/nursery context than
+    other parts
+If you don't believe you're doing any of these things, please file a bug:
+https://github.com/python-trio/trio/issues/new
+"""
+
+
 @attr.s(cmp=False, repr=False, slots=True)
 class CancelScope(metaclass=Final):
     """A *cancellation scope*: the link between a unit of cancellable
@@ -319,13 +389,60 @@ def _exc_filter(self, exc):
         return exc
 
     def _close(self, exc):
+        if self._cancel_status is None:
+            new_exc = RuntimeError(
+                "Cancel scope stack corrupted: attempted to exit {!r} "
+                "which had already been exited".format(self)
+            )
+            new_exc.__context__ = exc
+            return new_exc
         scope_task = current_task()
+        if scope_task._cancel_status is not self._cancel_status:
+            # Cancel scope mis-nesting: this cancel scope isn't the most
+            # recently opened by this task (that's still open). That is,
+            # our assumptions about context managers forming a stack
+            # have been violated. Try and make the best of it.
+            if self._cancel_status.abandoned_by_misnesting:
+                # We are an inner cancel scope that was still active when
+                # some outer scope was closed. The closure of that outer
+                # scope threw an error, so we don't need to throw another
+                # one; it would just confuse the traceback.
+                pass
+            elif not self._cancel_status.encloses(scope_task._cancel_status):
+                # This task isn't even indirectly contained within the
+                # cancel scope it's trying to close. Raise an error
+                # without changing any state.
+                new_exc = RuntimeError(
+                    "Cancel scope stack corrupted: attempted to exit {!r} "
+                    "from unrelated {!r}\n{}".format(
+                        self, scope_task, MISNESTING_ADVICE
+                    )
+                )
+                new_exc.__context__ = exc
+                return new_exc
+            else:
+                # Otherwise, there's some inner cancel scope(s) that
+                # we're abandoning by closing this outer one.
+                # CancelStatus.close() will take care of the plumbing;
+                # we just need to make sure we don't let the error
+                # pass silently.
+                new_exc = RuntimeError(
+                    "Cancel scope stack corrupted: attempted to exit {!r} "
+                    "in {!r} that's still within its child {!r}\n{}".format(
+                        self, scope_task, scope_task._cancel_status._scope,
+                        MISNESTING_ADVICE
+                    )
+                )
+                new_exc.__context__ = exc
+                exc = new_exc
+                scope_task._activate_cancel_status(self._cancel_status.parent)
+        else:
+            scope_task._activate_cancel_status(self._cancel_status.parent)
         if (
             exc is not None
             and not self._cancel_status.parent_cancellation_is_visible_to_us
         ):
             exc = MultiError.filter(self._exc_filter, exc)
-        scope_task._activate_cancel_status(self._cancel_status.parent)
         self._cancel_status.close()
         with self._might_change_registered_deadline():
             self._cancel_status = None
@@ -1155,6 +1272,29 @@ def _return_value_looks_like_wrong_library(value):
         return task
 
     def task_exited(self, task, outcome):
+        if (
+            task._cancel_status is not None
+            and task._cancel_status.abandoned_by_misnesting
+            and task._cancel_status.parent is None
+        ):
+            # The cancel scope surrounding this task's nursery was closed
+            # before the task exited. Force the task to exit with an error,
+            # since the error might not have been caught elsewhere. See the
+            # comments in CancelStatus.close().
+            try:
+                # Raise this, rather than just constructing it, to get a
+                # traceback frame included
+                raise RuntimeError(
+                    "Cancel scope stack corrupted: cancel scope surrounding "
+                    "{!r} was closed before the task exited\n{}".format(
+                        task, MISNESTING_ADVICE
+                    )
+                )
+            except RuntimeError as new_exc:
+                if isinstance(outcome, Error):
+                    new_exc.__context__ = outcome.error
+                outcome = Error(new_exc)
+
         task._activate_cancel_status(None)
         self.tasks.remove(task)
         if task is self.main_task:

trio/_core/tests/test_run.py

@@ -6,7 +6,7 @@
 import time
 import types
 import warnings
-from contextlib import contextmanager
+from contextlib import contextmanager, ExitStack
 from math import inf
 
 import attr
@@ -923,6 +923,76 @@ async def enter_scope():
     assert scope.cancel_called  # never become un-cancelled
 
 
+async def test_cancel_scope_misnesting():
+    outer = _core.CancelScope()
+    inner = _core.CancelScope()
+    with ExitStack() as stack:
+        stack.enter_context(outer)
+        with inner:
+            with pytest.raises(RuntimeError, match="still within its child"):
+                stack.close()
+        # No further error is raised when exiting the inner context
+
+    # If there are other tasks inside the abandoned part of the cancel tree,
+    # they get cancelled when the misnesting is detected
+    async def task1():
+        with pytest.raises(_core.Cancelled):
+            await sleep_forever()
+
+    # Even if inside another cancel scope
+    async def task2():
+        with _core.CancelScope():
+            with pytest.raises(_core.Cancelled):
+                await sleep_forever()
+
+    with ExitStack() as stack:
+        stack.enter_context(_core.CancelScope())
+        async with _core.open_nursery() as nursery:
+            nursery.start_soon(task1)
+            nursery.start_soon(task2)
+            await wait_all_tasks_blocked()
+            with pytest.raises(RuntimeError, match="still within its child"):
+                stack.close()
+
+    # Variant that makes the child tasks direct children of the scope
+    # that noticed the misnesting:
+    nursery_mgr = _core.open_nursery()
+    nursery = await nursery_mgr.__aenter__()
+    try:
+        nursery.start_soon(task1)
+        nursery.start_soon(task2)
+        nursery.start_soon(sleep_forever)
+        await wait_all_tasks_blocked()
+        nursery.cancel_scope.__exit__(None, None, None)
+    finally:
+        with pytest.raises(RuntimeError) as exc_info:
+            await nursery_mgr.__aexit__(*sys.exc_info())
+        assert "which had already been exited" in str(exc_info.value)
+        assert type(exc_info.value.__context__) is _core.MultiError
+        assert len(exc_info.value.__context__.exceptions) == 3
+        cancelled_in_context = False
+        for exc in exc_info.value.__context__.exceptions:
+            assert isinstance(exc, RuntimeError)
+            assert "closed before the task exited" in str(exc)
+            cancelled_in_context |= isinstance(
+                exc.__context__, _core.Cancelled
+            )
+        assert cancelled_in_context  # for the sleep_forever
+
+    # Trying to exit a cancel scope from an unrelated task raises an error
+    # without affecting any state
+    async def task3(task_status):
+        with _core.CancelScope() as scope:
+            task_status.started(scope)
+            await sleep_forever()
+
+    async with _core.open_nursery() as nursery:
+        scope = await nursery.start(task3)
+        with pytest.raises(RuntimeError, match="from unrelated"):
+            scope.__exit__(None, None, None)
+        scope.cancel()
+
+
 async def test_timekeeping():
     # probably a good idea to use a real clock for *one* test anyway...
     TARGET = 0.1