[RFC] Refactor cancellation for great justice

[oremanj]

Relevant to #886, #606, #285, #147, #70, #58, maybe others.

I was continuing my effort to shoehorn linked cancel scopes and graceful cancellation into CancelScope earlier today and it was feeling too much of a mess, so I decided to explore other options. This PR is the result. It makes major changes to Trio's cancellation internals, but barely any to Trio's cancellation semantics -- all tests pass except for one that is especially persnickety about cancel_called. No new tests or docs yet as I wanted to get feedback on the approach before polishing.

An overview:

• New class CancelBinding manages a single lexical context (a with block or a task) that might get a different cancellation treatment than its surroundings. "All plumbing, no policy."
• Each cancel binding has an effective deadline, a single task, and links to parent and child bindings. Each parent lexically encloses its children. The only cancel bindings with multiple children are the ones immediately surrounding nurseries, and they have one child binding per nursery child task plus maybe one in the nested child.
• Each cancel binding calculates its effective deadline based on its parent's effective deadline and some additional data. The actual calculation is performed by an associated CancelLogic instance (a small ABC).
• CancelScope now implements CancelLogic, providing the deadline/shield semantics we know and love. It manages potentially-multiple CancelBindings.
• Cancel stacks are gone. Instead, each task has an "active" (innermost) cancel binding, which changes as the task moves in and out of cancellation regions. The active cancel binding's effective deadline directly determines whether and when Cancelled is raised in the task.
• Runner.deadlines stores tasks instead of cancel scopes. There is no longer a meaningful state of "deadline is in the past but scope isn't cancelled yet" (this is what the sole failing test doesn't like). If the effective deadline of a task's active cancel binding is non-infinite and in the future, it goes in Runner.deadlines. If it's in the past, the task has a pending cancellation by definition.

Potential advantages:

• Cancellation becomes extensible without changes to _core, via users writing their own CancelLogic and wrapping a core CancelBinding(s) around it. We could even move CancelScope out of _core if we want to make a point.
• Nursery.start() is much simpler now.
• Splitting shielding into a separate object from cancellation becomes trivial if we decide to do that (they'd be two kinds of CancelLogic).
• Most operations that are performed frequently take constant time: checking whether you're cancelled, checking what your deadline is, entering and leaving a cancel binding. I haven't benchmarked, so it's possible we're losing on constant factors or something, but in theory this should be faster than the old approach.
• Since tasks now have well-defined root cancel bindings, I think Propagating cancellation through threads #606 becomes straightforward - just provide a way to spawn a system task whose cancel binding is a child of something other than the system nursery's cancel binding.

Caveats:

• We call current_time() rather a lot. Not sure if this is worth worrying about, and could probably be cached if so.
• All hell breaks loose if time moves backwards.
• Losing the deadline-in-past vs cancelled distinction makes "once it gets cancelled, it stays cancelled" harder to reason about / ensure. We provide a weaker form of it, by marking the level we expect a Cancelled exception to reach before we throw it in, and being willing to repeat that process (raising it out further) if the landscape changes before it gets there.
• I'm not thrilled about "cancel binding" as a name - other suggestions are welcome.
• There are probably bugs, because aren't there always?

Current core cancel logic expressed in the new system:

    def compute_effective_deadline(
        self, parent_effective_deadline, parent_extra_info, task
    ):
        incoming_deadline = inf if self._shield else parent_effective_deadline
        my_deadline = -inf if self._cancel_called else self._deadline
        return min(incoming_deadline, my_deadline), parent_extra_info

I think the extension to support a grace period would just be:

    def compute_effective_deadline(
        self, parent_effective_deadline, parent_extra_info, task
    ):
        parent_cleanup_deadline = parent_extra_info.get("effective_cleanup_deadline", parent_effective_deadline)
        if self._shield:
            parent_effective_deadline = parent_cleanup_deadline = inf
        my_cleanup_start = min(self._deadline, self._cancel_called_at)
        merged_cleanup_deadline = min(parent_cleanup_deadline, my_cleanup_start + self._grace_period)
        my_extra_info = parent_extra_info.set("effective_cleanup_deadline", merged_cleanup_deadline)
        if self._shield_during_cleanup:
            effective_deadline = merged_cleanup_deadline
        else:
            effective_deadline = min(parent_effective_deadline, my_cleanup_start)
        return effective_deadline, my_extra_info

Maybe that's not quite simple but it is miles better than what I was looking at before. :-)

[oremanj]

Some more thoughts on this after sleeping on it:

• The API definitely needs some work but I still think the core ideas are sound.
• Probably rename (CancelBinding, CancelScope) to (CancelScope, CancelSource), and leave CancelSource out of the first version, following the discussion on Cancel scopes with disjoint extents #886. (Remaining comments here will use the updated names, i.e., a with block that provides cancellation semantics is a CancelScope.)
• Think more about whether and how to re-strengthen "once cancelled, always cancelled", and in general make more guarantees/invariants about when a cancellation is actually irreversibly delivered.
• CancelLogic needs to provide a "can cancellations propagate through here?" flag. In current Trio, if you have

with move_on_at(current_time() - 2) as outer:
    with CancelScope(shield=True, deadline=current_time() - 1) as inner:
        await checkpoint()

then the shield means the cancellation gets caught by inner, not outer, even though both inner and outer are cancelled. In this PR as it currently stands, the cancellation would be caught by outer.

[njsmith]

As usual for my first comment on these large diffs, I have spent a bunch of time thinking about the high-level idea but haven't actually read the diff, so apologies if anything I say is way off.

It took me a while to wrap my head around it, but I think the basic idea of having separate objects that track the aggregate cancel status for each extent is pretty good. But, I think it's probably simpler and more performant to keep deadlines separate from the cancel status. So what I'm imagining is something like:

• We have a tree of objects, let's call them CancelStatus. Every piece of code in Trio has a unique CancelStatus object associated with it. The only state we need for this is a single attribute on the task object. No need for a stack. Their main purpose is to hold a single bit: is this code cancelled or not?

• They maintain this bit via a message-passing algorithm: when someone calls cancel_scope.cancel(), then the relevant CancelStatus object flips its bit to "cancelled", and then it sends a message to all of its immediate children informing them of this. In general they respond by flipping their bits as well and then sending another message to their children, but not necessarily (e.g. if there's a cancel shield in effect).

• At least for now, there are only two actual "behaviors" a CancelStatus can have: a regular one propagates cancellation downward. A shielding one... doesn't.

• These are totally hidden from users. We'll expose CancelScope (and maybe CancelShield?) objects, and they internally manage CancelStatus objects – manipulating their state, switching which one is active in __enter__ and __exit__, etc. Nurseries will also need an associated CancelStatus, which is used as the initial CancelStatus for tasks inside them.

• Deadline handling works much the way it does now: we keep table of CancelScopes with active deadlines, and at appropriate moments these trigger calls to cancel_scope.cancel(). Which then updates the assocated CancelStatus, etc., as usual.

There are two reasons I like this overall approach:

• It seems pretty simple to reason about. I always imagine cancel scopes as like, "shining a light" down the task tree, which might get intercepted by a shield. If we add a "soft cancel" state, then that's like a different colored light, that can pass through a lens and transform into a "hard cancel", etc. And that's pretty much exactly what a message-passing approach simulates. It's like a flood-fill algorithm. Also, having expired deadlines call .cancel() avoids all discomfort you noticed.

• It seems about as efficient as we can hope for? Right now we do these O(cancel stack depth) operations on every checkpoint, which is silly. OTOH, a good thing about the current code is that each time the user sets a deadline, that only adds one entry to the global deadline table. If we track the effective deadline of each CancelStatus separately, then we potentially have a lot more entries in the table, which is pretty expensive. (Networking libraries often spend huge amounts of effort optimizing the cost of putting a deadline in their table and then removing it again, because most deadlines never fire, and tracking them is expensive.) So in my proposal, we keep the simple deadline behavior, so setting/modifying a deadline remains O(1), but the per-checkpoint cost becomes O(1) and trivial (if task._cancel_status.cancel_requested: ...).

  We do seem to lose in one place: when a scope is cancelled, then currently we can find all the affected tasks in O(affected tasks), but with this version finding them requires O(affected tasks * cancel stack depth). But... I think this is an illusion. First, because the only reason we can find all the affected tasks that quickly is because for each cancel scope we maintain a set of all the tasks that have it on their stack... and we have to update that set every time a new task is created, so maintaining the sets is O(affected tasks * cancel stack depth), which we pay all the time, while in this version creating a new task is O(1) and we only pay the O(cancel stack depth) when a cancel actually happens. And secondly, in the current code, after you find the tasks... for each one you have to do an O(cancel stack depth) operation to figure out whether the cancellation actually affects it (checking for shields). In the proposed version, this cost is absorbed into the cost of finding the tasks in the first place.

  Overall, I think suspect this proposal is just about optimal algorithmically, though I haven't done a detailed analysis.

What do you think of this, at a high level?

A few more thoughts:

Here's an idea for managing deadline entries in the global table: we only put CancelScopes in there while they are (a) bound, (b) not cancelled, (c) have a non-trivial deadline. In particular, by leaving unbound scopes out, we avoid some pathological cases where user code drops a CancelScope on the floor without binding it, but it can't be GC'ed because it's pinned by the deadline table. (And I seriously don't want to try to invent a WeakTimerWheel or something.) But, problem: what if someone peeks at cancel_requested before they're bound? It feels weird for the deadline to not to get checked at all until they're bound. Answer: make cancel_called a property that goes ahead and checks the deadline before returning. Combined with #909, this should make it seamless. (And if a cancel scope has already been exited, then I think it makes sense to leave cancel_called as False forever, even if the deadline passes, because that matches reality – cancel was indeed never called.)

Cancellation becomes extensible without changes to _core, via users writing their own CancelLogic and wrapping a core CancelBinding(s) around it. We could even move CancelScope out of _core if we want to make a point.

I hope at this point you've heard enough of my design philosophy that you won't be too surprised when I say this sounds like a negative to me, not a positive :-). (But whatever, we don't have to expose it even if it is nice and clean.)

Nursery.start() is much simpler now.

Yes. Hallelujah.

[oremanj]

Many thanks for the detailed feedback!

As usual for my first comment on these large diffs, I have spent a bunch of time thinking about the high-level idea but haven't actually read the diff, so apologies if anything I say is way off.

That was exactly what I was hoping for in this case - the code is still pretty hacked together at this point. The system you've detailed is uncannily close to the one I was imagining / tried to implement, and many of the differences (especially wrt tracking cancelled-ness explicitly as opposed to well-I-guess-this-deadline-is-in-the-past-ly) are changes I was planning to make as well. (Your CancelStatus is my CancelBinding, and I like your name better.) I'm glad I wasn't barking entirely up the wrong tree!

We have a tree of objects, let's call them CancelStatus. Every piece of code in Trio has a unique CancelStatus object associated with it. The only state we need for this is a single attribute on the task object. No need for a stack. Their main purpose is to hold a single bit: is this code cancelled or not?

I think tracking graceful cancellation will require more state than this. For example:

with CancelScope() as outer:
    outer.cancel(grace_period=1)
    with shield_during_cleanup():
        await trio.sleep(0.5)
    with shield_during_cleanup():
        # this sleep() needs to wake up in 0.5 seconds
        await trio.sleep(1)

So far I've been imagining each cancel status tracking a tuple (effective deadline, effective cleanup deadline), and the type of cancel status inserted by shield_during_cleanup would set its effective deadline to its parent's effective cleanup deadline. But for that to work, the CancelStatus objects have to know about the deadlines. Your point about not incurring an explosion of entries in Runner.deadlines is well-taken, but I think if we go all the way down to "one entry per CancelScope" we start having to worry again about the same cancel scope needing to get multiple deadline callbacks (one to wake up the children that are not soft-shielded, another after the grace period to wake up the children that are soft-shielded).

Apart from that detail I really like this approach, and I agree with you about the simplicity of reasoning and the algorithmic probably-ideality.

Here's an idea for managing deadline entries in the global table: we only put CancelScopes in there while they are (a) bound, (b) not cancelled, (c) have a non-trivial deadline.

This is the same thing we're doing for CancelScope now, isn't it? With the current unbound cancel scope code, we do have the oddity that cancel_called doesn't become True until the with block is entered even if the deadline already expired, so we should probably implement your make-cancel_called-a-property suggestion there too. But I'm not sure if any of this is new problems stemming from the new organization or not.

Cancellation becomes extensible without changes to _core

I hope at this point you've heard enough of my design philosophy that you won't be too surprised when I say this sounds like a negative to me, not a positive :-)

I suspected you would say that, but I had to try :P. I'm fine not exposing the flexibility; I think even having the internal architecture be flexible will make it a lot easier to prototype improvements to the cancellation system in the future.

[njsmith]

So far I've been imagining each cancel status tracking a tuple (effective deadline, effective cleanup deadline), and the type of cancel status inserted by shield_during_cleanup would set its effective deadline to its parent's effective cleanup deadline. But for that to work, the CancelStatus objects have to know about the deadlines. Your point about not incurring an explosion of entries in Runner.deadlines is well-taken, but I think if we go all the way down to "one entry per CancelScope" we start having to worry again about the same cancel scope needing to get multiple deadline callbacks (one to wake up the children that are not soft-shielded, another after the grace period to wake up the children that are soft-shielded).

Well, this is borrowing trouble, since we don't know if we'll implement soft cancels yet :-). But my intuition would be: if we do, then we extend CancelStatus to track a tri-state UNCANCELLED, SOFT_CANCELLED, HARD_CANCELLED. And if we also have separate soft and hard deadlines, then I think letting CancelStatus register two wakeups in the table is fine?

This is the same thing we're doing for CancelScope now, isn't it? With the current unbound cancel scope code, we do have the oddity that cancel_called doesn't become True until the with block is entered even if the deadline already expired, so we should probably implement your make-cancel_called-a-property suggestion there too. But I'm not sure if any of this is new problems stemming from the new organization or not.

Yeah, I guess I just thought of it because with the PR's approach, using the deadline itself to track cancellation, then you sort of automagically get the deadline/cancel-status in sync all the time, so I was wondering if there was any way to keep that on the outside while using simpler state tracking internally.

I suspected you would say that, but I had to try :P. I'm fine not exposing the flexibility; I think even having the internal architecture be flexible will make it a lot easier to prototype improvements to the cancellation system in the future.

Yeah, for sure.

[oremanj]

Closing this in favor of the more polished and mild-mannered #958.