CancellationToken: Option to only trigger "drop guard" if dropped while panicking

[jamesmunns]

Is your feature request related to a problem? Please describe.

In my own version of CancellationToken, I wanted to have something like a DropGuard, but had avoided it because I didn't want to manually manage defusing them.

I realized that we can check if we are panicking during the drop constructor, and use that to gate whether the drop should signal a cancellation. This is the case for many of my projects where I have many long-lived tasks, and want to trigger a graceful shutdown if any one of the "peers" fail.

Describe the solution you'd like

It could be useful for CancellationToken to have one or both of the following APIs:

• A method that sets a flag, like .cancel_on_panic(true) that "arms" this check, without requiring a wrapper type (like DropGuard)
• Another wrapper type (like PanicCancellationToken) that always performs this check.

Here's a proof of concept using non-tokio-parts, but I thought I'd share in case it was useful:

https://gist.github.com/jamesmunns/c92458fd54af690079ac5becf4a49bc7

[jamesmunns]

By the way, I'm happy to look into implementing this, it would be good to get any feedback on which (or both) of the approaches would be interesting.

[Darksonn]

Normally I think this is done using disarm

let guard = token.guard();
this_may_panic().await;
guard.disarm();

but I suppose we could add a dedicated PanicGuard if this is too hard to use.

[jamesmunns]

Yep, the goal is to make this easier to use in cases where you might have multiple returns, and don't necessarily want to handle this in every exit case.

async fn example(&self, token: CancellationToken) -> Result<()> {
    let guard = token.guard();
    loop {
        match self.next() {
            A(a) => {
                let is_good = a.may_panic();
                if !is_good {
                    guard.disarm();
                    return Err(...);
                }
            },
            B(b) => {
                // uh oh! ? here can return without disarming!
                let is_good = b.may_panic()?;
            },
            // ...
        }
    }
}

vs

async fn example(&self, mut token: CancellationToken) -> Result<()> {
    let guard = token.cancel_on_panic(true);
    loop {
        match self.next() {
            A(a) => {
                let is_good = a.may_panic();
                if !is_good {
                    return Err(...);
                }
            },
            B(b) => {
                // No panic? no problem!
                let is_good = b.may_panic()?;
            },
            // ...
        }
    }
}

[Darksonn]

Aha, so you don't want to cancel when exiting on error from a question mark?

[jamesmunns]

Yep! That's the idea, I only want to send a cancellation message if a panic occurs, because I am using the cancellation token for the task fan-out of long running tasks.

For example if I have long lived worker tasks, and one of them reaches a critical error, I want to signal all of the rest of them to shut down, because without all of the worker tasks operating, my application can't continue.

By making this happen on panic, I can just "panic quickly" if any fatal error occurs (like the database becomes unreachable).

[Darksonn]

Since this is tokio-util, the bar for adding stuff isn't as high as the core Tokio crate. I'm okay with adding a PanicGuard.

[jamesmunns]

Just as a check, I'm guessing this is preferred to adding a field to the CancellationToken?

If it were a field, it could be inherited on cloning or creating a child, and you wouldn't need to create panic guards everywhere you want to use them.

This could be done by adding:

pub struct CancellationToken {
    inner: Arc<tree_node::TreeNode>,
    panic_on_cancel: bool, // Would require &mut self to set
    // OR
    panic_on_cancel: AtomicBool, // Would only require &self to set
}

// ...

impl Drop for CancellationToken {
    fn drop(&mut self) {
        // NEW
        if self.panic_on_cancel && std::thread::panicking() {
            self.cancel();
        }
        // EXISTING
        tree_node::decrease_handle_refcount(&self.inner);
    }
}

[Darksonn]

I would definitely prefer a new guard type. Adding a boolean would make all tokens in use out there larger for no reason.