How can I stop an async program from terminating when any task is running?

user_124486 0
Viewed 273

Let's say I have something like this,

async fn do_update() {
    // here we store it.
    let task = task::spawn(async {
        let duration = Duration::from_millis(10);
        let mut stream = tokio::time::interval(duration);
        stream.tick().await;
        loop {
            println!("Foo");
            stream.tick().await;
        }
    });
    // and here we await it.
    task.await;
}

If I call that, like do_update().await the loop runs forever. Which is what I want. But I don't want it to be the only thing that runs. I just want it to block termination until the task is resolved.

I want an event to run every 5 seconds without blocking the rest of the program. This works exactly as I want, if I put it at the end of the program. It seems that this was once done with shutdown_on_idle, but the solution now offered isn't valid to me for the same reason as this comment,

I cannot implement it like that because I spawn inside of another task.

In a system where tasks spawn tasks, how can I be sure that the executable is only shut down when there are no more tasks running?

rust-tokio async-await rust
1 Answers

I think you are looking for structured concurrency, specifically a global implicit scope. Sadly, nobody could find a good solution, so work on that has been abandoned for the time being.

In the meantime, here's one possible workaround. Call spawn_keep_alive instead of tokio::spawn for any tasks that should keep the program alive:

Playground

use parking_lot::Mutex;
use std::{
    future::Future,
    sync::atomic::{AtomicU32, Ordering},
    time::Duration,
};
use tokio::{sync::oneshot, task::JoinHandle};

static KEEPALIVE_COUNT: AtomicU32 = AtomicU32::new(0);
static KEEPALIVE_SENDER: Mutex<Option<oneshot::Sender<()>>> = parking_lot::const_mutex(None);

pub fn spawn_keep_alive<T>(task: T) -> JoinHandle<T::Output>
where
    T: Future + Send + 'static,
    T::Output: Send + 'static,
{
    KEEPALIVE_COUNT.fetch_add(1, Ordering::Relaxed);
    tokio::spawn(async {
        let result = task.await;
        if KEEPALIVE_COUNT.fetch_sub(1, Ordering::Relaxed) == 1 {
            let sender = KEEPALIVE_SENDER.try_lock().unwrap().take().unwrap();
            sender.send(()).unwrap();
        }
        result
    })
}

async fn do_update() {
    let mut stream = tokio::time::interval(Duration::from_millis(100));
    stream.tick().await;
    for _ in 0..10 {
        println!("Foo");
        stream.tick().await;
    }
    spawn_keep_alive(async {
        tokio::time::sleep(Duration::from_millis(1000)).await;
        println!("I'm aliveeee!");
    });
    spawn_keep_alive(async {
        tokio::time::sleep(Duration::from_millis(2000)).await;
        println!("Don't forget about me!");
    });
}

#[tokio::main]
async fn main() {
    let (send, recv) = oneshot::channel();
    *KEEPALIVE_SENDER.try_lock().unwrap() = Some(send);
    spawn_keep_alive(do_update());
    // Wait for all keep-alive tasks to finish
    recv.await.unwrap();
}

(KEEPALIVE_SENDER could probably be more efficient, but it's only used twice).

user_33988390
Related