From c93ee250541320c3a9d7901b6157fabc21ae8bb1 Mon Sep 17 00:00:00 2001
From: Willy Tarreau <w@1wt.eu>
Date: Thu, 27 Feb 2025 10:44:32 +0100
Subject: [PATCH] MINOR: cpu-topo: add "only-node" and "drop-node" to cpu-set

These are processed after the topology is detected, and they allow to
restrict binding to or evict CPUs matching the indicated node(s).
---
 doc/configuration.txt      |  2 ++
 include/haproxy/cpu_topo.h |  3 +++
 src/cpu_topo.c             | 37 ++++++++++++++++++++++++++++++++++++-
 src/haproxy.c              |  3 +++
 4 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/doc/configuration.txt b/doc/configuration.txt
index 156fc5156..7cd4c30ba 100644
--- a/doc/configuration.txt
+++ b/doc/configuration.txt
@@ -1955,6 +1955,8 @@ cpu-set <directive>...
                          command for example.
     - drop-cpu <set>     do not bind to CPUs in this set
     - only-cpu <set>     do not bind to CPUs not in this set
+    - drop-node <set>    do not bind to CPUs belonging to this NUMA node
+    - only-node <set>    do not bind to CPUs not belonging to this NUMA node
 
 crt-base <dir>
   Assigns a default directory to fetch SSL certificates from when a relative
diff --git a/include/haproxy/cpu_topo.h b/include/haproxy/cpu_topo.h
index ae86b02f2..97a8a783d 100644
--- a/include/haproxy/cpu_topo.h
+++ b/include/haproxy/cpu_topo.h
@@ -34,6 +34,9 @@ void cpu_fixup_topology(void);
 /* compose clusters */
 void cpu_compose_clusters(void);
 
+/* apply remaining topology-based cpu set restrictions */
+void cpu_refine_cpusets(void);
+
 /* Detects CPUs that are bound to the current process. Returns the number of
  * CPUs detected or 0 if the detection failed.
  */
diff --git a/src/cpu_topo.c b/src/cpu_topo.c
index d4bdd9f7e..e976432fc 100644
--- a/src/cpu_topo.c
+++ b/src/cpu_topo.c
@@ -32,6 +32,9 @@ struct cpu_set_cfg {
 	/* CPU numbers to accept / reject */
 	struct hap_cpuset only_cpus;
 	struct hap_cpuset drop_cpus;
+	/* node numbers to accept / reject */
+	struct hap_cpuset only_nodes;
+	struct hap_cpuset drop_nodes;
 } cpu_set_cfg;
 
 /* Detects CPUs that are online on the system. It may rely on FS access (e.g.
@@ -826,6 +829,19 @@ void cpu_compose_clusters(void)
 	cpu_reorder_by_index(ha_cpu_topo, cpu_topo_maxcpus);
 }
 
+/* apply remaining topology-based cpu set restrictions */
+void cpu_refine_cpusets(void)
+{
+	int cpu;
+
+	/* remove CPUs in the drop-node set or not in the only-node set */
+	for (cpu = 0; cpu <= cpu_topo_lastcpu; cpu++) {
+		if ( ha_cpuset_isset(&cpu_set_cfg.drop_nodes, ha_cpu_topo[cpu].no_id) ||
+		    !ha_cpuset_isset(&cpu_set_cfg.only_nodes, ha_cpu_topo[cpu].no_id))
+			ha_cpu_topo[cpu].st |= HA_CPU_F_DONT_USE;
+	}
+}
+
 /* CPU topology detection below, OS-specific */
 
 #if defined(__linux__)
@@ -1191,6 +1207,22 @@ static int cfg_parse_cpu_set(char **args, int section_type, struct proxy *curpx,
 				ha_cpuset_and(&cpu_set_cfg.only_cpus, &tmp_cpuset);
 			arg++;
 		}
+		else if (strcmp(args[arg], "drop-node") == 0 || strcmp(args[arg], "only-node") == 0) {
+			if (!*args[arg + 1]) {
+				memprintf(err, "missing node set");
+				goto parse_err;
+			}
+
+			cpu_set_str[0] = args[arg + 1];
+			if (parse_cpu_set(cpu_set_str, &tmp_cpuset, err) != 0)
+				goto parse_err;
+
+			if (*args[arg] == 'd') // nodes to drop
+				ha_cpuset_or(&cpu_set_cfg.drop_nodes, &tmp_cpuset);
+			else // nodes to keep
+				ha_cpuset_and(&cpu_set_cfg.only_nodes, &tmp_cpuset);
+			arg++;
+		}
 		else {
 			/* fall back with default error message */
 			memprintf(err, "'%s' passed an unknown directive '%s'", args[0], args[arg]);
@@ -1215,7 +1247,7 @@ static int cfg_parse_cpu_set(char **args, int section_type, struct proxy *curpx,
 
  leave_with_err:
 	/* complete with supported directives */
-	memprintf(err, "%s (only 'reset', 'only-cpu', 'drop-cpu' supported).", *err);
+	memprintf(err, "%s (only 'reset', 'only-cpu', 'drop-cpu', 'only-node', 'drop-node' supported).", *err);
  leave:
 	return -1;
 }
@@ -1260,10 +1292,13 @@ static int cpu_topo_alloc(void)
 	/* pre-inizialize the configured CPU sets */
 	ha_cpuset_zero(&cpu_set_cfg.drop_cpus);
 	ha_cpuset_zero(&cpu_set_cfg.only_cpus);
+	ha_cpuset_zero(&cpu_set_cfg.drop_nodes);
+	ha_cpuset_zero(&cpu_set_cfg.only_nodes);
 
 	/* preset all CPUs in the "only-XXX" sets */
 	for (cpu = 0; cpu < cpu_topo_maxcpus; cpu++) {
 		ha_cpuset_set(&cpu_set_cfg.only_cpus, cpu);
+		ha_cpuset_set(&cpu_set_cfg.only_nodes, cpu);
 	}
 
 	return 1;
diff --git a/src/haproxy.c b/src/haproxy.c
index 812eca5d9..dd02a096a 100644
--- a/src/haproxy.c
+++ b/src/haproxy.c
@@ -2071,6 +2071,9 @@ static void step_init_2(int argc, char** argv)
 
 	/* compose clusters */
 	cpu_compose_clusters();
+
+	/* refine topology-based CPU sets */
+	cpu_refine_cpusets();
 #endif
 
 	/* detect the optimal thread-groups and nbthreads if not set */