[detect_numa] Filter empty NUMA nodes (#29)

* Add development option to provide command for lscpu

* Filter out empty NUMA nodes by default, added option to include them in a count if desired

detect_numa

@@ -2,6 +2,8 @@
 
 get_nodes=0
 get_cpus=0
+remove_empty_nodes=1
+lscpu_command="lscpu"
 node=""
 
 usage()
@@ -23,13 +25,20 @@ usage()
     echo -e "\t-n/--node"
     echo -e "\t\tRestrict CPUs listed by --cpu-list to certain nodes"
     echo -e "\t\t\tSupports ranges (0-5), comma separated values (0,2,4), and single digits (1)"
+
+    echo -e "\t--include-empty"
+    echo -e "\t\tInclude empty NUMA nodes in --node-count output"
+
+    echo -e "\t-i/--input"
+    echo -e "\t\tProvide a command to gather information in lscpu format"
+    echo -e "\t\tThis is intended for development"
     exit 0
 }
 
 # Verifies that NUMA exists on the system
 has_numa()
 {
-    local numa_nodes=`lscpu | grep NUMA | grep CPU | wc -l`
+    local numa_nodes=`$lscpu_command | grep NUMA | grep CPU | wc -l`
     if [ -z "$numa_nodes" ] || [ $numa_nodes -le 0 ]; then
         echo 0
     else
@@ -40,9 +49,14 @@ has_numa()
 # Fetches the number of NUMA nodes
 get_num_nodes()
 {
-    local numa_nodes=`lscpu | grep NUMA | grep CPU | wc -l`
-    if [ "`has_numa`" -eq 0 ]; then
-        numa_nodes=1
+    local numa_nodes=1
+    if [ "`has_numa`" -eq 1 ]; then
+        local numa_node_info=`$lscpu_command | grep NUMA | grep CPU`
+        if [ $remove_empty_nodes -ne 0 ]; then
+            numa_nodes=`echo "$numa_node_info" | awk '{ print $4 }' | grep -v "^$" | wc -l`
+        else
+            numa_nodes=`echo "$numa_node_info" | wc -l`
+        fi
     fi
 
     echo $numa_nodes
@@ -65,7 +79,7 @@ fetch_cpus() {
         echo "Error: NUMA Node $node does not exist" > /dev/stderr
         exit 1
     fi
-    lscpu | grep -E "$prefix$node" | sort | awk '{ print $4 }'
+    $lscpu_command | grep -E "$prefix$node" | sort | awk '{ print $4 }'
 }
 
 NOARG_OPTS=(
@@ -76,18 +90,21 @@ NOARG_OPTS=(
     "cpu-list"
     "list-cpu"
     "node-count"
+    "include-empty"
 )
 
 ARG_OPTS=(
     "node"
     "n"
+    "i"
+    "input"
 )
 
 opts=$(getopt \
     --longoptions "$(printf "%s," "${NOARG_OPTS[@]}")" \
     --longoptions "$(printf "%s:," "${ARG_OPTS[@]}")" \
     --name "$(basename "$0")" \
-    --options "hn:" \
+    --options "hi:n:" \
     -- "$@"
 )
 
@@ -114,6 +131,14 @@ while [[ $# -gt 0 ]]; do
         get_nodes=1
         shift 1
     ;;
+    -i | --input)
+        lscpu_command=$2
+        shift 2
+    ;;
+    --include-empty)
+        remove_empty_nodes=0
+        shift 1
+    ;;
     --)
         break
     ;;

tests/resources/detect_numa/empty_numa.txt

@@ -0,0 +1,11 @@
+NUMA:
+  NUMA node(s):         9
+  NUMA node0 CPU(s):    0-71
+  NUMA node1 CPU(s):
+  NUMA node2 CPU(s):
+  NUMA node3 CPU(s):
+  NUMA node4 CPU(s):
+  NUMA node5 CPU(s):
+  NUMA node6 CPU(s):
+  NUMA node7 CPU(s):
+  NUMA node8 CPU(s):

tests/test_detect_numa

@@ -27,4 +27,10 @@ assert_raises "$exec_file --node $((num_nodes)) --cpu-list" 1 # num_nodes is inv
 # Ensure that the --node flag fails when a negative NUMA node is provided
 assert_raises "$exec_file --node -1 --cpu-list" 1
 
+# Ensure that empty NUMA nodes are not counted by default
+assert "$exec_file -i \"cat tests/resources/detect_numa/empty_numa.txt\"" 1
+
+# Ensure that --include-empty works as expected
+assert "$exec_file --include-empty -i \"cat tests/resources/detect_numa/empty_numa.txt\"" 9
+
 assert_end