*/
 static void flush_memcg_stats_dwork(struct work_struct *w);
 static DECLARE_DEFERRABLE_WORK(stats_flush_dwork, flush_memcg_stats_dwork);
-static atomic_t stats_flush_ongoing = ATOMIC_INIT(0);
 static u64 flush_last_time;
 
 #define FLUSH_TIME (2UL*HZ)
        }
 }
 
-static void do_flush_stats(void)
+static void do_flush_stats(struct mem_cgroup *memcg)
 {
-       /*
-        * We always flush the entire tree, so concurrent flushers can just
-        * skip. This avoids a thundering herd problem on the rstat global lock
-        * from memcg flushers (e.g. reclaim, refault, etc).
-        */
-       if (atomic_read(&stats_flush_ongoing) ||
-           atomic_xchg(&stats_flush_ongoing, 1))
-               return;
-
-       WRITE_ONCE(flush_last_time, jiffies_64);
-
-       cgroup_rstat_flush(root_mem_cgroup->css.cgroup);
+       if (mem_cgroup_is_root(memcg))
+               WRITE_ONCE(flush_last_time, jiffies_64);
 
-       atomic_set(&stats_flush_ongoing, 0);
+       cgroup_rstat_flush(memcg->css.cgroup);
 }
 
-void mem_cgroup_flush_stats(void)
+/*
+ * mem_cgroup_flush_stats - flush the stats of a memory cgroup subtree
+ * @memcg: root of the subtree to flush
+ *
+ * Flushing is serialized by the underlying global rstat lock. There is also a
+ * minimum amount of work to be done even if there are no stat updates to flush.
+ * Hence, we only flush the stats if the updates delta exceeds a threshold. This
+ * avoids unnecessary work and contention on the underlying lock.
+ */
+void mem_cgroup_flush_stats(struct mem_cgroup *memcg)
 {
-       if (memcg_should_flush_stats(root_mem_cgroup))
-               do_flush_stats();
+       if (mem_cgroup_disabled())
+               return;
+
+       if (!memcg)
+               memcg = root_mem_cgroup;
+
+       if (memcg_should_flush_stats(memcg))
+               do_flush_stats(memcg);
 }
 
-void mem_cgroup_flush_stats_ratelimited(void)
+void mem_cgroup_flush_stats_ratelimited(struct mem_cgroup *memcg)
 {
        /* Only flush if the periodic flusher is one full cycle late */
        if (time_after64(jiffies_64, READ_ONCE(flush_last_time) + 2*FLUSH_TIME))
-               mem_cgroup_flush_stats();
+               mem_cgroup_flush_stats(memcg);
 }
 
 static void flush_memcg_stats_dwork(struct work_struct *w)
         * Deliberately ignore memcg_should_flush_stats() here so that flushing
         * in latency-sensitive paths is as cheap as possible.
         */
-       do_flush_stats();
+       do_flush_stats(root_mem_cgroup);
        queue_delayed_work(system_unbound_wq, &stats_flush_dwork, FLUSH_TIME);
 }
 
         *
         * Current memory state:
         */
-       mem_cgroup_flush_stats();
+       mem_cgroup_flush_stats(memcg);
 
        for (i = 0; i < ARRAY_SIZE(memory_stats); i++) {
                u64 size;
        int nid;
        struct mem_cgroup *memcg = mem_cgroup_from_seq(m);
 
-       mem_cgroup_flush_stats();
+       mem_cgroup_flush_stats(memcg);
 
        for (stat = stats; stat < stats + ARRAY_SIZE(stats); stat++) {
                seq_printf(m, "%s=%lu", stat->name,
 
        BUILD_BUG_ON(ARRAY_SIZE(memcg1_stat_names) != ARRAY_SIZE(memcg1_stats));
 
-       mem_cgroup_flush_stats();
+       mem_cgroup_flush_stats(memcg);
 
        for (i = 0; i < ARRAY_SIZE(memcg1_stats); i++) {
                unsigned long nr;
        struct mem_cgroup *memcg = mem_cgroup_from_css(wb->memcg_css);
        struct mem_cgroup *parent;
 
-       mem_cgroup_flush_stats();
+       mem_cgroup_flush_stats(memcg);
 
        *pdirty = memcg_page_state(memcg, NR_FILE_DIRTY);
        *pwriteback = memcg_page_state(memcg, NR_WRITEBACK);
        int i;
        struct mem_cgroup *memcg = mem_cgroup_from_seq(m);
 
-       mem_cgroup_flush_stats();
+       mem_cgroup_flush_stats(memcg);
 
        for (i = 0; i < ARRAY_SIZE(memory_stats); i++) {
                int nid;
                        break;
                }
 
-               cgroup_rstat_flush(memcg->css.cgroup);
+               /*
+                * mem_cgroup_flush_stats() ignores small changes. Use
+                * do_flush_stats() directly to get accurate stats for charging.
+                */
+               do_flush_stats(memcg);
                pages = memcg_page_state(memcg, MEMCG_ZSWAP_B) / PAGE_SIZE;
                if (pages < max)
                        continue;
 static u64 zswap_current_read(struct cgroup_subsys_state *css,
                              struct cftype *cft)
 {
-       cgroup_rstat_flush(css->cgroup);
-       return memcg_page_state(mem_cgroup_from_css(css), MEMCG_ZSWAP_B);
+       struct mem_cgroup *memcg = mem_cgroup_from_css(css);
+
+       mem_cgroup_flush_stats(memcg);
+       return memcg_page_state(memcg, MEMCG_ZSWAP_B);
 }
 
 static int zswap_max_show(struct seq_file *m, void *v)