[RFC PATCH 5/8] mm/memcontrol: add per-cpu stock for tier charge/uncharge

From: Qiqi Liu <liuqiqi-UOlijcLmZ/[email protected]>

Add tier_stock_pcp, a per-cpu cache with NR_TIER_STOCK slots, to cache
charge surplus for specific (memcg, tier_id) pairs.

The fast path, consume_tier_stock(), attempts to satisfy charges from
this cache. On a miss, try_charge_memcg_tier() charges in units of
max(MEMCG_CHARGE_BATCH, nr_pages), falling back to an exact nr_pages
charge when the batch does not fit, and refill_tier_stock() restocks
the surplus. Uncharging is handled by the same helper: refill_tier_stock()
drains a mismatched slot and caches the new pair.

Cached stock is drained via drain_all_tier_stock() during css_offline,
and from the charge path when a tier charge retries under limit pressure.

This is purely an optimization; the semantics remain identical to
per-page charging.

Signed-off-by: Qiqi Liu <liuqiqi-UOlijcLmZ/[email protected]>
---
 mm/memcontrol.c | 211 ++++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 203 insertions(+), 8 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 30f24604010c..3aa55c287248 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2318,11 +2318,33 @@ void drain_all_stock(struct mem_cgroup *root_memcg)
 	mutex_unlock(&percpu_charge_mutex);
 }
 
+/*
+ * Per-CPU tier charge/uncharge stock: caches up to NR_TIER_STOCK (memcg,
+ * tier_id) pairs so small charges/uncharges avoid atomics.
+ */
+#define NR_TIER_STOCK 7
+struct tier_stock_pcp {
+	local_lock_t lock;
+	uint8_t nr_pages[NR_TIER_STOCK];
+	struct mem_cgroup *cached[NR_TIER_STOCK];
+	int tier_id[NR_TIER_STOCK];
+	struct work_struct work;
+	unsigned long flags;
+	uint8_t drain_idx;
+};
+static DEFINE_PER_CPU(struct tier_stock_pcp, tier_stock);
+
+static void __drain_tier_stock(struct tier_stock_pcp *stock, int i);
+
 static int memcg_hotplug_cpu_dead(unsigned int cpu)
 {
 	/* no need for the local lock */
+	int i;
+
 	drain_obj_stock(&per_cpu(obj_stock, cpu));
 	drain_stock_fully(&per_cpu(memcg_stock, cpu));
+	for (i = 0; i < NR_TIER_STOCK; i++)
+		__drain_tier_stock(&per_cpu(tier_stock, cpu), i);
 
 	return 0;
 }
@@ -2467,12 +2489,6 @@ static void memcg_charge_tier(struct mem_cgroup *memcg, struct folio *folio,
 	memcg_charge_tier_id(memcg, node_to_tier_id(folio_nid(folio)), nr_pages);
 }
 
-static void memcg_uncharge_tier(struct mem_cgroup *memcg, struct folio *folio,
-			       unsigned long nr_pages)
-{
-	memcg_uncharge_tier_id(memcg, node_to_tier_id(folio_nid(folio)), nr_pages);
-}
-
 static void memcg_free_tier_counters(struct mem_cgroup *memcg)
 {
 	struct memcg_tier_counter *tc, *tmp;
@@ -2956,14 +2972,173 @@ static int try_charge_memcg(struct mem_cgroup *memcg, gfp_t gfp_mask,
 	return 0;
 }
 
+/* Drain cached tier charge from a per-CPU stock slot.
+ * The drained count is returned to the tier's page_counter,
+ * and the memcg css reference is put.
+ */
+static void __drain_tier_stock(struct tier_stock_pcp *stock, int i)
+{
+	struct memcg_tier_counter *tc;
+	struct mem_cgroup *old = READ_ONCE(stock->cached[i]);
+	uint8_t stock_pages;
+
+	if (!old)
+		return;
+	stock_pages = READ_ONCE(stock->nr_pages[i]);
+	if (stock_pages) {
+		rcu_read_lock();
+		tc = memcg_tier_counter_find(old, READ_ONCE(stock->tier_id[i]));
+		if (tc)
+			page_counter_uncharge(&tc->counter, stock_pages);
+		rcu_read_unlock();
+		WRITE_ONCE(stock->nr_pages[i], 0);
+	}
+	css_put(&old->css);
+	WRITE_ONCE(stock->cached[i], NULL);
+	WRITE_ONCE(stock->tier_id[i], -1);
+}
+
+static void drain_local_tier_stock(struct work_struct *dummy)
+{
+	struct tier_stock_pcp *stock;
+	int i;
+
+	if (WARN_ON_ONCE(!in_task()))
+		return;
+	local_lock(&tier_stock.lock);
+	stock = this_cpu_ptr(&tier_stock);
+	for (i = 0; i < NR_TIER_STOCK; i++)
+		__drain_tier_stock(stock, i);
+	clear_bit(FLUSHING_CACHED_CHARGE, &stock->flags);
+	local_unlock(&tier_stock.lock);
+}
+
+/* Drain per-CPU tier stock entries matching @memcg on all CPUs. */
+static void drain_all_tier_stock(struct mem_cgroup *memcg)
+{
+	int cpu, curcpu, i;
+
+	if (!mutex_trylock(&percpu_charge_mutex))
+		return;
+	migrate_disable();
+	curcpu = smp_processor_id();
+	for_each_online_cpu(cpu) {
+		struct tier_stock_pcp *stock = &per_cpu(tier_stock, cpu);
+
+		if (test_bit(FLUSHING_CACHED_CHARGE, &stock->flags))
+			continue;
+		for (i = 0; i < NR_TIER_STOCK; i++) {
+			if (READ_ONCE(stock->cached[i]) != memcg)
+				continue;
+			if (!test_and_set_bit(FLUSHING_CACHED_CHARGE,
+				      &stock->flags)) {
+				if (cpu == curcpu)
+					drain_local_tier_stock(&stock->work);
+				else
+					schedule_drain_work(cpu, &stock->work);
+			}
+			break;
+		}
+	}
+	migrate_enable();
+	mutex_unlock(&percpu_charge_mutex);
+}
+
+/* Consume @nr_pages from the per-CPU tier stock if a matching slot has enough surplus */
+static bool consume_tier_stock(struct mem_cgroup *memcg, int tier_id,
+				      unsigned int nr_pages)
+{
+	struct tier_stock_pcp *stock;
+	bool ret = false;
+	int i;
+	uint8_t pages;
+
+	BUILD_BUG_ON(MEMCG_CHARGE_BATCH > S8_MAX);
+
+	if (nr_pages > MEMCG_CHARGE_BATCH)
+		return false;
+
+	local_lock(&tier_stock.lock);
+	stock = this_cpu_ptr(&tier_stock);
+
+	for (i = 0; i < NR_TIER_STOCK; i++) {
+		if (READ_ONCE(stock->cached[i]) == memcg &&
+		    READ_ONCE(stock->tier_id[i]) == tier_id) {
+			pages = READ_ONCE(stock->nr_pages[i]);
+			if (pages >= nr_pages) {
+				WRITE_ONCE(stock->nr_pages[i], pages - nr_pages);
+				ret = true;
+			}
+			break;
+		}
+	}
+
+	local_unlock(&tier_stock.lock);
+	return ret;
+}
+
+/* Refund @nr_pages to the per-CPU tier stock. */
+static void refill_tier_stock(struct mem_cgroup *memcg, int tier_id,
+				      unsigned int nr_pages)
+{
+	struct memcg_tier_counter *tc;
+	struct tier_stock_pcp *stock;
+	int empty_slot = -1;
+	uint8_t pages;
+	int i;
+
+	/* Too big to cache: direct uncharge, leave the stock untouched. */
+	if (nr_pages > MEMCG_CHARGE_BATCH) {
+		rcu_read_lock();
+		tc = memcg_tier_counter_find(memcg, tier_id);
+		if (tc)
+			page_counter_uncharge(&tc->counter, nr_pages);
+		rcu_read_unlock();
+		return;
+	}
+
+	local_lock(&tier_stock.lock);
+	stock = this_cpu_ptr(&tier_stock);
+
+	for (i = 0; i < NR_TIER_STOCK; i++) {
+		if (!READ_ONCE(stock->cached[i]) && empty_slot == -1)
+			empty_slot = i;
+		if (READ_ONCE(stock->cached[i]) == memcg &&
+		    READ_ONCE(stock->tier_id[i]) == tier_id) {
+			pages = READ_ONCE(stock->nr_pages[i]) + nr_pages;
+			WRITE_ONCE(stock->nr_pages[i], pages);
+			if (pages > MEMCG_CHARGE_BATCH)
+				__drain_tier_stock(stock, i);
+			goto out;
+		}
+	}
+
+	/* Mismatch: pick a slot (empty or evict), drain, cache new. */
+	i = empty_slot;
+	if (i == -1) {
+		i = stock->drain_idx++;
+		if (stock->drain_idx == NR_TIER_STOCK)
+			stock->drain_idx = 0;
+	}
+	__drain_tier_stock(stock, i);
+	css_get(&memcg->css);
+	WRITE_ONCE(stock->cached[i], memcg);
+	WRITE_ONCE(stock->tier_id[i], tier_id);
+	WRITE_ONCE(stock->nr_pages[i], nr_pages);
+out:
+	local_unlock(&tier_stock.lock);
+}
+
 static int try_charge_memcg_tier(struct mem_cgroup *memcg, gfp_t gfp_mask,
 	unsigned int nr_pages, int tier_id)
 {
 	struct memcg_tier_counter *tc;
 	struct page_counter *counter;
+	unsigned int batch = max(MEMCG_CHARGE_BATCH, nr_pages);
 	int nr_retries = MAX_RECLAIM_RETRIES;
 	unsigned long nr_reclaimed = 0;
 	bool passed_oom = false;
+	bool drained = false;
 	nodemask_t nodes, *nmp = NULL;
 
 	if (tier_id < 0)
@@ -2976,9 +3151,16 @@ static int try_charge_memcg_tier(struct mem_cgroup *memcg, gfp_t gfp_mask,
 		return 0;
 
 retry:
-	if (page_counter_try_charge(&tc->counter, nr_pages, &counter))
+	if (consume_tier_stock(memcg, tier_id, nr_pages))
+		return 0;
+	if (page_counter_try_charge(&tc->counter, batch, &counter))
 		goto success;
 
+	if (batch > nr_pages) {
+		batch = nr_pages;
+		goto retry;
+	}
+
 	/* Over max -> reclaim. */
 	if (unlikely(current->flags & PF_MEMALLOC))
 		goto force;
@@ -2994,6 +3176,13 @@ static int try_charge_memcg_tier(struct mem_cgroup *memcg, gfp_t gfp_mask,
 
 	if (page_counter_read(&tc->counter) + nr_pages <= READ_ONCE(tc->counter.max))
 		goto retry;
+
+	if (!drained) {
+		drain_all_tier_stock(memcg);
+		drained = true;
+		goto retry;
+	}
+
 	if (gfp_mask & __GFP_NORETRY)
 		goto nomem;
 	if (nr_reclaimed && nr_pages <= (1 << PAGE_ALLOC_COSTLY_ORDER))
@@ -3011,6 +3200,8 @@ static int try_charge_memcg_tier(struct mem_cgroup *memcg, gfp_t gfp_mask,
 	}
 	goto nomem;
 success:
+	if (batch > nr_pages)
+		refill_tier_stock(memcg, tier_id, batch - nr_pages);
 	do {
 		struct memcg_tier_counter *tc_this;
 
@@ -4575,6 +4766,7 @@ static void mem_cgroup_css_offline(struct cgroup_subsys_state *css)
 	lru_gen_offline_memcg(memcg);
 
 	drain_all_stock(memcg);
+	drain_all_tier_stock(memcg);
 
 	mem_cgroup_private_id_put(memcg, 1);
 }
@@ -5565,7 +5757,8 @@ static void uncharge_folio(struct folio *folio, struct uncharge_gather *ug)
 		/* LRU pages aren't accounted at the root level */
 		if (!obj_cgroup_is_root(objcg)) {
 			ug->nr_memory += nr_pages;
-			memcg_uncharge_tier(obj_cgroup_memcg(objcg), folio,
+			refill_tier_stock(obj_cgroup_memcg(objcg),
+				node_to_tier_id(folio_nid(folio)),
 				nr_pages);
 		}
 		ug->pgpgout++;
@@ -5893,6 +6086,8 @@ int __init mem_cgroup_init(void)
 			  drain_local_memcg_stock);
 		INIT_WORK(&per_cpu_ptr(&obj_stock, cpu)->work,
 			  drain_local_obj_stock);
+		INIT_WORK(&per_cpu(tier_stock, cpu).work,
+			  drain_local_tier_stock);
 	}
 
 	memcg_size = struct_size_t(struct mem_cgroup, nodeinfo, nr_node_ids);
-- 
2.43.0