Process elements in futures::stream::Stream concurrently with reference to shared tasks - rust

Question

I am trying to connect to multiple peers and process a series of asynchronous tasks for each peer. Once a task is completed by any peer it should short-circuit the task for all other peers even if the task is in progress.
In the following example I have 3 tasks:
use std::cell::RefCell;
use std::sync::Arc;
use std::rc::Rc;
use futures::stream::StreamExt;
use rand::Rng;
#[derive(Clone)]
struct Task {
id: u32,
}
impl Task {
fn process<'a>(&'a self, peer: &'a str) -> impl futures::future::Future<Output = ()> + 'a {
let delay = core::time::Duration::from_secs(rand::thread_rng().gen_range(5..10));
async move {
tokio::time::sleep(delay).await;
println!("{} completed task #{} first", peer, self.id);
}
}
}
async fn process_task(task: &Task, done_with_task: &Rc<RefCell<bool>>, notify_task: &Arc<tokio::sync::Notify>, peer: &str) {
if !*done_with_task.borrow() {
tokio::select! {
process = task.process(peer) => {
done_with_task.replace(true);
notify_task.notify_waiters();
}
cancel = notify_task.notified() => {}
}
}
}
#[tokio::main]
async fn main() {
let peers = vec!["peer A", "peer B", "peer C"];
let peers = futures::stream::iter(peers);
let (tx, rx) = tokio::sync::mpsc::channel(100);
let rx = tokio_stream::wrappers::ReceiverStream::new(rx);
let rx = peers.chain(rx);
let tasks = vec![
Task {id: 1},
Task {id: 2},
Task {id: 3},
];
let notify_task_1 = Arc::new(tokio::sync::Notify::new());
let done_with_task_1 = Rc::new(RefCell::new(false));
let notify_task_2 = Arc::new(tokio::sync::Notify::new());
let done_with_task_2 = Rc::new(RefCell::new(false));
let notify_task_3 = Arc::new(tokio::sync::Notify::new());
let done_with_task_3 = Rc::new(RefCell::new(false));
let handle_conn_fut = rx.for_each_concurrent(0, |peer| {
let mut tasks = tasks.clone().into_iter();
let notify_task_1 = notify_task_1.clone();
let notify_task_2 = notify_task_2.clone();
let notify_task_3 = notify_task_3.clone();
let done_with_task_1 = done_with_task_1.clone();
let done_with_task_2 = done_with_task_2.clone();
let done_with_task_3 = done_with_task_3.clone();
async move {
println!("connecting to {}", peer);
// task #1
let current_task = tasks.next().unwrap();
process_task(&current_task, &done_with_task_1, &notify_task_1, peer).await;
println!("task #{} done for {}", current_task.id, peer);
// task #2
let current_task = tasks.next().unwrap();
process_task(&current_task, &done_with_task_2, &notify_task_2, peer).await;
println!("task #{} done for {}", current_task.id, peer);
// task #3
let current_task = tasks.next().unwrap();
process_task(&current_task, &done_with_task_3, &notify_task_3, peer).await;
println!("task #{} done for {}", current_task.id, peer);
}
});
let create_new_conn_fut = async move {
// a new peer connects after 10 seconds
tokio::time::sleep(core::time::Duration::from_secs(10)).await;
tx.send("peer D").await.unwrap()
};
// awaits both futures in parallell
futures::future::join(handle_conn_fut, create_new_conn_fut).await;
}
output:
connecting to peer A
connecting to peer B
connecting to peer C
peer A completed task #1 first
task #1 done for peer A
task #1 done for peer C
task #1 done for peer B
connecting to peer D
task #1 done for peer D
peer B completed task #2 first
task #2 done for peer B
task #2 done for peer D
task #2 done for peer C
task #2 done for peer A
peer B completed task #3 first
task #3 done for peer B
task #3 done for peer A
task #3 done for peer C
task #3 done for peer D
The output is correct, but I am trying to find a more idiomatic way to write this, without explicitly needing to handle each task.
I have thought of storing the tasks in a futures::stream::FuturesOrdered, but I am not sure how to go about this.
With abuse of notation, I would like something like this to work in theory:
let tasks = Rc::new(RefCell::new(tasks)) // ?
let handle_conn_fut = rx.for_each_concurrent(0, |peer| {
async move {
println!("connecting to {}", peer);
while let Some(task) = tasks.borrow().next() {
tokio::select! {
complete = tasks.process(peer) => {
// mutate tasks
tasks.borrow_mut().pop() // ?
}
state_change = tasks.is_mutated() => {} // ?
}
}
}
});

Answer 1

Solved this by storing the current task in a tokio::sync::watch::channel which makes it possible to await changes to the task with has_changed. The backlog or future tasks needs to stored in mutable memory location that can be moved into different async blocks, so I used Rc<RefCell<Vec<Task>>> for this. One could maybe use the Rc<RefCell<Vec<Task>>> directly together with tokio::sync::Notify instead of a channel, but I had some trouble with borrowing and mutating the Rc<RefCell<Vec<Task>>> simultaneously in process_taskas explained here: https://doc.rust-lang.org/stable/std/cell/struct.RefCell.html#panics-4
use std::cell::RefCell;
use std::rc::Rc;
use futures::stream::StreamExt;
use rand::Rng;
#[derive(Clone, Debug)]
struct Task {
id: u32,
}
impl Task {
fn process<'a>(&'a self, peer: &'a str) -> impl futures::future::Future<Output = ()> + 'a {
let delay = core::time::Duration::from_secs(rand::thread_rng().gen_range(5..10));
async move {
tokio::time::sleep(delay).await;
println!("{} completed task #{} first", peer, self.id);
}
}
}
async fn process_task(task: &Task, peer: &str, task_backlog: &Rc<RefCell<Vec<Task>>>, tx_task: &tokio::sync::watch::Sender<Task>, rx_task: &mut tokio::sync::watch::Receiver<Task>) {
tokio::select! {
process = task.process(peer) => {
let new_task = task_backlog.borrow_mut().pop().unwrap();
tx_task.send(new_task).unwrap()
}
cancel = rx_task.changed() => {}
}
}
#[tokio::main]
async fn main() {
let peers = vec!["peer A", "peer B", "peer C"];
let peers = futures::stream::iter(peers);
let (tx, rx) = tokio::sync::mpsc::channel(100);
let rx = tokio_stream::wrappers::ReceiverStream::new(rx);
let rx = peers.chain(rx);
let task_backlog = Rc::new(RefCell::new(vec![
Task {id: 3},
Task {id: 2},
]));
let (tx_task, rx_task) = tokio::sync::watch::channel(Task { id: 1 });
let tx_task_ref = &tx_task;
let handle_conn_fut = rx.for_each_concurrent(0, |peer| {
let task_backlog = task_backlog.clone();
let mut rx_task = rx_task.clone();
let tx_task = tx_task_ref.clone();
async move {
println!("connecting to {}", peer);
loop {
let current_task = rx_task.borrow_and_update().clone();
process_task(&current_task, peer, &task_backlog, &tx_task, &mut rx_task).await;
println!("task #{} done for {}", current_task.id, peer);
}
}
});
let create_new_conn_fut = async move {
// a new peer connects after 10 seconds
tokio::time::sleep(core::time::Duration::from_secs(10)).await;
tx.send("peer D").await.unwrap()
};
// awaits both futures in parallell
futures::future::join(handle_conn_fut, create_new_conn_fut).await;
}