[PATCH v8 7/7] mm/memcontrol: restrict memory stock to leaf cgroups on v2

[Joshua Hahn <joshua.hahnjy@gmail.com>]

In the old memcg_stock world, non-leaf cgroups in cgroup v2 were never
directly charged (except for reparented objcgs), since processes can
only live in leaf cgroups. This meant that stocking was exclusive to
leaf cgroups.

In the new page_counter_stock world, stock belongs to each page_counter.
To get parity with vanilla stock, prevent stock from accumulating in
page_counters once they have children. This makes sure that stock does
not get stranded in the hierarchy, unable to be used by processes that
actually charge. When a parent becomes a leaf again, re-enable stock.

Prevent stock from accumulating in non-leaf page_counters in three ways:
1. Clear the parent's memory stock pointer to make sure the stock cannot
   get refilled in the future.
2. Schedule drains for non-isolated CPUs that stock the new parent.
3. CPUs that are not drained due to either isolation or racing with a
   refill are cleaned up lazily. The next local stock miss on that CPU
   scans for any stock that has been disabled, and will drain its stock.
   Note that this happens under the local stock lock, but this is rare.

Signed-off-by: Joshua Hahn <joshua.hahnjy@gmail.com>
---
 mm/memcontrol.c   | 28 ++++++++++++++++++++++++++++
 mm/page_counter.c |  8 +++++++-
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index aac987ebaab6f..a7549fd06d93e 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2185,6 +2185,22 @@ static void schedule_stock_drain(struct page_counter_stock_pcp __percpu *stock,
 		clear_bit(FLUSHING_CACHED_CHARGE, &pcp_stock->flags);
 }
 
+static void mem_cgroup_disable_stock(struct mem_cgroup *memcg)
+{
+	struct page_counter_stock_pcp __percpu *stock;
+	int cpu, curcpu;
+
+	stock = xchg(&memcg->memory.stock, NULL);
+	if (!stock)
+		return;
+
+	migrate_disable();
+	curcpu = smp_processor_id();
+	for_each_online_cpu(cpu)
+		schedule_stock_drain(stock, &memcg->css, cpu, curcpu);
+	migrate_enable();
+}
+
 /*
  * Drains all per-CPU charge caches for given root_memcg resp. subtree
  * of the hierarchy under it.
@@ -4308,6 +4324,9 @@ static int mem_cgroup_css_online(struct cgroup_subsys_state *css)
 	 */
 	xa_store(&mem_cgroup_private_ids, memcg->private_id, memcg->private_id_objcg, GFP_KERNEL);
 
+	if (cgroup_subsys_on_dfl(memory_cgrp_subsys) && css->parent)
+		mem_cgroup_disable_stock(parent_mem_cgroup(memcg));
+
 	return 0;
 free_objcg:
 	for_each_node(nid) {
@@ -4360,6 +4379,15 @@ static void mem_cgroup_css_offline(struct cgroup_subsys_state *css)
 	drain_all_stock(memcg);
 
 	mem_cgroup_private_id_kill(memcg);
+
+	if (cgroup_subsys_on_dfl(memory_cgrp_subsys)) {
+		struct mem_cgroup *parent = parent_mem_cgroup(memcg);
+
+		if (parent && !css_is_dying(&parent->css) &&
+		    atomic_read(&parent->css.online_cnt) == 2 &&
+		    !mem_cgroup_is_root(parent))
+			WRITE_ONCE(parent->memory.stock, &memory_stock);
+	}
 }
 
 static void mem_cgroup_css_released(struct cgroup_subsys_state *css)
diff --git a/mm/page_counter.c b/mm/page_counter.c
index e1feb2bd1e225..7a3b7cea0ff94 100644
--- a/mm/page_counter.c
+++ b/mm/page_counter.c
@@ -158,6 +158,7 @@ static void page_counter_drain_stock(struct page_counter_stock_pcp *stock,
  * Consume the cached charge if enough nr_pages are present, otherwise return
  * failure. Also return failure for charge requests larger than
  * PAGE_COUNTER_STOCK_BATCH or if the local lock is already taken.
+ * Drain disabled cached counters while looking up the requested counter.
  */
 static bool page_counter_consume_stock(struct page_counter *counter,
 				       unsigned long nr_pages)
@@ -176,12 +177,17 @@ static bool page_counter_consume_stock(struct page_counter *counter,
 	for (i = 0; i < NR_PAGE_COUNTER_STOCK; i++) {
 		struct page_counter *cached = pcp_stock->cached[i];
 
+		/* Drain any stock that is disabled but hasn't drained yet */
+		if (cached && !READ_ONCE(cached->stock)) {
+			page_counter_drain_stock(pcp_stock, i);
+			continue;
+		}
 		if (counter != cached)
 			continue;
 
 		stock_pages = pcp_stock->nr_pages[i];
 		if (unlikely(stock_pages < nr_pages))
-			break;
+			continue;
 
 		stock_pages -= nr_pages;
 		WRITE_ONCE(pcp_stock->nr_pages[i], stock_pages);
-- 
2.53.0-Meta