Repository: python-trio/trio

Issue:
Reduce or eliminate quadratic behavior in cancellation system?
Currently, checking for cancellation is O(n), where n is the depth of the task's cancel stack. We do this check frequently, including every time a task yields and in `yield_if_cancelled`. This isn't *so* bad because generally cancel stacks are not terribly deep, but it is worrisome. (Note that the stack is always at least as deep as the supervision tree, so while a good practice is to only have one explicit timeout at the top of your code, that doesn't mean cancellation stacks are only 1 deep.)

This isn't really urgent until we start seeing cancellation-related operations showing up in profiles, but I suspect that eventually it might be worth reworking this code to improve the asymptotics. It's a bit tricky given the richness of our semantics (in particular shielding), but we can certainly do better than we are right now.

A simple change would be to have a per-task cache that says where the task is cancelled, and recompute it (a) any time one of the cancel scopes is cancelled, (b) any time one of the cancel states shielding state is mutated, (c) any time the cancel stack is popped. (I guess we can slightly optimize further by observing that `shield = True` cannot create a cancellation where there wasn't one before, and vice-versa for `shield = False`, and popping can never create a cancellation. But I suspect it's very rare for shields to be used for multiple tasks, or to be set/unset very often, so it doesn't much matter how fast that operation is -- the big win, if any, will be from reducing the overhead on yield. And popping is always restricted to just one task.)


Discussion:

Resolving patch:
diff --git a/trio/_core/_run.py b/trio/_core/_run.py
--- a/trio/_core/_run.py
+++ b/trio/_core/_run.py
@@ -124,6 +124,126 @@ def deadline_to_sleep_time(self, deadline):
 ################################################################
 
 
+@attr.s(cmp=False, slots=True)
+class CancelStatus:
+    """Tracks the cancellation status for a contiguous extent
+    of code that will become cancelled, or not, as a unit.
+
+    Each task has at all times a single "active" CancelStatus whose
+    cancellation state determines whether checkpoints executed in that
+    task raise Cancelled. Each 'with CancelScope(...)' context is
+    associated with a particular CancelStatus.  When a task enters
+    such a context, a CancelStatus is created which becomes the active
+    CancelStatus for that task; when the 'with' block is exited, the
+    active CancelStatus for that task goes back to whatever it was
+    before.
+
+    CancelStatus objects are arranged in a tree whose structure
+    mirrors the lexical nesting of the cancel scope contexts.  When a
+    CancelStatus becomes cancelled, it notifies all of its direct
+    children, who become cancelled in turn (and continue propagating
+    the cancellation down the tree) unless they are shielded. (There
+    will be at most one such child except in the case of a
+    CancelStatus that immediately encloses a nursery.) At the leaves
+    of this tree are the tasks themselves, which get woken up to deliver
+    an abort when their direct parent CancelStatus becomes cancelled.
+
+    You can think of CancelStatus as being responsible for the
+    "plumbing" of cancellations as oppposed to CancelScope which is
+    responsible for the origination of them.
+
+    """
+
+    # Our associated cancel scope. Can be any object with attributes
+    # `deadline`, `shield`, and `cancel_called`, but in current usage
+    # is always a CancelScope object. Must not be None.
+    _scope = attr.ib()
+
+    # True iff the tasks in self._tasks should receive cancellations
+    # when they checkpoint. Always True when scope.cancel_called is True;
+    # may also be True due to a cancellation propagated from our
+    # parent.  Unlike scope.cancel_called, this does not necessarily stay
+    # true once it becomes true. For example, we might become
+    # effectively cancelled due to the cancel scope two levels out
+    # becoming cancelled, but then the cancel scope one level out
+    # becomes shielded so we're not effectively cancelled anymore.
+    effectively_cancelled = attr.ib(default=False)
+
+    # 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()).
+    _parent = attr.ib(default=None, repr=False)
+
+    # All of the CancelStatuses that have this CancelStatus as their parent.
+    _children = attr.ib(factory=set, init=False, repr=False)
+
+    # Tasks whose cancellation state is currently tied directly to
+    # the cancellation state of this CancelStatus object. Don't modify
+    # this directly; instead, use Task._activate_cancel_status().
+    # Invariant: all(task._cancel_status is self for task in self._tasks)
+    _tasks = attr.ib(factory=set, init=False, repr=False)
+
+    def __attrs_post_init__(self):
+        if self._parent is not None:
+            self._parent._children.add(self)
+            self.recalculate()
+
+    # parent/children/tasks accessors are used by TaskStatus.started()
+
+    @property
+    def parent(self):
+        return self._parent
+
+    @parent.setter
+    def parent(self, parent):
+        if self._parent is not None:
+            self._parent._children.remove(self)
+        self._parent = parent
+        if self._parent is not None:
+            self._parent._children.add(self)
+            self.recalculate()
+
+    @property
+    def children(self):
+        return frozenset(self._children)
+
+    @property
+    def tasks(self):
+        return frozenset(self._tasks)
+
+    def close(self):
+        assert not self._tasks and not self._children
+        self.parent = None  # now we're not a child of self.parent anymore
+
+    @property
+    def parent_cancellation_is_visible_to_us(self):
+        return (
+            self._parent is not None and not self._scope.shield
+            and self._parent.effectively_cancelled
+        )
+
+    def recalculate(self):
+        new_state = (
+            self._scope.cancel_called
+            or self.parent_cancellation_is_visible_to_us
+        )
+        if new_state != self.effectively_cancelled:
+            self.effectively_cancelled = new_state
+            if new_state:
+                for task in self._tasks:
+                    task._attempt_delivery_of_any_pending_cancel()
+            for child in self._children:
+                child.recalculate()
+
+    def effective_deadline(self):
+        if self.effectively_cancelled:
+            return -inf
+        if self._parent is None or self._scope.shield:
+            return self._scope.deadline
+        return min(self._scope.deadline, self._parent.effective_deadline())
+
+
 @attr.s(cmp=False, repr=False, slots=True)
 class CancelScope:
     """A *cancellation scope*: the link between a unit of cancellable
@@ -164,10 +284,10 @@ class CancelScope:
     has been entered yet, and changes take immediate effect.
     """
 
-    _tasks = attr.ib(factory=set, init=False)
+    _cancel_status = attr.ib(default=None, init=False)
     _has_been_entered = attr.ib(default=False, init=False)
-    _effective_deadline = attr.ib(default=inf, init=False)
-    cancel_called = attr.ib(default=False, init=False)
+    _registered_deadline = attr.ib(default=inf, init=False)
+    _cancel_called = attr.ib(default=False, init=False)
     cancelled_caught = attr.ib(default=False, init=False)
 
     # Constructor arguments:
@@ -183,11 +303,33 @@ def __enter__(self):
             )
         self._has_been_entered = True
         if current_time() >= self._deadline:
-            self.cancel_called = True
-        with self._might_change_effective_deadline():
-            self._add_task(task)
+            self.cancel()
+        with self._might_change_registered_deadline():
+            self._cancel_status = CancelStatus(
+                scope=self, parent=task._cancel_status
+            )
+            task._activate_cancel_status(self._cancel_status)
         return self
 
+    def _exc_filter(self, exc):
+        if isinstance(exc, Cancelled):
+            self.cancelled_caught = True
+            return None
+        return exc
+
+    def _close(self, exc):
+        scope_task = current_task()
+        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
+        return exc
+
     @enable_ki_protection
     def __exit__(self, etype, exc, tb):
         # NB: NurseryManager calls _close() directly rather than __exit__(),
@@ -213,18 +355,16 @@ def __exit__(self, etype, exc, tb):
                 value.__context__ = old_context
 
     def __repr__(self):
-        if self._tasks:
-            binding = "bound to {} task{}".format(
-                len(self._tasks), "s" if len(self._tasks) > 1 else ""
-            )
+        if self._cancel_status is not None:
+            binding = "active"
         elif self._has_been_entered:
             binding = "exited"
         else:
             binding = "unbound"
 
-        if self.cancel_called:
+        if self._cancel_called:
             state = ", cancelled"
-        elif self.deadline == inf:
+        elif self._deadline == inf:
             state = ""
         else:
             try:
@@ -233,8 +373,8 @@ def __repr__(self):
                 state = ""
             else:
                 state = ", deadline is {:.2f} seconds {}".format(
-                    abs(self.deadline - now),
-                    "from now" if self.deadline >= now else "ago"
+                    abs(self._deadline - now),
+                    "from now" if self._deadline >= now else "ago"
                 )
 
         return "<trio.CancelScope at {:#x}, {}{}>".format(
@@ -243,17 +383,17 @@ def __repr__(self):
 
     @contextmanager
     @enable_ki_protection
-    def _might_change_effective_deadline(self):
+    def _might_change_registered_deadline(self):
         try:
             yield
         finally:
-            old = self._effective_deadline
-            if self.cancel_called or not self._tasks:
+            old = self._registered_deadline
+            if self._cancel_status is None or self._cancel_called:
                 new = inf
             else:
                 new = self._deadline
             if old != new:
-                self._effective_deadline = new
+                self._registered_deadline = new
                 runner = GLOBAL_RUN_CONTEXT.runner
                 if old != inf:
                     del runner.deadlines[old, id(self)]
@@ -288,7 +428,7 @@ def deadline(self):
 
     @deadline.setter
     def deadline(self, new_deadline):
-        with self._might_change_effective_deadline():
+        with self._might_change_registered_deadline():
             self._deadline = float(new_deadline)
 
     @property
@@ -321,9 +461,8 @@ def shield(self, new_value):
         if not isinstance(new_value, bool):
             raise TypeError("shield must be a bool")
         self._shield = new_value
-        if not self._shield:
-            for task in self._tasks:
-                task._attempt_delivery_of_any_pending_cancel()
+        if self._cancel_status is not None:
+            self._cancel_status.recalculate()
 
     @enable_ki_protection
     def cancel(self):
@@ -332,50 +471,44 @@ def cancel(self):
         This method is idempotent, i.e., if the scope was already
         cancelled then this method silently does nothing.
         """
-        if self.cancel_called:
+        if self._cancel_called:
             return
-        with self._might_change_effective_deadline():
-            self.cancel_called = True
-        for task in self._tasks:
-            task._attempt_delivery_of_any_pending_cancel()
-
-    def _add_task(self, task):
-        self._tasks.add(task)
-        task._cancel_stack.append(self)
-
-    def _remove_task(self, task):
-        self._tasks.remove(task)
-        assert task._cancel_stack[-1] is self
-        task._cancel_stack.pop()
-
-    # Used by the nursery.start trickiness
-    def _tasks_removed_by_adoption(self, tasks):
-        with self._might_change_effective_deadline():
-            self._tasks.difference_update(tasks)
-
-    # Used by the nursery.start trickiness
-    def _tasks_added_by_adoption(self, tasks):
-        self._tasks.update(tasks)
-
-    def _exc_filter(self, exc):
-        if isinstance(exc, Cancelled):
-            self.cancelled_caught = True
-            return None
-        return exc
-
-    def _close(self, exc):
-        scope_task = current_task()
-        if (
-            exc is not None and self.cancel_called
-            and scope_task._pending_cancel_scope() is self
-        ):
-            exc = MultiError.filter(self._exc_filter, exc)
-        with self._might_change_effective_deadline():
-            self._remove_task(scope_task)
-        return exc
+        with self._might_change_registered_deadline():
+            self._cancel_called = True
+        if self._cancel_status is not None:
+            self._cancel_status.recalculate()
 
-
-@deprecated("0.10.0", issue=607, instead="trio.CancelScope")
+    @property
+    def cancel_called(self):
+        """Readonly :class:`bool`. Records whether cancellation has been
+        requested for this scope, either by an explicit call to
+        :meth:`cancel` or by the deadline expiring.
+
+        This attribute being True does *not* necessarily mean that the
+        code within the scope has been, or will be, affected by the
+        cancellation. For example, if :meth:`cancel` was called after
+        the last checkpoint in the ``with`` block, when it's too late to
+        deliver a :exc:`~trio.Cancelled` exception, then this attribute
+        will still be True.
+
+        This attribute is mostly useful for debugging and introspection.
+        If you want to know whether or not a chunk of code was actually
+        cancelled, then :attr:`cancelled_caught` is usually more
+        appropriate.
+        """
+        if self._cancel_status is not None or not self._has_been_entered:
+            # Scope is active or not yet entered: make sure cancel_called
+            # is true if the deadline has passed. This shouldn't
+            # be able to actually change behavior, since we check for
+            # deadline expiry on scope entry and at every checkpoint,
+            # but it makes the value returned by cancel_called more
+            # closely match expectations.
+            if not self._cancel_called and current_time() >= self._deadline:
+                self.cancel()
+        return self._cancel_called
+
+
+@deprecated("0.11.0", issue=607, instead="trio.CancelScope")
 def open_cancel_scope(*, deadline=inf, shield=False):
     """Returns a context manager which creates a new cancellation scope."""
     return CancelScope(deadline=deadline, shield=shield)
@@ -410,62 +543,43 @@ def started(self, value=None):
         # will eventually exit on its own, and we don't want to risk moving
         # children that might have propagating Cancelled exceptions into
         # a place with no cancelled cancel scopes to catch them.
-        if _pending_cancel_scope(self._old_nursery._cancel_stack) is not None:
+        if self._old_nursery._cancel_status.effectively_cancelled:
             return
 
         # Can't be closed, b/c we checked in start() and then _pending_starts
         # should keep it open.
         assert not self._new_nursery._closed
 
-        # otherwise, find all the tasks under the old nursery, and move them
-        # under the new nursery instead. This means:
-        # - changing parents of direct children
-        # - changing cancel stack of all direct+indirect children
-        # - changing cancel stack of all direct+indirect children's nurseries
-        # - checking for cancellation in all changed cancel stacks
-        old_stack = self._old_nursery._cancel_stack
-        new_stack = self._new_nursery._cancel_stack
-        # LIFO todo stack for depth-first traversal
-        todo = list(self._old_nursery._children)
-        munged_tasks = []
-        while todo:
-            task = todo.pop()
-            # Direct children need to be reparented
-            if task._parent_nursery is self._old_nursery:
-                self._old_nursery._children.remove(task)
-                task._parent_nursery = self._new_nursery
-                self._new_nursery._children.add(task)
-            # Everyone needs their cancel scopes fixed up...
-            assert task._cancel_stack[:len(old_stack)] == old_stack
-            task._cancel_stack[:len(old_stack)] = new_stack
-            # ...and their nurseries' cancel scopes fixed up.
-            for nursery in task._child_nurseries:
-                assert nursery._cancel_stack[:len(old_stack)] == old_stack
-                nursery._cancel_stack[:len(old_stack)] = new_stack
-                # And then add all the nursery's children to our todo list
-                todo.extend(nursery._children)
-            # And make a note to check for cancellation later
-            munged_tasks.append(task)
-
-        # Tell all the cancel scopes about the change. (There are probably
-        # some scopes in common between the two stacks, so some scopes will
-        # get the same tasks removed and then immediately re-added. This is
-        # fine though.)
-        for cancel_scope in old_stack:
-            cancel_scope._tasks_removed_by_adoption(munged_tasks)
-        for cancel_scope in new_stack:
-            cancel_scope._tasks_added_by_adoption(munged_tasks)
+        # Move tasks from the old nursery to the new
+        tasks = self._old_nursery._children
+        self._old_nursery._children = set()
+        for task in tasks:
+            task._parent_nursery = self._new_nursery
+            self._new_nursery._children.add(task)
+
+        # Move all children of the old nursery's cancel status object
+        # to be underneath the new nursery instead. This includes both
+        # tasks and child cancel status objects.
+        # NB: If the new nursery is cancelled, reparenting a cancel
+        # status to be underneath it can invoke an abort_fn, which might
+        # do something evil like cancel the old nursery. We thus break
+        # everything off from the old nursery before we start attaching
+        # anything to the new.
+        cancel_status_children = self._old_nursery._cancel_status.children
+        cancel_status_tasks = set(self._old_nursery._cancel_status.tasks)
+        cancel_status_tasks.discard(self._old_nursery._parent_task)
+        for cancel_status in cancel_status_children:
+            cancel_status.parent = None
+        for task in cancel_status_tasks:
+            task._activate_cancel_status(None)
+        for cancel_status in cancel_status_children:
+            cancel_status.parent = self._new_nursery._cancel_status
+        for task in cancel_status_tasks:
+            task._activate_cancel_status(self._new_nursery._cancel_status)
 
         # That should have removed all the children from the old nursery
         assert not self._old_nursery._children
 
-        # After all the delicate surgery is done, check for cancellation in
-        # all the tasks that had their cancel scopes munged. This can trigger
-        # arbitrary abort() callbacks, so we put it off until our internal
-        # data structures are all self-consistent again.
-        for task in munged_tasks:
-            task._attempt_delivery_of_any_pending_cancel()
-
         # And finally, poke the old nursery so it notices that all its
         # children have disappeared and can exit.
         self._old_nursery._check_nursery_closed()
@@ -533,13 +647,13 @@ class Nursery:
     def __init__(self, parent_task, cancel_scope):
         self._parent_task = parent_task
         parent_task._child_nurseries.append(self)
-        # the cancel stack that children inherit - we take a snapshot, so it
+        # the cancel status that children inherit - we take a snapshot, so it
         # won't be affected by any changes in the parent.
-        self._cancel_stack = list(parent_task._cancel_stack)
+        self._cancel_status = parent_task._cancel_status
         # the cancel scope that directly surrounds us; used for cancelling all
         # children.
         self.cancel_scope = cancel_scope
-        assert self.cancel_scope is self._cancel_stack[-1]
+        assert self.cancel_scope._cancel_status is self._cancel_status
         self._children = set()
         self._pending_excs = []
         # The "nested child" is how this code refers to the contents of the
@@ -643,19 +757,6 @@ def __del__(self):
 ################################################################
 
 
-def _pending_cancel_scope(cancel_stack):
-    # Return the outermost exception that is is not outside a shield.
-    pending_scope = None
-    for scope in cancel_stack:
-        # Check shield before _exc, because shield should not block
-        # processing of *this* scope's exception
-        if scope.shield:
-            pending_scope = None
-        if pending_scope is None and scope.cancel_called:
-            pending_scope = scope
-    return pending_scope
-
-
 @attr.s(cmp=False, hash=False, repr=False)
 class Task:
     _parent_nursery = attr.ib()
@@ -711,10 +812,18 @@ def child_nurseries(self):
     # Cancellation
     ################
 
-    _cancel_stack = attr.ib(default=attr.Factory(list), repr=False)
+    # The CancelStatus object that is currently active for this task.
+    # Don't change this directly; instead, use _activate_cancel_status().
+    _cancel_status = attr.ib(default=None, repr=False)
 
-    def _pending_cancel_scope(self):
-        return _pending_cancel_scope(self._cancel_stack)
+    def _activate_cancel_status(self, cancel_status):
+        if self._cancel_status is not None:
+            self._cancel_status._tasks.remove(self)
+        self._cancel_status = cancel_status
+        if self._cancel_status is not None:
+            self._cancel_status._tasks.add(self)
+            if self._cancel_status.effectively_cancelled:
+                self._attempt_delivery_of_any_pending_cancel()
 
     def _attempt_abort(self, raise_cancel):
         # Either the abort succeeds, in which case we will reschedule the
@@ -733,7 +842,7 @@ def _attempt_abort(self, raise_cancel):
     def _attempt_delivery_of_any_pending_cancel(self):
         if self._abort_func is None:
             return
-        if self._pending_cancel_scope() is None:
+        if not self._cancel_status.effectively_cancelled:
             return
 
         def raise_cancel():
@@ -1035,8 +1144,7 @@ def _return_value_looks_like_wrong_library(value):
 
         if nursery is not None:
             nursery._children.add(task)
-            for scope in nursery._cancel_stack:
-                scope._add_task(task)
+            task._activate_cancel_status(nursery._cancel_status)
 
         if self.instruments:
             self.instrument("task_spawned", task)
@@ -1046,8 +1154,7 @@ def _return_value_looks_like_wrong_library(value):
         return task
 
     def task_exited(self, task, outcome):
-        while task._cancel_stack:
-            task._cancel_stack[-1]._remove_task(task)
+        task._activate_cancel_status(None)
         self.tasks.remove(task)
         if task is self.main_task:
             self.main_task_outcome = outcome
@@ -1669,15 +1776,7 @@ def current_effective_deadline():
         float: the effective deadline, as an absolute time.
 
     """
-    task = current_task()
-    deadline = inf
-    for scope in task._cancel_stack:
-        if scope._shield:
-            deadline = inf
-        if scope.cancel_called:
-            deadline = -inf
-        deadline = min(deadline, scope._deadline)
-    return deadline
+    return current_task()._cancel_status.effective_deadline()
 
 
 async def checkpoint():
@@ -1715,7 +1814,7 @@ async def checkpoint_if_cancelled():
     """
     task = current_task()
     if (
-        task._pending_cancel_scope() is not None or
+        task._cancel_status.effectively_cancelled or
         (task is task._runner.main_task and task._runner.ki_pending)
     ):
         await _core.checkpoint()


Test patch:
diff --git a/trio/_core/tests/test_run.py b/trio/_core/tests/test_run.py
--- a/trio/_core/tests/test_run.py
+++ b/trio/_core/tests/test_run.py
@@ -546,11 +546,7 @@ async def test_cancel_scope_repr(mock_clock):
     scope = _core.CancelScope()
     assert "unbound" in repr(scope)
     with scope:
-        assert "bound to 1 task" in repr(scope)
-        async with _core.open_nursery() as nursery:
-            nursery.start_soon(sleep, 10)
-            assert "bound to 2 tasks" in repr(scope)
-            nursery.cancel_scope.cancel()
+        assert "active" in repr(scope)
         scope.deadline = _core.current_time() - 1
         assert "deadline is 1.00 seconds ago" in repr(scope)
         scope.deadline = _core.current_time() + 10
@@ -871,6 +867,16 @@ async def sleep_until_cancelled(scope):
         await wait_all_tasks_blocked()
         scope.cancel()
 
+    # Shield before entry
+    scope = _core.CancelScope()
+    scope.shield = True
+    with _core.CancelScope() as outer, scope:
+        outer.cancel()
+        await _core.checkpoint()
+        scope.shield = False
+        with pytest.raises(_core.Cancelled):
+            await _core.checkpoint()
+
     # Can't reuse
     with _core.CancelScope() as scope:
         await _core.checkpoint()
@@ -907,6 +913,15 @@ async def enter_scope():
         assert "single 'with' block" in str(exc_info.value)
         nursery.cancel_scope.cancel()
 
+    # If not yet entered, cancel_called is true when the deadline has passed
+    # even if cancel() hasn't been called yet
+    scope = _core.CancelScope(deadline=_core.current_time() + 1)
+    assert not scope.cancel_called
+    scope.deadline -= 1
+    assert scope.cancel_called
+    scope.deadline += 1
+    assert scope.cancel_called  # never become un-cancelled
+
 
 async def test_timekeeping():
     # probably a good idea to use a real clock for *one* test anyway...