added addr allocation policy

2 files changed, 204 insertions, 16 deletions

diff --git a/src/address_allocation_policy.rs b/src/address_allocation_policy.rs
new file mode 100644
index 0000000..758ab21
--- /dev/null
+++ b/src/address_allocation_policy.rs
@@ -0,0 +1,156 @@
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum AddressAllocationPolicy {
+    PerCpu(u32),
+    PerMachine(u32),
+    Total(u32),
+}
+
+#[derive(Debug)]
+pub struct InvalidAddressAllocationPolicy(String);
+
+impl std::fmt::Display for InvalidAddressAllocationPolicy {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("invalid address allocation policy: ")?;
+        f.write_str(&self.0)
+    }
+}
+
+impl std::error::Error for InvalidAddressAllocationPolicy {}
+
+impl From<std::num::ParseIntError> for InvalidAddressAllocationPolicy {
+    fn from(value: std::num::ParseIntError) -> Self {
+        Self(value.to_string())
+    }
+}
+
+impl std::str::FromStr for AddressAllocationPolicy {
+    type Err = InvalidAddressAllocationPolicy;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if let Some(n) = s.strip_suffix("/cpu") {
+            Ok(Self::PerCpu(n.parse()?))
+        } else if let Some(n) = s.strip_suffix("/machine") {
+            Ok(Self::PerMachine(n.parse()?))
+        } else {
+            Ok(Self::Total(s.parse()?))
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::str::FromStr;
+
+    #[test]
+    fn test_per_cpu_parsing() {
+        assert_eq!(
+            AddressAllocationPolicy::from_str("10/cpu").unwrap(),
+            AddressAllocationPolicy::PerCpu(10)
+        );
+        assert_eq!(
+            AddressAllocationPolicy::from_str("1/cpu").unwrap(),
+            AddressAllocationPolicy::PerCpu(1)
+        );
+        assert_eq!(
+            AddressAllocationPolicy::from_str("1000/cpu").unwrap(),
+            AddressAllocationPolicy::PerCpu(1000)
+        );
+    }
+
+    #[test]
+    fn test_per_machine_parsing() {
+        assert_eq!(
+            AddressAllocationPolicy::from_str("20/machine").unwrap(),
+            AddressAllocationPolicy::PerMachine(20)
+        );
+        assert_eq!(
+            AddressAllocationPolicy::from_str("1/machine").unwrap(),
+            AddressAllocationPolicy::PerMachine(1)
+        );
+        assert_eq!(
+            AddressAllocationPolicy::from_str("500/machine").unwrap(),
+            AddressAllocationPolicy::PerMachine(500)
+        );
+    }
+
+    #[test]
+    fn test_total_parsing() {
+        assert_eq!(
+            AddressAllocationPolicy::from_str("100").unwrap(),
+            AddressAllocationPolicy::Total(100)
+        );
+        assert_eq!(
+            AddressAllocationPolicy::from_str("1").unwrap(),
+            AddressAllocationPolicy::Total(1)
+        );
+        assert_eq!(
+            AddressAllocationPolicy::from_str("9999").unwrap(),
+            AddressAllocationPolicy::Total(9999)
+        );
+    }
+
+    #[test]
+    fn test_invalid_number_formats() {
+        assert!(AddressAllocationPolicy::from_str("-5/cpu").is_err());
+        assert!(AddressAllocationPolicy::from_str("abc/cpu").is_err());
+        assert!(AddressAllocationPolicy::from_str("10.5/machine").is_err());
+        assert!(AddressAllocationPolicy::from_str("xyz").is_err());
+        assert!(AddressAllocationPolicy::from_str("").is_err());
+    }
+
+    #[test]
+    fn test_invalid_suffixes() {
+        assert!(AddressAllocationPolicy::from_str("10/node").is_err());
+        assert!(AddressAllocationPolicy::from_str("10/core").is_err());
+        assert!(AddressAllocationPolicy::from_str("10/").is_err());
+    }
+
+    #[test]
+    fn test_zero_values() {
+        assert_eq!(
+            AddressAllocationPolicy::from_str("0/cpu").unwrap(),
+            AddressAllocationPolicy::PerCpu(0)
+        );
+        assert_eq!(
+            AddressAllocationPolicy::from_str("0/machine").unwrap(),
+            AddressAllocationPolicy::PerMachine(0)
+        );
+        assert_eq!(
+            AddressAllocationPolicy::from_str("0").unwrap(),
+            AddressAllocationPolicy::Total(0)
+        );
+    }
+
+    #[test]
+    fn test_large_numbers() {
+        assert_eq!(
+            AddressAllocationPolicy::from_str("4294967295/cpu").unwrap(),
+            AddressAllocationPolicy::PerCpu(u32::MAX)
+        );
+        assert_eq!(
+            AddressAllocationPolicy::from_str("4294967295/machine").unwrap(),
+            AddressAllocationPolicy::PerMachine(u32::MAX)
+        );
+        assert_eq!(
+            AddressAllocationPolicy::from_str("4294967295").unwrap(),
+            AddressAllocationPolicy::Total(u32::MAX)
+        );
+    }
+
+    #[test]
+    fn test_overflow() {
+        assert!(AddressAllocationPolicy::from_str("4294967296/cpu").is_err());
+        assert!(AddressAllocationPolicy::from_str("9999999999999/machine").is_err());
+        assert!(AddressAllocationPolicy::from_str("18446744073709551616").is_err());
+    }
+
+    #[test]
+    fn test_whitespace_handling() {
+        assert!(AddressAllocationPolicy::from_str(" 10/cpu").is_err());
+        assert!(AddressAllocationPolicy::from_str("10/cpu ").is_err());
+        assert!(AddressAllocationPolicy::from_str("10 /cpu").is_err());
+        assert!(AddressAllocationPolicy::from_str("10/ cpu").is_err());
+    }
+}
+
diff --git a/src/main.rs b/src/main.rs
index 47b4493..f05dad9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -18,10 +18,12 @@ use tokio::{
 };
 
 use crate::{
+    address_allocation_policy::AddressAllocationPolicy,
     context::{Context, ExecutionNode},
     latency_matrix::LatencyMatrix,
 };
 
+pub mod address_allocation_policy;
 pub mod context;
 pub mod latency_matrix;
 pub mod machine;
@@ -69,7 +71,7 @@ struct NetUpArgs {
     #[clap(flatten)]
     common: Common,
     #[clap(long)]
-    addr_per_cpu: u32,
+    addresses: AddressAllocationPolicy,
     #[clap(long)]
     latency_matrix: PathBuf,
 }
@@ -92,7 +94,7 @@ struct NetPreviewArgs {
     machine: Vec<Machine>,
 
     #[clap(long)]
-    addr_per_cpu: u32,
+    addresses: AddressAllocationPolicy,
 
     #[clap(long)]
     latency_matrix: PathBuf,
@@ -153,7 +155,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.addr_per_cpu);
+    let configs = machine_generate_configs(&matrix, &machines, &args.addresses);
     machines_net_container_build(&context, &machines).await?;
     machines_clean(&context, &machines).await?;
     machines_configure(&context, &configs).await?;
@@ -197,7 +199,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.addr_per_cpu);
+    let configs = machine_generate_configs(&matrix, &machines, &args.addresses);
 
     for config in configs {
         (0..20).for_each(|_| print!("-"));
@@ -757,38 +759,68 @@ fn machine_address_for_idx(machine: Machine, idx: u32) -> Ipv4Addr {
 fn machine_generate_configs(
     matrix: &LatencyMatrix,
     machines: &[Machine],
-    addr_per_cpu: u32,
+    addr_policy: &AddressAllocationPolicy,
 ) -> Vec<MachineConfig> {
     let mut configs = Vec::default();
     let mut addresses = Vec::default();
     let mut address_to_index = HashMap::<Ipv4Addr, usize>::default();
+    let mut addresses_per_machine = HashMap::<Machine, Vec<Ipv4Addr>>::default();
+    machines.iter().for_each(|&m| {
+        addresses_per_machine.insert(m, Default::default());
+    });
 
     // gather all addresses across all machines
-    for &machine in machines {
-        for i in 0..(addr_per_cpu * machine.cpus()) {
-            let address = machine_address_for_idx(machine, i);
-            addresses.push(address);
-            address_to_index.insert(address, addresses.len() - 1);
+    match addr_policy {
+        AddressAllocationPolicy::PerCpu(n) => {
+            for &machine in machines {
+                for i in 0..(n * machine.cpus()) {
+                    let address = machine_address_for_idx(machine, i);
+                    addresses.push(address);
+                }
+            }
         }
+        AddressAllocationPolicy::PerMachine(n) => {
+            for &machine in machines {
+                for i in 0..*n {
+                    let address = machine_address_for_idx(machine, i);
+                    addresses.push(address);
+                }
+            }
+        }
+        AddressAllocationPolicy::Total(n) => {
+            let mut counter = 0;
+            while counter < *n {
+                let machine = machines[(counter as usize) % machines.len()];
+                let address = machine_address_for_idx(machine, counter / (machines.len() as u32));
+                addresses.push(address);
+                counter += 1;
+            }
+        }
+    }
+    for (idx, &address) in addresses.iter().enumerate() {
+        let machine = machine_from_addr(address).expect("we should only generate valid addresses");
+        address_to_index.insert(address, idx);
+        addresses_per_machine
+            .entry(machine)
+            .or_default()
+            .push(address);
     }
 
     for &machine in machines {
-        let mut machine_addresses = Vec::default();
+        let machine_addresses = &addresses_per_machine[&machine];
         let mut machine_ip_commands = Vec::default();
         let mut machine_tc_commands = Vec::default();
         let mut machine_nft_script = String::default();
 
         machine_ip_commands.push(format!("route add 10.0.0.0/8 dev {}", machine.interface()));
-        for i in 0..(addr_per_cpu * machine.cpus()) {
-            let address = machine_address_for_idx(machine, i);
-            machine_addresses.push(address);
+        for address in machine_addresses.iter() {
             machine_ip_commands.push(format!("addr add {address}/32 dev {}", machine.interface()));
         }
 
         let mut latencies_set = HashSet::<u32>::default();
         let mut latencies_buckets = Vec::<u32>::default();
         let mut latencies_addr_pairs = HashMap::<u32, Vec<(Ipv4Addr, Ipv4Addr)>>::default();
-        for &addr in &machine_addresses {
+        for &addr in machine_addresses {
             let addr_idx = address_to_index[&addr];
             for other_idx in (0..addresses.len()).filter(|i| *i != addr_idx) {
                 let other = addresses[other_idx];
@@ -863,7 +895,7 @@ fn machine_generate_configs(
 
         configs.push(MachineConfig {
             machine,
-            addresses: machine_addresses,
+            addresses: machine_addresses.clone(),
             nft_script: machine_nft_script,
             tc_commands: machine_tc_commands,
             ip_commands: machine_ip_commands,