rustpython-asynciorust-tokio

How does Rust handle resource cleanup when tokio::select! cancels a future?


In Rust, when using tokio::select!, once one of the branches completes, the other branches are implicitly canceled and no longer polled. I am trying to understand how resource cleanup is handled in such cases.

Consider the following example from tokio official documentation:

use tokio::net::TcpStream;
use tokio::sync::oneshot;

#[tokio::main]
async fn main() {
    let (tx, rx) = oneshot::channel();

    // Spawn a task that sends a message over the oneshot
    tokio::spawn(async move {
        tx.send("done").unwrap();
    });

    tokio::select! {
        socket = TcpStream::connect("localhost:3465") => {
            println!("Socket connected {:?}", socket);
        }
        msg = rx => {
            println!("Received message first {:?}", msg);
        }
    }
}

Here, if rx completes first, the TcpStream::connect future gets canceled. My questions are:

  1. How is resource cleanup handled for canceled futures?
    • For example, if TcpStream::connect completes but is then canceled due to tokio::select!, what happens to the established connection?
    • Will it remain open indefinitely if not explicitly handled?
  2. Does Rust provide a mechanism similar to Python’s CancelledError for handling cleanup in an async task?
    • In Python’s asyncio, when a task is canceled, you can catch asyncio.CancelledError and perform cleanup. Is there an equivalent in Rust?
  3. How does Drop handle async cleanup?
    • Since resource cleanup (like closing a TCP connection) is usually async, and Drop does not allow async operations, how should one properly clean up async resources in such cases?

Would appreciate an explanation of best practices for handling this kind of situation in Rust’s async ecosystem.


Solution

  • How is resource cleanup handled for canceled futures?

    It works the same way with tokio::select! as any other mechanism of storing and polling a futures. When a future is dropped, its drop glue runs, including any manual Drop implementations.

    For example, if TcpStream::connect completes but is then canceled due to tokio::select!, what happens to the established connection?

    The connection will be closed.

    Does Rust provide a mechanism similar to Python’s CancelledError for handling cleanup in an async task?

    Yes, dropping is the standard way to handle this. Futures that need to know when they are canceled can implement Drop. This is usually only required by types that directly handle native resources -- note that the future type itself may not need to implement Drop if it contains a value of a type that handles this disposal itself. If your future is built on other futures, you usually don't have to do anything yourself -- the futures owned by your future are transitively dropped, and so the whole "future tree" is cleaned up automatically.

    While Rust async can be somewhat complex in other ways, handling cancellation is one of the easier things since it's almost always completely automatic unless you're managing native resources yourself, without any wrapper type that handles disposal for you.

    Since resource cleanup (like closing a TCP connection) is usually async, and Drop does not allow async operations, how should one properly clean up async resources in such cases?

    Just because futures provide async functionality does not mean that drop glue needs to be async. For example, a TCP connection could be synchronously closed in the Drop implementation. This is what tokio's TcpStream will do, because under the hood it wraps a std::net::TcpStream and does not customize/override that type's drop behavior.

    Alternatively, the Drop implementation could arrange for the cleanup to happen asynchronously, either by using runtime-internal machinery, or by e.g. tokio::spawning the cleanup code, or otherwise queuing the cleanup to happen later. For example, database connections belonging to a pool often need to do some cleanup (such as rolling back any outstanding transaction) before returning the connection to the pool.

    There has been discussion around adding an AsyncDrop trait that works like Drop but allows the implementation to return a future, but this does not yet exist.