rustrust-tokio

How to remotely shut down running tasks with Tokio

I have a UDP socket that is receiving data

pub async fn start()  -> Result<(), std::io::Error> {
    loop {
        let mut data  = vec![0; 1024];
        socket.recv_from(&mut data).await?;
    }
}

This code is currently blocked on the .await when there is no data coming in. I want to gracefully shut down my server from my main thread, so how do I send a signal to this .await that it should stop sleeping and shut down instead?

Solution

  • Note: The Tokio website has a page on graceful shutdown. The modern recommendation is to use a cancellation token like described on that page. The answer below predates cancellation tokens.

    If you have more than one task to kill, you should use a broadcast channel to send shutdown messages. You can use it together with tokio::select!.

    use tokio::sync::broadcast::Receiver;

// You may want to log errors rather than return them in this function.
pub async fn start(kill: Receiver<()>) -> Result<(), std::io::Error> {
    tokio::select! {
        output = real_start() => output,
        _ = kill.recv() => Err(...),
    }
}

pub async fn real_start() -> Result<(), std::io::Error> {
    loop {
        let mut data  = vec![0; 1024];
        socket.recv_from(&mut data).await?;
    }
}

    Then to kill all the tasks, send a message on the channel.

    To kill only a single task, you can use the JoinHandle::abort method, which will kill the task as soon as possible. Note that this method is available only in Tokio 1.x and 0.3.x, and to abort a task using Tokio 0.2.x, see the next section below.

    let task = tokio::spawn(start());

...

task.abort();

    As an alternative to JoinHandle::abort, you can use abortable from the futures crate. When you spawn the task, you do the following:

    let (task, handle) = abortable(start());
tokio::spawn(task);

    Then later you can kill the task by calling the abort method.

    handle.abort();

    Of course, a channel with select! can also be used to kill a single task, perhaps combined with an oneshot channel rather than a broadcast channel.

    All of these methods guarantee that the real_start method is killed at an .await. It is not possible to kill the task while it is running code between two .awaits. You can read more about why this is here.

    The mini-redis project contains an accessible real-world example of graceful shutdown of a server. Additionally, the Tokio tutorial has chapters on both select and channels.