CODEWITHSUNDEEP
rustasync-await
Tim
Tim

Reputation: 4823

Polling many futures of different types

I'm trying to understand how to implement polling multiple futures with different types. For context, I'm calling an API that will return something like:

[{"type": "source_a", "id": 123}, {"type": "source_b", "id": 234}, ...]

Each type in the API response requires a call to another API, with each API returning different data types. The code I've written works something like this:

async fn get_data(sources: Vec<Source>) -> Data {
    let mut data = Default::default();

    for source in sources {
        if source.kind == "source_a" {
            let source_data = get_source_a(source).await;

            process_source_a(source_data, &mut data);
        } else if source.kind == "source_b" {
            ...
        }
    }

    data
}

This won't run concurrently, it will simply fetch sources one at a time and process them. How can I rewrite this so that each source is fetched concurrently and then processed once data is available? Speaking Rustily, I think what I want is to execute a closure that mutably borrows data when the future is ready. Should I be looking at something like an Arc<RefCell<Data>>?

Upvotes: 2

Views: 1629

Answers (1)

user4815162342
user4815162342

Reputation: 155446

To process the futures in parallel, you need to await something like join_all, which will run them concurrently and return when they are all done. For this to work, you have to resolve two issues:

  • join_all requires futures of the same type, so you need to box them or otherwise unify them.
  • data needs to be accessed by multiple async blocks, so it needs to be protected by Arc and Mutex.

The first issue can be solved simply by spawning the async fns as tasks, which has the added advantage of potentially running them in parallel (in addition to them being run concurrently). The example below uses tokio::spawn, but it would be almost exactly the same for async_std. Since there is no reproducible example, I can't test the code, but it could look like this:

async fn get_data(sources: Vec<Source>) -> Data {
    let data = Arc::new(Mutex::new(Data::default()));
    let mut tasks = vec![];

    for source in sources {
        if source.kind == "source_a" {
            let data = Arc::clone(&data);
            tasks.push(tokio::task::spawn(async move {
                let source_data = get_source_a(source).await;
                process_source_a(source_data, &mut data.lock().unwrap());
            }));
        } else if source.kind == "source_b" {
            // ...
        }
    }

    // Wait for all sources to finish and propagate the panic if any.
    // With async_std this wouldn't require the `for_each()`.
    futures::future::join_all(tasks)
        .await
        .for_each(|x| x.unwrap());

    // As all tasks are done, there should be no references to `data` at
    // this point, so we can extract it out of the Arc<Mutex<_>> wrapping.
    data.try_unwrap().unwrap().into_inner()
}

Upvotes: 4

Related Questions