[PATCH v3] mm/bpf: Always flush memcg stats when explicitly requested

[Runli <mingyu.he@shopee.com>]

memcg_rstat_updated() stops walking the ancestors when a memcg exceeds
the flush threshold. A concurrent flush can leave an ancestor below the
threshold, so the early exit prevents it from receiving further updates.

Consider root and leaf on a 4-CPU system, with a threshold of 256.
CPU 0 is flushing another CPU's stats after CPU 1's stats have already
been flushed. The counts below are shared stats_updates counters:

  CPU 0 (root flush)                CPU 1 (leaf updates)
  ------------------               --------------------
  clear leaf->stats_updates
                                   add 64 to leaf and root
  clear root->stats_updates
  finish flush
                                   leaf = 64, root = 0
                                   4 more batches of 64:
                                     leaf = 320, root = 256
                                   next update:
                                     leaf > threshold
                                     break; root is not updated

Root remains at the threshold, which is insufficient to trigger a flush,
while further leaf updates keep taking the early exit. The periodic
forced flush restores progress, but until then the root's aggregated
statistics can fall far behind. For example, page cache grows at rate of
1 GiB/s, a two-second wait can leave the reported usage about 2 GiB below
the actual usage.

Make bpf_mem_cgroup_flush_stats() force a flush without waiting for the
periodic flusher, while preserving the memcg flush bookkeeping.

Suggested-by: Shakeel Butt <shakeel.butt@linux.dev>
Link: https://lore.kernel.org/all/avKy9oMwDRjKjfzm@linux.dev/
Signed-off-by: Runli <mingyu.he@shopee.com>
---
Changes since v2:
- Add mem_cgroup_flush_stats_force() and use it in the BPF kfunc to
  preserve the memcg flush bookkeeping.

v2: https://lore.kernel.org/all/20261111071434.34547-1-mingyu.he@shopee.com/

Changes since v1:
- Make the BPF kfunc flush unconditionally, as suggested by Shakeel.
- Leave memcg_rstat_updated() unchanged.

v1: https://lore.kernel.org/all/20260911032905.63683-1-mingyu.he@shopee.com/

 include/linux/memcontrol.h |  5 +++++
 mm/bpf_memcontrol.c        |  5 +++--
 mm/memcontrol.c            | 17 +++++++++++++++++
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 69e1bcc67c58..e4131a368d8c 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -988,6 +988,7 @@ void mod_memcg_lruvec_state(struct lruvec *lruvec,
 			    enum node_stat_item idx, int val);
 
 void mem_cgroup_flush_stats(struct mem_cgroup *memcg);
+void mem_cgroup_flush_stats_force(struct mem_cgroup *memcg);
 void mem_cgroup_flush_stats_ratelimited(struct mem_cgroup *memcg);
 
 void mod_lruvec_kmem_state(void *p, enum node_stat_item idx, int val);
@@ -1464,6 +1465,10 @@ static inline void mem_cgroup_flush_stats(struct mem_cgroup *memcg)
 {
 }
 
+static inline void mem_cgroup_flush_stats_force(struct mem_cgroup *memcg)
+{
+}
+
 static inline void mem_cgroup_flush_stats_ratelimited(struct mem_cgroup *memcg)
 {
 }
diff --git a/mm/bpf_memcontrol.c b/mm/bpf_memcontrol.c
index d8f579c28560..f04be9479e38 100644
--- a/mm/bpf_memcontrol.c
+++ b/mm/bpf_memcontrol.c
@@ -154,11 +154,12 @@ __bpf_kfunc unsigned long bpf_mem_cgroup_page_state(struct mem_cgroup *memcg, in
  * bpf_mem_cgroup_flush_stats - Flush memory cgroup's statistics
  * @memcg: memory cgroup
  *
- * Propagate memory cgroup's statistics up the cgroup tree.
+ * Flush the memory cgroup subtree and propagate its statistics up the cgroup
+ * tree, regardless of the number of pending updates.
  */
 __bpf_kfunc void bpf_mem_cgroup_flush_stats(struct mem_cgroup *memcg)
 {
-	mem_cgroup_flush_stats(memcg);
+	mem_cgroup_flush_stats_force(memcg);
 }
 
 /**
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 84f4a8fd4d16..0057e0ab5a45 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -824,6 +824,23 @@ void mem_cgroup_flush_stats(struct mem_cgroup *memcg)
 	__mem_cgroup_flush_stats(memcg, false);
 }
 
+/**
+ * mem_cgroup_flush_stats_force - unconditionally flush memory cgroup statistics
+ * @memcg: root of the subtree to flush, or NULL for the root memory cgroup
+ *
+ * Flush the statistics regardless of the number of pending updates.
+ */
+void mem_cgroup_flush_stats_force(struct mem_cgroup *memcg)
+{
+	if (mem_cgroup_disabled())
+		return;
+
+	if (!memcg)
+		memcg = root_mem_cgroup;
+
+	__mem_cgroup_flush_stats(memcg, true);
+}
+
 void mem_cgroup_flush_stats_ratelimited(struct mem_cgroup *memcg)
 {
 	/* Only flush if the periodic flusher is one full cycle late */
-- 
2.39.2 (Apple Git-143)