How to wait all spawned async tasks

user_1395740 0
Viewed 2371

I have an async method that uses tokio::fs to explore a directory:

use failure::Error;
use futures::Future;
use std::path::PathBuf;
use tokio::prelude::*;

fn visit_async(path: PathBuf) -> Box<Future<Item = (), Error = Error> + Send> {
    let task = tokio::fs::read_dir(path)
        .flatten_stream()
        .for_each(move |entry| {
            let path = entry.path();
            if path.is_dir() {
                let task = visit_async(entry.path());
                tokio::spawn(task.map_err(drop));
            } else {
                println!("File: {:?}", path);
            }
            future::ok(())
        })
        .map_err(Error::from);
    Box::new(task)
}

I need to execute another future after all the the future returned by this method ends as well as all the tasks spawned by it. Is there a better way that just starting another runtime?

let t = visit_async(PathBuf::from(".")).map_err(drop);
tokio::run(t);

tokio::run(future::ok(()));
rust future runtime asynchronous
2 Answers

I'd strive to avoid using tokio::spawn() here, and try to wrap it all into a single future (in general, I think you only do tokio::spawn when you don't care about the result or execution, which we do here). That should make it easy to wait for completion. I haven't tested this, but something along these lines might do the trick:

    let task = tokio::fs::read_dir(path)
        .flatten_stream()
        .for_each(move |entry| {
            let path = entry.path();
            if path.is_dir() {
                let task = visit_async(entry.path());
                future::Either::A(task)
            } else {
                println!("File: {:?}", path);
                future::Either::B(future::ok(()))
            }
        })
        .map_err(Error::from)
        .and_then(|_| {
            // Do some work once all tasks complete
        });
    Box::new(task)

This will cause the asynchronous tasks to execute in sequence. You could use and_then instead of for_each to execute them in parallel and then into_future().and_then(|_| { ... }) to tuck on some action to execute afterwards.

Emil_H_737720

There's another issue with parallel descent in the FS: you may run out of file descriptors.

There is a way to solve both issues by creating a tokio::sync::Semaphore to limit the concurrent number of these tasks. After you are done spawning all of them, you can use Semaphore::acquire_many with the same value you used at creation, to block until all other tasks are finished.

For correctness, you should acquire the semaphore before spawning the task, and then pass the SemaphorePermit to the task (and make sure it doesn't get dropped before you are done). If you acquire the semaphore inside the tasks, there is a risk the main task might acquire all the permits before the first sub-task has a chance to run.

Since you can only move a SemaphorePermit<'static> inside the task, you will need to have a &'static Semaphore, for instance using lazy_static! or Box::leak.

user_2759840
Related