mm, swap: delay and unify memcg lookup and charging for swapin

Instead of checking the cgroup private ID during page table walk in
swap_pte_batch(), move the memcg lookup into __swap_cache_add_check()
under the cluster lock.

The first pre-alloc check is speculative and skips the memcg check since
the post-alloc stable check ensures all slots covered by the folio belong
to the same memcg.  It is very rare for contiguous and aligned entries
across a contiguous region of a page table of the same process or shmem
mapping to belong to different memcgs.

This also prepares for recording the memcg info in the cluster's table.
Also make the order check and fallback more compact.

There should be no user-observable behavior change.

Link: https://lore.kernel.org/20260517-swap-table-p4-v5-8-88ae43e064c7@tencent.com
Signed-off-by: Kairui Song <kasong@tencent.com>
Acked-by: Chris Li <chrisl@kernel.org>
Cc: Baolin Wang <baolin.wang@linux.alibaba.com>
Cc: Baoquan He <bhe@redhat.com>
Cc: Barry Song <baohua@kernel.org>
Cc: Chengming Zhou <chengming.zhou@linux.dev>
Cc: David Hildenbrand <david@kernel.org>
Cc: Hugh Dickins <hughd@google.com>
Cc: Johannes Weiner <hannes@cmpxchg.org>
Cc: Kemeng Shi <shikemeng@huaweicloud.com>
Cc: Lorenzo Stoakes <ljs@kernel.org>
Cc: Muchun Song <muchun.song@linux.dev>
Cc: Nhat Pham <nphamcs@gmail.com>
Cc: Roman Gushchin <roman.gushchin@linux.dev>
Cc: Shakeel Butt <shakeel.butt@linux.dev>
Cc: Youngjun Park <youngjun.park@lge.com>
Cc: Zi Yan <ziy@nvidia.com>
Signed-off-by: Andrew Morton <akpm@linux-foundation.org>

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 7d08128de1fdc44e507075ca675380a370fb56e2..a013f37f24aa05ac704764ef3f3d00d17ff7162a 100644 (file)
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -646,8 +646,8 @@ static inline int mem_cgroup_charge(struct folio *folio, struct mm_struct *mm,
 
 int mem_cgroup_charge_hugetlb(struct folio* folio, gfp_t gfp);
 
-int mem_cgroup_swapin_charge_folio(struct folio *folio, struct mm_struct *mm,
-                                 gfp_t gfp, swp_entry_t entry);
+int mem_cgroup_swapin_charge_folio(struct folio *folio, unsigned short id,
+                                  struct mm_struct *mm, gfp_t gfp);
 
 void __mem_cgroup_uncharge(struct folio *folio);
 
@@ -1137,7 +1137,7 @@ static inline int mem_cgroup_charge_hugetlb(struct folio* folio, gfp_t gfp)
 }
 
 static inline int mem_cgroup_swapin_charge_folio(struct folio *folio,
-                       struct mm_struct *mm, gfp_t gfp, swp_entry_t entry)
+                unsigned short id, struct mm_struct *mm, gfp_t gfp)
 {
        return 0;
 }

diff --git a/mm/internal.h b/mm/internal.h
index 09931b1e535f3f71887b5b6473f93ed21a41c7e7..9dbd8e3c991f3ee15b2be40bcd7229daff80581c 100644 (file)
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -451,24 +451,16 @@ 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;
-       const softleaf_t entry = softleaf_from_pte(pte);
        pte_t *ptep = start_ptep + 1;
-       unsigned short cgroup_id;
 
        VM_WARN_ON(max_nr < 1);
-       VM_WARN_ON(!softleaf_is_swap(entry));
+       VM_WARN_ON(!softleaf_is_swap(softleaf_from_pte(pte)));
 
-       cgroup_id = lookup_swap_cgroup_id(entry);
        while (ptep < end_ptep) {
-               softleaf_t entry;
-
                pte = ptep_get(ptep);
 
                if (!pte_same(pte, expected_pte))
                        break;
-               entry = softleaf_from_pte(pte);
-               if (lookup_swap_cgroup_id(entry) != cgroup_id)
-                       break;
                expected_pte = pte_next_swp_offset(expected_pte);
                ptep++;
        }

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index c3d0f79dc84e1cb979ce5c60bf3eb2df974f5320..1b58b314cb186884d0a8937a6f126e74c207d717 100644 (file)
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -5079,27 +5079,25 @@ out:
 
 /**
  * mem_cgroup_swapin_charge_folio - Charge a newly allocated folio for swapin.
- * @folio: folio to charge.
+ * @folio: the folio to charge
+ * @id: memory cgroup id
  * @mm: mm context of the victim
  * @gfp: reclaim mode
- * @entry: swap entry for which the folio is allocated
  *
  * This function charges a folio allocated for swapin. Please call this before
  * adding the folio to the swapcache.
  *
  * Returns 0 on success. Otherwise, an error code is returned.
  */
-int mem_cgroup_swapin_charge_folio(struct folio *folio, struct mm_struct *mm,
-                                 gfp_t gfp, swp_entry_t entry)
+int mem_cgroup_swapin_charge_folio(struct folio *folio, unsigned short id,
+                                  struct mm_struct *mm, gfp_t gfp)
 {
        struct mem_cgroup *memcg;
-       unsigned short id;
        int ret;
 
        if (mem_cgroup_disabled())
                return 0;
 
-       id = lookup_swap_cgroup_id(entry);
        rcu_read_lock();
        memcg = mem_cgroup_from_private_id(id);
        if (!memcg || !css_tryget_online(&memcg->css))

diff --git a/mm/swap_state.c b/mm/swap_state.c
index 7a80494fa37fbd439040c97689ee8a23492b8a74..bdd949ae004485a0e0a95df062db3a0c8a0bdb71 100644 (file)
--- a/mm/swap_state.c
+++ b/mm/swap_state.c
@@ -142,17 +142,21 @@ void *swap_cache_get_shadow(swp_entry_t entry)
  * @ci: The locked swap cluster
  * @targ_entry: The target swap entry to check, will be rounded down by @nr
  * @nr: Number of slots to check, must be a power of 2
- * @shadowp: Returns the shadow value if one exists in the range.
+ * @shadowp: Returns the shadow value if one exists in the range
+ * @memcg_id: Returns the memory cgroup id, NULL to ignore cgroup check
  *
  * Check if all slots covered by given range have a swap count >= 1.
- * Retrieves the shadow if there is one.
+ * Retrieves the shadow if there is one. If @memcg_id is not NULL, also
+ * checks if all slots belong to the same cgroup and return the cgroup
+ * private id.
  *
  * Context: Caller must lock the cluster.
  * Return: 0 if success, error code if failed.
  */
 static int __swap_cache_add_check(struct swap_cluster_info *ci,
                                  swp_entry_t targ_entry,
-                                 unsigned long nr, void **shadowp)
+                                 unsigned long nr, void **shadowp,
+                                 unsigned short *memcg_id)
 {
        unsigned int ci_off, ci_end;
        unsigned long old_tb;
@@ -172,19 +176,24 @@ static int __swap_cache_add_check(struct swap_cluster_info *ci,
                return -EEXIST;
        if (!__swp_tb_get_count(old_tb))
                return -ENOENT;
-       if (swp_tb_is_shadow(old_tb) && shadowp)
+       if (shadowp && swp_tb_is_shadow(old_tb))
                *shadowp = swp_tb_to_shadow(old_tb);
+       if (memcg_id)
+               *memcg_id = lookup_swap_cgroup_id(targ_entry);
 
        if (nr == 1)
                return 0;
 
+       targ_entry.val = round_down(targ_entry.val, nr);
        ci_off = round_down(ci_off, nr);
        ci_end = ci_off + nr;
        do {
                old_tb = __swap_table_get(ci, ci_off);
                if (unlikely(swp_tb_is_folio(old_tb) ||
-                            !__swp_tb_get_count(old_tb)))
+                            !__swp_tb_get_count(old_tb) ||
+                            (memcg_id && *memcg_id != lookup_swap_cgroup_id(targ_entry))))
                        return -EBUSY;
+               targ_entry.val++;
        } while (++ci_off < ci_end);
 
        return 0;
@@ -400,6 +409,7 @@ static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
        swp_entry_t entry;
        struct folio *folio;
        void *shadow = NULL;
+       unsigned short memcg_id;
        unsigned long address, nr_pages = 1UL << order;
        struct vm_area_struct *vma = vmf ? vmf->vma : NULL;
 
@@ -408,7 +418,7 @@ static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
 
        /* Check if the slot and range are available, skip allocation if not */
        spin_lock(&ci->lock);
-       err = __swap_cache_add_check(ci, targ_entry, nr_pages, NULL);
+       err = __swap_cache_add_check(ci, targ_entry, nr_pages, NULL, NULL);
        spin_unlock(&ci->lock);
        if (unlikely(err))
                return ERR_PTR(err);
@@ -431,7 +441,7 @@ static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
 
        /* Double check the range is still not in conflict */
        spin_lock(&ci->lock);
-       err = __swap_cache_add_check(ci, targ_entry, nr_pages, &shadow);
+       err = __swap_cache_add_check(ci, targ_entry, nr_pages, &shadow, &memcg_id);
        if (unlikely(err)) {
                spin_unlock(&ci->lock);
                folio_put(folio);
@@ -443,8 +453,8 @@ static struct folio *__swap_cache_alloc(struct swap_cluster_info *ci,
        __swap_cache_do_add_folio(ci, folio, entry);
        spin_unlock(&ci->lock);
 
-       if (mem_cgroup_swapin_charge_folio(folio, vmf ? vmf->vma->vm_mm : NULL,
-                                          gfp, entry)) {
+       if (mem_cgroup_swapin_charge_folio(folio, memcg_id,
+                                          vmf ? vmf->vma->vm_mm : NULL, gfp)) {
                spin_lock(&ci->lock);
                __swap_cache_do_del_folio(ci, folio, entry, shadow);
                spin_unlock(&ci->lock);