[RFC PATCH v3 03/17] sched/topology: Add an iterator for cached NUMA node affinity rows

[Jianyong Wu]

Provide for_each_affi_node() to walk the cached node affinity sequence.
Visit the local node first, then remote nodes by raw NUMA distance and
triangular tie rank. Include possible nodes without CPUs and skip holes
in the node ID space.

Pin one immutable row under RCU for the entire walk so rebuilding the
table cannot change the ordering midway through a migration decision.
If the table is unavailable, visit only the local node.

Signed-off-by: Jianyong Wu <wujianyong@xxxxxxxx>
---
 kernel/sched/sched.h    | 24 ++++++++++++++++++++++++
 kernel/sched/topology.c | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 4dba127f1d4f..5d3fa5d70d16 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -4140,6 +4140,30 @@ sched_cache_numa_key_cmp(struct sched_cache_numa_key a,
 	return 0;
 }
 
+struct affi_node_order_iter {
+	const int *order;
+	int root;
+	int nr_nodes;
+	int index;
+};
+
+/* Hold RCU continuously from begin() through the last next() call. */
+extern struct affi_node_order_iter affi_node_order_begin(int cpu);
+extern int affi_node_order_next(struct affi_node_order_iter *iter);
+
+/*
+ * Walk possible nodes, local first, then by raw NUMA distance and triangular
+ * tie rank. Pin one immutable row under RCU for the entire walk. The body
+ * must not sleep; break and return release the scoped lock. If the cached
+ * order is unavailable, visit only the local node.
+ */
+#define for_each_affi_node(cpu, node)					\
+	scoped_guard(rcu)						\
+		for (struct affi_node_order_iter __sn_iter =		\
+			     affi_node_order_begin((cpu));		\
+		     ({ int __sn_node = affi_node_order_next(&__sn_iter); \
+			__sn_node >= 0 && ((node) = __sn_node, true); });)
+
 /* Node IDs in, zero-based rank out; -1 if unavailable or invalid. */
 extern int affi_node_rank_of(int prefer_node, int node);
 
diff --git a/kernel/sched/topology.c b/kernel/sched/topology.c
index 6b7b38eb7cd3..8d0fd0f8a67a 100644
--- a/kernel/sched/topology.c
+++ b/kernel/sched/topology.c
@@ -1101,6 +1101,38 @@ void __init sched_cache_init_node_order(void)
 	sched_cache_numa_distance_update_end(true);
 }
 
+struct affi_node_order_iter affi_node_order_begin(int cpu)
+{
+	struct affi_node_order_topology *topo;
+	struct affi_node_order_iter iter = { };
+
+	RCU_LOCKDEP_WARN(!rcu_read_lock_any_held(), "node order read without RCU");
+	if (cpu < 0 || cpu >= nr_cpu_ids || !cpu_possible(cpu))
+		return iter;
+
+	iter.root = cpu_to_node(cpu);
+	if (iter.root < 0 || iter.root >= nr_node_ids || !node_possible(iter.root))
+		return iter;
+
+	iter.nr_nodes = 1;
+	topo = rcu_dereference_all(affi_node_order);
+	if (topo && iter.root < topo->size) {
+		iter.nr_nodes = topo->nr_nodes;
+		iter.order = topo->order + (size_t)iter.root * iter.nr_nodes;
+	}
+	return iter;
+}
+
+int affi_node_order_next(struct affi_node_order_iter *iter)
+{
+	int idx = iter->index;
+
+	if (idx >= iter->nr_nodes)
+		return -1;
+	iter->index++;
+	return iter->order ? iter->order[idx] : iter->root;
+}
+
 /*
  * Return @node's zero-based rank in @prefer_node's affinity order, with
  * @prefer_node itself ranked first. Both arguments are node IDs, not CPUs.
-- 
2.34.1