Async-std: Using Barrier with Channels, the proper way?

[ansrivas]

Hi @stjepang thank you for the wonderful channels implementation.
cc: @yoshuawuyts

I am trying to migrate one of my go-implementations to async-std channels, so this is what I came up with, but wondering if this translation from go's waitgroup and channels to rust's Barrier and channels is correct.

I have tested the following implementation and it works, but just want to make sure that my understanding is correct.

$ cat Cargo.toml

[dependencies]
tokio = { version = "=0.2.0-alpha.6", default-features = false, features = ["rt-full", "tcp"] }
tokio-executor = "=0.2.0-alpha.6"

[dependencies.async-std] 
version = "0.99.11"
features = ["unstable"]

$ cat main.rs

use async_std::sync::{Arc, Barrier};
use async_std::task;
use std::time::Duration;

use async_std::sync::channel;
use async_std::sync::{Receiver, Sender};

async fn worker(input_chan: Receiver<i32>, results_chan: Sender<i32>) {
    while let Some(v) = input_chan.recv().await {

        // Simulate some work here
                task::sleep(Duration::from_secs(1)).await;

        println!("Publishing results of heavy computation {}", v);

                // Send the results back
        results_chan.send(v).await;
    }
}

#[tokio::main]
async fn main() {
  let max_worker = 5;
  let max_tasks = 10;

  // This channel's capacity should be equal to the 
  // amount of results we expect the workers to send back.
  // In case the capacity is lesser,  the worker task will never be able
  // to send the result and be stuck foreever
  let (s_results, r_results) = channel(max_tasks as usize);

  // This channel will only receive input tasks one by one, 
  // which will be picked up by the workers
  let (s, r) = channel(1 as usize);

  let barrier = Arc::new(Barrier::new(max_worker));
  let mut handles = Vec::with_capacity(max_worker);

  // Spawn 5 workers
  for _ in 0..max_worker {
    let c = barrier.clone();
    let receiver_chan_input = r.clone();
    let sender_chan_results = s_results.clone();

    handles.push(task::spawn(async move {
            worker(receiver_chan_input, sender_chan_results).await;
            c.wait().await;
    }));
   }

  // Send arbitrary amount of jobs, here we are sending numbers from 1 to 10
  // This will also block until a worker has picked up a job
  // because the size of channel is 1
  for idx in 1..max_tasks {
    s.send(idx).await
  }

  // Drop the sender part of the channel after publishing all the inputs for the job.
  drop(s);

  // Now wait and make sure that futures are all computed
  // This also means that the workers have successfully published
  // results of their computation to the results-channel.
  for handle in handles {
    handle.await;
  }
  // Now drop the s_results channel so that we mark this sender as closed.
  drop(s_results);

  let mut output = 0; 
  while let Some(x) = r_results.recv().await {
    output += x;
  }

  println!(
    "Final output of summing numbers between 0 and {} is {}",
    max_tasks, output
  );
}

[ghost]

Hey, this looks like solid code! :)

As another suggestion (not saying you should do it, it's just an idea), Go's waitgroup can be simulated with our channels:

• wg.Add(1) is like sender.clone()
• wg.Done() is like drop(sender)
• wg.Wait() is like receiver.recv().await

[ansrivas]

@stjepang Thank you, also just putting this here as an example.

use async_std::task;
use std::time::Duration;

use async_std::sync::channel;
use async_std::sync::{Receiver, Sender};

async fn worker(wg: Sender<i32>, id: i32) {
    println!("Starting work {}", id);
    task::sleep(Duration::from_secs(1)).await;
    println!("Done work {}", id);
    // Similar to wg.Done();
    drop(wg);
}

#[tokio::main]
async fn main() {
    let max_tasks = 10;
    let (sender, receiver) = channel(1 as usize);
    for idx in 0..max_tasks {
        // Similar to wg.Add(1)
        let wg = sender.clone();
        // spawn tasks ( similar goroutines )
        task::spawn(async move {
            worker(wg, idx).await;
        });
    }
    // Drop the original sender so that receiver
    // knows the channel is closed
    drop(sender);
    // Similar to wg.Done()
    receiver.recv().await;
}