added concurrency limit via OAR_P2P_CONCURRENCY_LIMIT

the env var OAR_P2P_CONCURRENCY_LIMIT limits the number of parallel
"operations" being done on the cluster machines. so, if it is set to 3,
then we only work on 3 machines at time. setting to 0 means unlimited.

2 files changed, 69 insertions, 35 deletions

diff --git a/src/machine.rs b/src/machine.rs
index 958a03c..ce5068f 100644
--- a/src/machine.rs
+++ b/src/machine.rs
@@ -1,5 +1,8 @@
+use std::sync::Arc;
+
 use eyre::{Context as _, Result};
 use futures::{StreamExt as _, stream::FuturesUnordered};
+use tokio::sync::Semaphore;
 
 macro_rules! define_machines {
     ($(($name:ident, $idx:expr, $hostname:expr, $cpus:expr, $interface:expr)),*) => {
@@ -141,23 +144,64 @@ define_machines!(
     (Vulpix1, 56, "vulpix-1", 112, todo!())
 );
 
-pub async fn for_each<F, FUT, RET>(machines: impl IntoIterator<Item = &Machine>, f: F) -> Result<()>
+pub async fn for_each<F, FUT, RET>(
+    machines: impl IntoIterator<Item = &Machine>,
+    f: F,
+) -> Result<Vec<(Machine, RET)>>
+where
+    F: Fn(Machine) -> FUT,
+    RET: Send + 'static,
+    FUT: std::future::Future<Output = Result<RET>>,
+{
+    let limit = match std::env::var("OAR_P2P_CONCURRENCY_LIMIT") {
+        Ok(value) => {
+            tracing::trace!("parsing concurrency limit value '{value}'");
+            let limit = value
+                .parse()
+                .expect("invalid value for OAR_P2P_CONCURRENCY_LIMIT");
+            tracing::debug!("using concurrency limit = {limit}");
+            limit
+        }
+        Err(_) => 0,
+    };
+    for_each_with_limit(machines, limit, f).await
+}
+
+pub async fn for_each_with_limit<F, FUT, RET>(
+    machines: impl IntoIterator<Item = &Machine>,
+    limit: usize,
+    f: F,
+) -> Result<Vec<(Machine, RET)>>
 where
     F: Fn(Machine) -> FUT,
+    RET: Send + 'static,
     FUT: std::future::Future<Output = Result<RET>>,
 {
+    let sem = Arc::new(Semaphore::new(if limit == 0 {
+        Semaphore::MAX_PERMITS
+    } else {
+        limit
+    }));
     let mut futures = FuturesUnordered::new();
 
     for &machine in machines {
         let fut = f(machine);
-        let fut = async move { (machine, fut.await) };
+        let sem = sem.clone();
+        let fut = async move {
+            let _permit = sem.acquire().await.unwrap();
+            (machine, fut.await)
+        };
         futures.push(fut);
     }
 
+    let mut results = Vec::default();
     while let Some((machine, result)) = futures.next().await {
-        if let Err(err) = result {
-            return Err(err).with_context(|| format!("running task on machine {machine}"));
+        match result {
+            Ok(value) => results.push((machine, value)),
+            Err(err) => {
+                return Err(err).with_context(|| format!("running task on machine {machine}"));
+            }
         }
     }
-    Ok(())
+    Ok(results)
 }
diff --git a/src/main.rs b/src/main.rs
index cc8afd0..68cbbdd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,7 +15,6 @@ use serde::Deserialize;
 use tokio::{
     io::{AsyncReadExt as _, AsyncWriteExt as _},
     process::Command,
-    task::JoinSet,
 };
 
 use crate::{
@@ -253,14 +252,14 @@ async fn cmd_net_down(args: NetDownArgs) -> Result<()> {
 async fn cmd_net_show(args: NetShowArgs) -> Result<()> {
     let context = context_from_common(&args.common).await?;
     let machines = oar::job_list_machines(&context).await?;
-    let mut set = JoinSet::default();
-    for machine in machines {
+    let results = machine::for_each(machines.iter(), |machine| {
         let context = context.clone();
-        set.spawn(async move { (machine, machine_list_addresses(&context, machine).await) });
-    }
+        async move { machine_list_addresses(&context, machine).await }
+    })
+    .await?;
+
     let mut addresses = Vec::default();
-    for (machine, result) in set.join_all().await {
-        let addrs = result?;
+    for (machine, addrs) in results {
         for addr in addrs {
             addresses.push((machine, addr));
         }
@@ -625,44 +624,35 @@ async fn machine_containers_clean(ctx: &Context, machine: Machine) -> Result<()>
 #[tracing::instrument(ret, err, skip_all)]
 async fn machines_clean(ctx: &Context, machines: &[Machine]) -> Result<()> {
     tracing::info!("cleaning machines: {machines:?}");
-    let mut set = JoinSet::default();
-    for &machine in machines {
+    machine::for_each(machines, |machine| {
         let ctx = ctx.clone();
-        set.spawn(async move { machine_clean(&ctx, machine).await });
-    }
-    let results = set.join_all().await;
-    for result in results {
-        result?;
-    }
+        async move { machine_clean(&ctx, machine).await }
+    })
+    .await?;
     Ok(())
 }
 
 #[tracing::instrument(ret, err, skip_all)]
 async fn machines_net_container_build(ctx: &Context, machines: &[Machine]) -> Result<()> {
     tracing::info!("building networking container for machines: {machines:?}");
-    let mut set = JoinSet::default();
-    for &machine in machines {
+    machine::for_each(machines, |machine| {
         let ctx = ctx.clone();
-        set.spawn(async move { machine_net_container_build(&ctx, machine).await });
-    }
-    for result in set.join_all().await {
-        result?;
-    }
+        async move { machine_net_container_build(&ctx, machine).await }
+    })
+    .await?;
     Ok(())
 }
 
 #[tracing::instrument(ret, err, skip_all)]
 async fn machines_configure(ctx: &Context, configs: &[MachineConfig]) -> Result<()> {
     tracing::info!("configuring machines");
-    let mut set = JoinSet::default();
-    for config in configs {
+    let machines = configs.iter().map(|c| &c.machine);
+    machine::for_each(machines, |machine| {
         let ctx = ctx.clone();
-        let config = config.clone();
-        set.spawn(async move { machine_configure(&ctx, &config).await });
-    }
-    for result in set.join_all().await {
-        result?;
-    }
+        let config = configs.iter().find(|c| c.machine == machine).unwrap();
+        async move { machine_configure(&ctx, &config).await }
+    })
+    .await?;
     Ok(())
 }