In the old memcg_stock world, non-leaf cgroups in cgroup v2 usually did
not get directly charged, since new processes can only live in leaf
cgroups. This meant that stocking was usuaully only for leaf cgroups.

In the new page_counter_stock world, stock belongs to each page_counter.
To prevent unused stock from accumulating in non-leaf cgroups, prevent
stock refills on cgroups once they have children, and schedule drains on
non-isolated CPUs that stock the parent.

CPUs that are not drained because they are isolated or because there was
a race with a refill may keep a stale slot for the parent. This is OK,
since the slot is never consumed or refilled, and will eventually get
drained like any other stale stock through stock eviction or draining.

One behavioral change is that processes that do charge directly to
interior cgroups now do not stock. Reparented objcgs or swap-in
charges used to refill and leave stock in non-leaf cgroups, but we drop
this in favor of not stranding unusable charges in parent cgroups.

Signed-off-by: Joshua Hahn <joshua.hahnjy@gmail.com>
---
 mm/memcontrol.c | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index cb1f46b457ae8..15b5e0692c436 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -2189,6 +2189,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.
@@ -4314,6 +4330,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) {
@@ -4366,6 +4385,20 @@ 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);
+
+		/*
+		 * The parent's online_cnt counts itself and its children.
+		 * Once the current css is offlined, a parent with an online_cnt
+		 * of 2 will become a leaf again.
+		 */
+		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)
-- 
2.53.1-Meta