reusing futures::stream::Stream and modifying state of each element - asynchronous

Question

I have a futures::stream::Stream which produces elements in the form <(State, impl std::fmt::Binary)> (Binary is an arbitrary placeholder for a trait I want to use):
let peers = (0..10).map(move |peer| async move {
let delay = core::time::Duration::from_secs(2); // should be 'random'
tokio::time::sleep(delay).await;
if peer % 2 == 0 {
stream::iter(Ok::<i32, std::io::Error>(peer).into_iter())
} else {
let custom_error = std::io::Error::new(std::io::ErrorKind::Other, "oh no!");
stream::iter(Err::<i32, std::io::Error>(custom_error).into_iter())
}
})
.collect::<FuturesUnordered<_>>()
.flatten()
.map(|peer| (State { foo: "foo".into(), bar: "bar".into() }, peer));
The peers above should correspond to a stream of peers which are connected successfully. Since I do not know until runtime how many peers are connected I can't store this in a Vec<(State, impl ...)> or similar.
Is it possible to somehow do a series of tasks concurrently which modifies the State internally for each peer where the task completion is determined by the first peer that completes the task? Similar to a race for each task.
I thought the following might work:
use futures::{
stream::{self, FuturesUnordered},
StreamExt, FutureExt,
};
#[derive(Debug, Clone)]
struct State {
foo: String,
bar: String
}
#[tokio::main]
async fn main() {
let futures = (0..10).map(move |peer| async move {
let mut delay = core::time::Duration::from_secs(2);
if peer == 0 {
delay = core::time::Duration::from_secs(100); // slow peer
}
tokio::time::sleep(delay).await;
if peer % 2 == 0 {
stream::iter(Ok::<i32, std::io::Error>(peer).into_iter())
} else {
let custom_error = std::io::Error::new(std::io::ErrorKind::Other, "oh no!");
stream::iter(Err::<i32, std::io::Error>(custom_error).into_iter())
}
})
.collect::<FuturesUnordered<_>>()
.flatten()
.map(|peer| (State { foo: "foo".into(), bar: "bar".into() }, peer));
// first task
let notify = std::rc::Rc::new(tokio::sync::Notify::new());
let futures = futures.map(|(mut state, x)| {
let notify = notify.clone();
async move {
tokio::select! {
biased;
_ = async {
println!("processing task #1 for peer {:b} with state {:?}", x, state);
let delay = core::time::Duration::from_secs(2);
tokio::time::sleep(delay).await;
state.foo = "test1".to_owned();
notify.notify_waiters();
} => {
(state, x)
}
_ = notify.notified() => { (state, x) }
}
}
}).buffer_unordered(10).collect::<Vec<(State, _)>>().await;
// second task
let notify = std::rc::Rc::new(tokio::sync::Notify::new());
let futures = stream::iter(futures);
let futures = futures.map(|(mut state, x)| {
let notify = notify.clone();
async move {
tokio::select! {
biased;
_ = async {
println!("processing task #2 for peer {:b} with state {:?}", x, state);
let delay = core::time::Duration::from_secs(2);
tokio::time::sleep(delay).await;
notify.notify_waiters();
} => {
(state, x)
}
_ = notify.notified() => { (state, x) }
}
}
}).buffer_unordered(10).collect::<Vec<(State, _)>>().await;
}
playground link
But it will be stuck on the first task because it is waiting for the slow peer with 100 seconds delay. Ideally, I want to prematurely finish the collect once the task is done. I have tried using take_until with notify.notified():
let futures = futures.map(|(mut state, x)| {
let notify = notify.clone();
async move {
tokio::select! {
...
}
}
}).buffer_unordered(10).take_until(notify.notified()).collect::<Vec<(State, _)>>().await;
but this will discard the other peers and leave only 1 peer in futures. I think this is because the outer notify.notified() takes precedence over the inner notify.notified() used in the tokio::select! statement.
Is there a way to reuse a futures::stream::Stream and simultaneously modify the elements which I have tried doing above?
Or is there a more idiomatic solution to what I am trying to achieve here?