Detect connection closed on endpoint with spawned task?

[robert-sjoblom]

Summary

Hi,

First of all, thank you for axum. It's mostly been a breeze working with it, and so far it's exceeded our expectations.

However, I have a case I can't solve in an elegant way. We have an API endpoint that runs for a long time and eventually returns a response. However, if the connection disconnects, the spawned task from the handler continues to run, while the handler task itself is cancelled.

Here's the setup (abbreviated):

async fn handler() -> Result<String, Error> {
  let (tx, rx) = oneshot::channel();

  tokio::spawn({
    async move { poll_until_ready(value, tx).await }
  });

  Ok(Json(wait_for_completion(rx).await?))
}

async fn poll_until_ready(values, tx) {
  loop {
    // loop until we finish or fail
  }
}

async fn wait_for_completion(rx) -> Result<String, Error> {
  let timeout = Duration::from_secs(60 * 30);
  match tokio::time::timeout(timeout, rx).await {
    Ok(Ok(received_val)) => received_val,
    Ok(Err(received_err)) => Err...,
    Err(elapsed) => Err...
  }
}

I couldn't find a good (readable) way to keep the code inside the handler, which is why I put the loop on a different thread and just await the message instead. Perhaps I'm solving it in a bad way, but I can't figure out how to handle the handler being dropped (the loop would cancel if it was in the handler function itself, of course).

I read #1094 but it doesn't really answer my use-case.

Grateful for any assistance,
Robert.

axum version

0.7.7

[mladedav]

I think #1094 is exactly what you should do in cases like this.

That said you could also check in your worker thread whether the channel was closed. If the handler finished, it dropped the receiver and you can check that through the sender.

[robert-sjoblom]

You're not wrong, I was blinded by my own initial solution. Thank you!

[jplatte]

What's the actual problem with calling poll_until_ready without a separate async task? If you definitely need the channel, how about using futures::future::join(poll_until_ready(value, tx), wait_for_completion(rx).await) (or maybe try_join)?

[robert-sjoblom]

Honestly, it's a case of code blindness. I'd figured out that a channel was one way to do it and then never questioned the approach. There's no issue calling poll_until_ready from the handler, as you correctly pointed out.

Thank you!