Fixes to net-shaper, and net.sh option to start/stop shaper (#6981)

* Fixes to net-shaper, and net.sh option to start/stop shaper

* fix shellcheck

* more shellchecks

net-shaper/src/main.rs

@@ -50,6 +50,7 @@ fn run(
     status_err_msg: &str,
     ignore_err: bool,
 ) -> bool {
+    println!("Running {:?}", std::process::Command::new(cmd).args(args));
     let output = std::process::Command::new(cmd)
         .args(args)
         .output()
@@ -73,8 +74,8 @@ fn run(
     }
 }
 
-fn insert_iptables_rule(my_partition: usize) -> bool {
+fn insert_iptables_rule(tos: u8) -> bool {
-    let my_tos = my_partition.to_string();
+    let my_tos = tos.to_string();
 
     // iptables -t mangle -A PREROUTING -p udp -j TOS --set-tos <my_parition_index>
     run(
@@ -83,7 +84,7 @@ fn insert_iptables_rule(my_partition: usize) -> bool {
             "-t",
             "mangle",
             "-A",
-            "PREROUTING",
+            "OUTPUT",
             "-p",
             "udp",
             "-j",
@@ -198,7 +199,15 @@ fn identify_my_partition(partitions: &[u8], index: u64, size: u64) -> usize {
         }
     }
 
-    my_partition + 1
+    my_partition
+}
+
+fn partition_id_to_tos(partition: usize) -> u8 {
+    if partition < 4 {
+        2u8.pow(partition as u32 + 1)
+    } else {
+        0
+    }
 }
 
 fn shape_network(matches: &ArgMatches) {
@@ -217,12 +226,12 @@ fn shape_network(matches: &ArgMatches) {
 
     assert!(my_index < network_size);
 
-    let my_partition = identify_my_partition(&topology.partitions, my_index, network_size);
+    let my_partition = identify_my_partition(&topology.partitions, my_index + 1, network_size);
     println!("My partition is {}", my_partition);
 
     flush_iptables_rule();
 
-    if !insert_iptables_rule(my_partition) {
+    if !insert_iptables_rule(partition_id_to_tos(my_partition)) {
         return;
     }
 
@@ -234,26 +243,28 @@ fn shape_network(matches: &ArgMatches) {
 
     topology.interconnects.iter().for_each(|i| {
         if i.b as usize == my_partition {
-            let handle = (i.a + 1).to_string();
+            let tos = partition_id_to_tos(i.a as usize);
-            let tos_string = i.a.to_string();
+            if tos == 0 {
-            let class = format!("1:{}", i.a);
+                println!("Incorrect value of TOS/Partition in config {}", i.a);
+                return;
+            }
+            let tos_string = tos.to_string();
+            let class = format!("1:{}", i.a + 1);
             if !insert_tc_netem(
                 interface.as_str(),
                 class.as_str(),
-                handle.as_str(),
+                tos_string.as_str(),
                 i.config.as_str(),
             ) {
-                flush_iptables_rule();
                 delete_tc_root(interface.as_str());
                 return;
             }
 
             if !insert_tos_filter(interface.as_str(), class.as_str(), tos_string.as_str()) {
-                flush_iptables_rule();
                 delete_tc_netem(
                     interface.as_str(),
                     class.as_str(),
-                    handle.as_str(),
+                    tos_string.as_str(),
                     i.config.as_str(),
                 );
                 delete_tc_root(interface.as_str());

net/net.sh

@@ -102,6 +102,8 @@ Operate a configured testnet
  netem-specific options:
    --config            - Netem configuration (as a double quoted string)
    --parition          - Percentage of network that should be configured with netem
+   --config-file       - Configuration file for partition and netem configuration
+   --netem-cmd         - Optional command argument to netem. Default is "add". Use "cleanup" to remove rules.
 
  update-specific options:
    --platform linux|osx|windows       - Deploy the tarball using 'solana-install deploy ...' for the
@@ -147,6 +149,8 @@ gpuMode=auto
 maybeUseMove=""
 netemPartition=""
 netemConfig=""
+netemConfigFile=""
+netemCommand="add"
 
 command=$1
 [[ -n $command ]] || usage
@@ -216,6 +220,12 @@ while [[ -n $1 ]]; do
     elif [[ $1 = --config ]]; then
       netemConfig=$2
       shift 2
+    elif [[ $1 == --config-file ]]; then
+      netemConfigFile=$2
+      shift 2
+    elif [[ $1 == --netem-cmd ]]; then
+      netemCommand=$2
+      shift 2
     elif [[ $1 = --gpu-mode ]]; then
       gpuMode=$2
       case "$gpuMode" in
@@ -996,23 +1006,33 @@ logs)
   done
   ;;
 netem)
-  num_nodes=$((${#validatorIpList[@]}*netemPartition/100))
+  if [[ -n $netemConfigFile ]]; then
-  if [[ $((${#validatorIpList[@]}*netemPartition%100)) -gt 0 ]]; then
+    for ipAddress in "${validatorIpList[@]}"; do
-    num_nodes=$((num_nodes+1))
+      "$here"/scp.sh "$netemConfigFile" solana@"$ipAddress":~/solana
-  fi
+    done
-  if [[ "$num_nodes" -gt "${#validatorIpList[@]}" ]]; then
+    for i in "${!validatorIpList[@]}"; do
-    num_nodes=${#validatorIpList[@]}
+      "$here"/ssh.sh solana@"${validatorIpList[$i]}" 'solana/scripts/net-shaper.sh' \
-  fi
+      "$netemCommand" solana/"$netemConfigFile" "${#validatorIpList[@]}" "$i"
+    done
+  else
+    num_nodes=$((${#validatorIpList[@]}*netemPartition/100))
+    if [[ $((${#validatorIpList[@]}*netemPartition%100)) -gt 0 ]]; then
+      num_nodes=$((num_nodes+1))
+    fi
+    if [[ "$num_nodes" -gt "${#validatorIpList[@]}" ]]; then
+      num_nodes=${#validatorIpList[@]}
+    fi
 
-  # Stop netem on all nodes
+    # Stop netem on all nodes
-  for ipAddress in "${validatorIpList[@]}"; do
+    for ipAddress in "${validatorIpList[@]}"; do
-    "$here"/ssh.sh solana@"$ipAddress" 'solana/scripts/netem.sh delete < solana/netem.cfg || true'
+      "$here"/ssh.sh solana@"$ipAddress" 'solana/scripts/netem.sh delete < solana/netem.cfg || true'
-  done
+    done
 
-  # Start netem on required nodes
+    # Start netem on required nodes
-  for ((i=0; i<num_nodes; i++ )); do :
+    for ((i=0; i<num_nodes; i++ )); do :
-    "$here"/ssh.sh solana@"${validatorIpList[$i]}" "echo $netemConfig > solana/netem.cfg; solana/scripts/netem.sh add \"$netemConfig\""
+      "$here"/ssh.sh solana@"${validatorIpList[$i]}" "echo $netemConfig > solana/netem.cfg; solana/scripts/netem.sh add \"$netemConfig\""
-  done
+    done
+  fi
   ;;
 *)
   echo "Internal error: Unknown command: $command"