mm: check all swaps belong to same swap_cgroup in swap_pte_batch()

Right now, it is possible two folios are contiguous in swap slots but they
don't belong to one memcg.  In this case, even we return a large nr, we
can't really batch free all slots.

Link: https://lkml.kernel.org/r/20240815215308.55233-1-21cnbao@gmail.com
Signed-off-by: Barry Song <v-songbaohua@oppo.com>
Reported-by: Yosry Ahmed <yosryahmed@google.com>
Reported-by: Chris Li <chrisl@kernel.org>
Signed-off-by: Andrew Morton <akpm@linux-foundation.org>

diff --git a/mm/internal.h b/mm/internal.h
index 1159b04e76a301f9c17c455736d3d43b36fbb039..26c25d18466db446ecd80bec87af841a8954786b 100644 (file)
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -15,6 +15,7 @@
 #include <linux/rmap.h>
 #include <linux/swap.h>
 #include <linux/swapops.h>
+#include <linux/swap_cgroup.h>
 #include <linux/tracepoint-defs.h>
 
 /* Internal core VMA manipulation functions. */
@@ -275,18 +276,22 @@ static inline int swap_pte_batch(pte_t *start_ptep, int max_nr, pte_t pte)
 {
        pte_t expected_pte = pte_next_swp_offset(pte);
        const pte_t *end_ptep = start_ptep + max_nr;
+       swp_entry_t entry = pte_to_swp_entry(pte);
        pte_t *ptep = start_ptep + 1;
+       unsigned short cgroup_id;
 
        VM_WARN_ON(max_nr < 1);
        VM_WARN_ON(!is_swap_pte(pte));
-       VM_WARN_ON(non_swap_entry(pte_to_swp_entry(pte)));
+       VM_WARN_ON(non_swap_entry(entry));
 
+       cgroup_id = lookup_swap_cgroup_id(entry);
        while (ptep < end_ptep) {
                pte = ptep_get(ptep);
 
                if (!pte_same(pte, expected_pte))
                        break;
-
+               if (lookup_swap_cgroup_id(pte_to_swp_entry(pte)) != cgroup_id)
+                       break;
                expected_pte = pte_next_swp_offset(expected_pte);
                ptep++;
        }