added the --matrix-wrap flag

I got to the point of needed more than 10k rows and since my biggest
latency matrix is 10k in size this option allows the values to wrap so
we can create a network bigger than the matrix.

1 files changed, 24 insertions, 9 deletions

diff --git a/src/main.rs b/src/main.rs
index 303d4f3..9f50595 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -114,6 +114,9 @@ struct NetUpArgs {
     /// 99.91315 110.71518 93.90618 84.67561 0.0 {n}
     #[clap(long)]
     latency_matrix: PathBuf,
+
+    #[clap(long)]
+    matrix_wrap: bool,
 }
 
 #[derive(Debug, Args)]
@@ -141,6 +144,9 @@ struct NetPreviewArgs {
 
     #[clap(long)]
     latency_matrix: PathBuf,
+
+    #[clap(long)]
+    matrix_wrap: bool,
 }
 
 #[derive(Debug, Args)]
@@ -245,7 +251,7 @@ async fn cmd_net_up(args: NetUpArgs) -> Result<()> {
     let matrix = LatencyMatrix::parse(&matrix_content, latency_matrix::TimeUnit::Milliseconds)
         .context("parsing latency matrix")?;
     let machines = oar::job_list_machines(&context).await?;
-    let configs = machine_generate_configs(&matrix, &machines, &args.addresses)?;
+    let configs = machine_generate_configs(&matrix, args.matrix_wrap, &machines, &args.addresses)?;
     machines_net_container_build(&context, &machines).await?;
     machines_clean(&context, &machines).await?;
     machines_configure(&context, &configs).await?;
@@ -307,7 +313,7 @@ async fn cmd_net_preview(args: NetPreviewArgs) -> Result<()> {
     let matrix = LatencyMatrix::parse(&matrix_content, latency_matrix::TimeUnit::Milliseconds)
         .context("parsing latency matrix")?;
     let machines = args.machine;
-    let configs = machine_generate_configs(&matrix, &machines, &args.addresses)?;
+    let configs = machine_generate_configs(&matrix, args.matrix_wrap, &machines, &args.addresses)?;
 
     for config in configs {
         (0..20).for_each(|_| print!("-"));
@@ -955,6 +961,7 @@ fn machine_address_for_idx(machine: Machine, idx: u32) -> Ipv4Addr {
 
 fn machine_generate_configs(
     matrix: &LatencyMatrix,
+    matrix_wrap: bool,
     machines: &[Machine],
     addr_policy: &AddressAllocationPolicy,
 ) -> Result<Vec<MachineConfig>> {
@@ -1003,12 +1010,14 @@ fn machine_generate_configs(
             .push(address);
     }
 
-    if addresses.len() > matrix.dimension() {
-        return Err(eyre::eyre!(
-            "latency matrix is too small, size is {} but {} was required",
-            matrix.dimension(),
-            addresses.len()
-        ));
+    if !matrix_wrap {
+        if addresses.len() > matrix.dimension() {
+            return Err(eyre::eyre!(
+                "latency matrix is too small, size is {} but {} was required",
+                matrix.dimension(),
+                addresses.len()
+            ));
+        }
     }
 
     for &machine in machines {
@@ -1029,7 +1038,13 @@ fn machine_generate_configs(
             let addr_idx = address_to_index[&addr];
             for other_idx in (0..addresses.len()).filter(|i| *i != addr_idx) {
                 let other = addresses[other_idx];
-                let latency = matrix.latency(addr_idx, other_idx);
+                let latency = match matrix_wrap {
+                    true => matrix.latency(
+                        addr_idx % matrix.dimension(),
+                        other_idx % matrix.dimension(),
+                    ),
+                    false => matrix.latency(addr_idx, other_idx),
+                };
                 let latency_millis = u32::try_from(latency.as_millis()).unwrap();
                 if !latencies_set.contains(&latency_millis) {
                     latencies_set.insert(latency_millis);