From 2987513b5eabae607da193093c4dff81b3ce017e Mon Sep 17 00:00:00 2001
From: Daniel Hodges <hodges.daniel.scott@gmail.com>
Date: Mon, 18 Nov 2024 06:41:10 -0800
Subject: [PATCH] scx_layered: Add netdev IRQ balancing node support

Add support for setting netdev IRQ balancing that is NUMA aware.

Signed-off-by: Daniel Hodges <hodges.daniel.scott@gmail.com>
---
 scheds/rust/scx_layered/src/main.rs | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/scheds/rust/scx_layered/src/main.rs b/scheds/rust/scx_layered/src/main.rs
index faa9562..f9ac689 100644
--- a/scheds/rust/scx_layered/src/main.rs
+++ b/scheds/rust/scx_layered/src/main.rs
@@ -1221,6 +1221,7 @@ struct Scheduler<'a> {
     nr_layer_cpus_ranges: Vec<(usize, usize)>,
     processing_dur: Duration,
 
+    topo: Topology,
     netdevs: BTreeMap<String, NetDev>,
     stats_server: StatsServer<StatsReq, StatsRes>,
 }
@@ -1536,6 +1537,7 @@ impl<'a> Scheduler<'a> {
             proc_reader,
             skel,
 
+            topo,
             netdevs,
             stats_server,
         };
@@ -1563,12 +1565,29 @@ impl<'a> Scheduler<'a> {
         }
 
         for (iface, netdev) in self.netdevs.iter_mut() {
+            let node = self
+                .topo
+                .nodes()
+                .into_iter()
+                .take_while(|n| n.id() == netdev.node)
+                .next()
+                .ok_or_else(|| anyhow!("Failed to get netdev node"))?;
+            let node_cpus = node.span();
             for (irq, irqmask) in netdev.irqs.iter_mut() {
                 irqmask.clear();
                 for cpu in available_cpus.iter_ones() {
+                    if !node_cpus.test_cpu(cpu) {
+                        continue;
+                    }
                     let _ = irqmask.set_cpu(cpu);
                 }
                 trace!("{} updating irq {} cpumask {:?}", iface, irq, irqmask);
+                // If no CPUs are available in the node then spread the load across the node
+                if irqmask.weight() == 0 {
+                    for cpu in node_cpus.as_raw_bitvec().iter_ones() {
+                        let _ = irqmask.set_cpu(cpu);
+                    }
+                }
             }
             netdev.apply_cpumasks()?;
         }