Can a future be cancelled when it hits an .await and is rescheduled by Tokio?

[mighty drift]

I have a piece of code that wraps a stream (cancellation safe) takes a message, processes it, makes 1 send call on a MPSC channel (where it .await's) and then returns the value to the caller. Now I traced a bug were I can see a new message is taken from the stream, but the function never returns the value.

Yet the caller is calling this function in a while loop with a tokio::select (where it awaits this call and an app shudown signal) and it doesn't produce any errors and the app isn't being shutdown.

So my only guess is that the future is somehow cancelled and the in flight message is lost in my function body while trying to call the send method.

Any insights or ways to debug/fix would be great!!

[civic canopy]

If your task is an axum request handler, then you might get cancelled if the connection is closed.

or similar cases like that

[mighty drift]

No the stream I read from is from a streamconsumer from rdkafka. The rest is "native" Rust (no frameworks)

[civic canopy]

Your task can definitely be cancelled at any .await, but it doesn't happen randomly. Cancellation only happens if you do so explicitly or if the runtime is shut down.

[mighty drift]

Hmm... Then it looks like it is not cancelled. But then I don't understand why the value doesn't get returned and at the same time it continues with processing new messages indicating this future should have been finished.

[mighty drift]

Ok, so I've put a atomic bool before and after the send.await in my wrapper function, setting the bool to true before the send call and back to false right after. And it indeed returns an error after running for a while that is already true before it is going to call the send method.

So its clearly either blocked and my wrapper method is called again (code wise this would be impossible) or cancelled. Any ideas how I can troubleshoot something like this?

Ow... I think I found it...

Oh man... Ok... I'm stupid.... 🙈

The issue is that I have a timer.tick() call in one of the branches of the select from where I call my wrapper. This is of course clearly the reason why the future gets cancelled! Its a bit of a long nasty select, but still stupid that I didn't see that any sooner.

But now how to prevent that? I guess I can't right? I cannot prevent a future from being cancelled at an await point, can I? Could I somehow do something with the stream of the channel to "glue" them together? Like conditionally taking the next message but only "commit" that I have the message once returned by my wrapper?

Probably not, but a bit in the dark about how to get around this one...

[civic canopy]

Can you share the code that is getting cancelled?

[mighty drift]

Yeah I can create a small snippet, hold on...

It's something like this:

    pub async fn recv(&self) -> Result<(BorrowedMessage<'_>, Offset), KafkaError> {
        tokio::select! {
            result = self.inner.recv() => match result {
                Ok(message) => {
                    let offset = Offset::from(&message);

                      self.offsets_tx
                        .send(OffsetMessage::InitOffset(offset.clone()))
                        .await
                        .map_err(|err| {
                            KafkaError::SetPartitionOffset(RDKafkaErrorCode::Fail)
                        })?;

                      return Ok((message, offset));
                }
                Err(err) => return Err(err),
            },
            _ = self.shutdown.signalled() => {
                return Err(KafkaError::Canceled);
            }
        }
    }

[civic canopy]

You could acquire a permit on offsets_tx before you receive. That way, there's guaranteed space in the channel and you can send without an await

[mighty drift]

And If there isn't an await it cannot be cancelled? Meaning that who ever is calling this will either the result of this future or the method isn't even polled. Of its polled, but cancelled before the inner.recv() finished... Do I understand all that correct?

Then in the end it's not so difficult to solve this as I thought it would be 😊

[civic canopy]

Cancellation can only happen at an await, yes.

[mighty drift]

Check... So the only interesting part here is that I should get the permit before the select call right. As otherwise I'll still be awaiting that call.

Think I got it, thanks (again!) for your help and insights 🙏

[civic canopy]

Well, you might still need signalled() to take effect during the reserve call

[mighty drift]

Yes, certainly. So it will be a second tokio select above the current one I figured.

[civic canopy]

You can keep one select! if you await on this:

async fn recv_with_permit(&self) -> Result<(Permit, Msg)> {
    let permit = self.offsets_tx.reserve().await?;
    let msg = self.inner.recv().await?;
    Ok((permit, msg))
}