[RFC PATCH 3/8] mm/memcontrol: add per-tier page counter infrastructure and lifecycle

From: Qiqi Liu <[email protected]>

Introduce struct memcg_tier_counter, a hierarchical page_counter for
each (memcg, tier) pair. Each counter tracks memory usage on a specific
NUMA memory tier within a cgroup.

The counters are managed under the memcg lifecycle:
  - Lookup: lockless and RCU-protected.
  - Allocation: GFP_KERNEL, during css_alloc and memory tier hotplug.
  - Parent hierarchy: established explicitly at creation time.
  - Destruction: in css_free, alongside the memcg.

During css_alloc, counters are pre-allocated for the tiers of all
currently online nodes, ensuring the charge hot path never needs to
allocate memory.

This patch only lays the groundwork; the counters are not yet
integrated into charge/uncharge. The actual accounting logic is added
in a subsequent patch.

Signed-off-by: Qiqi Liu <[email protected]>
---
 include/linux/memcontrol.h |  11 +++
 mm/memcontrol.c            | 141 +++++++++++++++++++++++++++++++++++++
 2 files changed, 152 insertions(+)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index e1f46a0016fc..c0f5929a87bf 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -199,6 +199,14 @@ struct obj_cgroup {
  * statistics based on the statistics developed by Rik Van Riel for clock-pro,
  * to help the administrator determine what knobs to tune.
  */
+
+struct memcg_tier_counter {
+	struct page_counter counter;
+	int tier_id;
+	struct list_head list;
+	struct rcu_head rcu;
+};
+
 struct mem_cgroup {
 	struct cgroup_subsys_state css;
 
@@ -320,6 +328,9 @@ struct mem_cgroup {
 	spinlock_t event_list_lock;
 #endif /* CONFIG_MEMCG_V1 */
 
+	spinlock_t tier_lock;
+	struct list_head tier_counters;
+
 	struct mem_cgroup_per_node *nodeinfo[];
 };
 
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 6dc4888a90f3..70efe01bc36f 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -28,6 +28,8 @@
 #include <linux/cgroup-defs.h>
 #include <linux/page_counter.h>
 #include <linux/memcontrol.h>
+#include <linux/memory-tiers.h>
+#include <linux/node.h>
 #include <linux/cgroup.h>
 #include <linux/cpuset.h>
 #include <linux/sched/mm.h>
@@ -2360,6 +2362,89 @@ static void high_work_func(struct work_struct *work)
 	reclaim_high(memcg, MEMCG_CHARGE_BATCH, GFP_KERNEL);
 }
 
+/*
+ * Find a tier counter for a given memcg and tier ID.
+ *
+ * Context: Caller must hold either:
+ * - The RCU read lock, for lockless lookup in the fast (charge) path.
+ * - memcg->tier_lock, for modifications in the slow (creation/free) path.
+ */
+static struct memcg_tier_counter *
+memcg_tier_counter_find(struct mem_cgroup *memcg, int tier_id)
+{
+	struct memcg_tier_counter *tc;
+
+	list_for_each_entry_rcu(tc, &memcg->tier_counters, list,
+				lockdep_is_held(&memcg->tier_lock))
+		if (tc->tier_id == tier_id)
+			return tc;
+	return NULL;
+}
+
+static struct page_counter *
+memcg_tier_parent_link(struct mem_cgroup *parent, int tier_id)
+{
+	struct memcg_tier_counter *ptc;
+	struct page_counter *pc;
+
+	if (!parent || mem_cgroup_is_root(parent))
+		return NULL;
+
+	rcu_read_lock();
+	ptc = memcg_tier_counter_find(parent, tier_id);
+	pc = ptc ? &ptc->counter : NULL;
+	rcu_read_unlock();
+	return pc;
+}
+
+static int memcg_tier_counter_create(struct mem_cgroup *memcg,
+				     struct mem_cgroup *parent, int tier_id)
+{
+	struct page_counter *parent_cnt;
+	struct memcg_tier_counter *new;
+
+	/* Fast path (lockless RCU read): already exists -> nothing to do. */
+	rcu_read_lock();
+	if (memcg_tier_counter_find(memcg, tier_id)) {
+		rcu_read_unlock();
+		return 0;
+	}
+	rcu_read_unlock();
+
+	parent_cnt = memcg_tier_parent_link(parent, tier_id);
+
+	new = kzalloc_obj(*new);
+	if (!new)
+		return -ENOMEM;
+
+	new->tier_id = tier_id;
+	page_counter_init(&new->counter, parent_cnt, false);
+	page_counter_set_high(&new->counter, PAGE_COUNTER_MAX);
+	INIT_LIST_HEAD(&new->list);
+
+	spin_lock(&memcg->tier_lock);
+	if (memcg_tier_counter_find(memcg, tier_id)) {
+		spin_unlock(&memcg->tier_lock);
+		kfree(new);
+		return 0;
+	}
+	list_add_tail_rcu(&new->list, &memcg->tier_counters);
+	spin_unlock(&memcg->tier_lock);
+	return 0;
+}
+
+static void memcg_free_tier_counters(struct mem_cgroup *memcg)
+{
+	struct memcg_tier_counter *tc, *tmp;
+
+	spin_lock(&memcg->tier_lock);
+	list_for_each_entry_safe(tc, tmp, &memcg->tier_counters, list) {
+		list_del_rcu(&tc->list);
+		kfree_rcu(tc, rcu);
+	}
+	spin_unlock(&memcg->tier_lock);
+}
+
 /*
  * Clamp the maximum sleep time per allocation batch to 2 seconds. This is
  * enough to still cause a significant slowdown in most cases, while still
@@ -4129,6 +4214,8 @@ static struct mem_cgroup *mem_cgroup_alloc(struct mem_cgroup *parent)
 		goto fail;
 
 	INIT_WORK(&memcg->high_work, high_work_func);
+	spin_lock_init(&memcg->tier_lock);
+	INIT_LIST_HEAD(&memcg->tier_counters);
 	vmpressure_init(&memcg->vmpressure);
 	INIT_LIST_HEAD(&memcg->memory_peaks);
 	INIT_LIST_HEAD(&memcg->swap_peaks);
@@ -4178,6 +4265,20 @@ mem_cgroup_css_alloc(struct cgroup_subsys_state *parent_css)
 
 		page_counter_init(&memcg->memory, &parent->memory, memcg_on_dfl);
 		page_counter_init(&memcg->swap, &parent->swap, false);
+
+		{
+			int nid, tid;
+
+			for_each_online_node(nid) {
+				tid = node_to_tier_id(nid);
+				if (tid != -1 && memcg_tier_counter_create(memcg,
+							mem_cgroup_from_css(parent_css), tid)) {
+					memcg_free_tier_counters(memcg);
+					mem_cgroup_free(memcg);
+					return ERR_PTR(-ENOMEM);
+				}
+			}
+		}
 #ifdef CONFIG_MEMCG_V1
 		memcg->memory.track_failcnt = !memcg_on_dfl;
 		WRITE_ONCE(memcg->oom_kill_disable, READ_ONCE(parent->oom_kill_disable));
@@ -4338,6 +4439,7 @@ static void mem_cgroup_css_free(struct cgroup_subsys_state *css)
 
 	vmpressure_cleanup(&memcg->vmpressure);
 	cancel_work_sync(&memcg->high_work);
+	memcg_free_tier_counters(memcg);
 	memcg1_remove_from_trees(memcg);
 	free_shrinker_info(memcg);
 	mem_cgroup_free(memcg);
@@ -5534,6 +5636,41 @@ __setup("cgroup.memory=", cgroup_memory);
  * basically everything that doesn't depend on a specific mem_cgroup structure
  * should be initialized from here.
  */
+#if defined(CONFIG_MEMORY_HOTPLUG) && defined(CONFIG_NUMA)
+/*
+ * Below MEMTIER_HOTPLUG_PRI (100): run after memory-tiers has set or
+ * cleared the node's tier association.
+ */
+#define MEMCG_TIER_NODE_PRI		90
+
+/* Memory hotplug callback: a node came online with a potentially new tier
+ * (e.g. CXL hotplug). Ensure every online memcg has a counter for this tier.
+ * Existing tiers hit in the RCU lookup, so this path does not allocate.
+ */
+static int __meminit memcg_tier_hotplug_cb(struct notifier_block *self,
+					    unsigned long action, void *_arg)
+{
+	struct node_notify *nn = _arg;
+	struct mem_cgroup *memcg;
+	int tid;
+
+	if (action != NODE_ADDED_FIRST_MEMORY)
+		return notifier_from_errno(0);
+
+	tid = node_to_tier_id(nn->nid);
+	if (tid < 0)
+		return notifier_from_errno(0);
+
+	for_each_mem_cgroup(memcg) {
+		if (!mem_cgroup_is_root(memcg) &&
+		    memcg_tier_counter_create(memcg, parent_mem_cgroup(memcg), tid))
+			pr_warn_ratelimited("memcg: tier %d counter alloc failed;"
+					     " tier accounting degraded\n", tid);
+	}
+	return notifier_from_errno(0);
+}
+#endif
+
 int __init mem_cgroup_init(void)
 {
 	unsigned int memcg_size;
@@ -5553,6 +5690,10 @@ int __init mem_cgroup_init(void)
 	memcg_wq = alloc_workqueue("memcg", WQ_PERCPU, 0);
 	WARN_ON(!memcg_wq);
 
+#if defined(CONFIG_MEMORY_HOTPLUG) && defined(CONFIG_NUMA)
+	hotplug_node_notifier(memcg_tier_hotplug_cb, MEMCG_TIER_NODE_PRI);
+#endif
+
 	for_each_possible_cpu(cpu) {
 		INIT_WORK(&per_cpu_ptr(&memcg_stock, cpu)->work,
 			  drain_local_memcg_stock);
-- 
2.43.0