[PATCH v1] mm/rmap: fix migration entries being treated as partial unmaps

[Joanne Koong <joannelkoong@gmail.com>]

When try_to_migrate_one() unmaps a large anonymous folio that is mapped by
PTEs, it replaces the PTEs with migration entries one at a time and drops
the rmap for each page with folio_remove_rmap_pte(). After the first PTE
is replaced, __folio_remove_rmap() sees a folio that is no longer fully
mapped but still has pages mapped, and calls
deferred_split_folio(folio, true). That marks the folio
PG_partially_mapped, bumps split_deferred and nr_anon_partially_mapped,
and puts it on the deferred split queue, even though nothing was actually
unmapped. The migration entries are turned back into mappings once the
split or migration is done.

This happens whenever a fully mapped, PTE-mapped anon large folio is
migrated or split. Its effects:

 - Migrated folios can later be split for no reason. migrate_folio_move()
   copies the source folio's deferred split state to the destination, but
   reads it after the unmap, when PG_partially_mapped has already been
   wrongly set. So a fully mapped mTHP ends up queued as partially mapped
   after migration, and the shrinker can later split it without checking
   whether it is underused.

 - split_deferred is overcounted. Every such split increments it, so the
   stat counts folios that were never partially mapped.

 - A folio whose split fails can later be split for no reason. If the
   split fails at folio_ref_freeze(), the folio is remapped but stays on
   the deferred split queue marked partially mapped. The deferred split
   shrinker then splits it without checking whether it is underused,
   breaking up a folio that is fully mapped and possibly fully in use.

Mappings replaced by migration entries are not really unmapped, so don't
queue the folio for deferred split when try_to_migrate_one() installs a
migration entry. Pass a flag down to __folio_remove_rmap() and have
try_to_migrate_one() clear it for those PTEs only. PTEs that it replaces
with a hwpoison entry or discards because they are unused are really
unmapped and keep the current behaviour, as do all other callers.

One consequence is that if part of the folio is really unmapped while the
migration entries are in place, zapping the migration entries does not go
through the rmap, so the folio is not marked partially mapped afterwards.
For a split this does not matter, as the folio is split anyway. For a
migrated folio it means the destination is not queued until a later unmap
notices it, which seems preferable to queueing every migrated folio.

Fixes: 8422acdc97ed ("mm: introduce a pageflag for partially mapped folios")
Fixes: a2e0c0668a34 ("mm: migrate: requeue destination folio on deferred split queue")
Reported-by: Sashiko <sashiko-bot@kernel.org>
Closes: https://lore.kernel.org/all/sashiko-outbox-198104@kernel.org/
Signed-off-by: Joanne Koong <joannelkoong@gmail.com>
---
This bug was found by Sashiko in its review of the mTHP shrinker series
(https://lore.kernel.org/linux-mm/sashiko-outbox-198104@kernel.org/). The
migration / rmap code is still new to me. This patch is a fix with the help of
AI, so I'd appreciate a careful look from people who know this code better.

This was tested in an x86 QEMU guest with 64K mTHP. Splitting 64 fully mapped
folios through the split_huge_pages debugfs file no longer increases
split_deferred (was +64), and migrating 64 to another node no longer leaves
nr_anon_partially_mapped up (was +64). split_huge_page_test, migration, and
khugepaged selftests pass.

 mm/rmap.c | 29 +++++++++++++++++++++++------
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/mm/rmap.c b/mm/rmap.c
index fa34e1ccbd8f..e5d17481c0c0 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -1799,7 +1799,7 @@ void folio_add_file_rmap_pud(struct folio *folio, struct page *page,
 
 static __always_inline void __folio_remove_rmap(struct folio *folio,
 		struct page *page, int nr_pages, struct vm_area_struct *vma,
-		enum pgtable_level level)
+		enum pgtable_level level, bool may_defer_split)
 {
 	atomic_t *mapped = &folio->_nr_pages_mapped;
 	int last = 0, nr = 0, nr_pmdmapped = 0;
@@ -1888,8 +1888,13 @@ static __always_inline void __folio_remove_rmap(struct folio *folio,
 	 *
 	 * Device private folios do not support deferred splitting and
 	 * shrinker based scanning of the folios to free.
+	 *
+	 * Skip this if the caller is only replacing the mappings with
+	 * migration entries. The folio is not really becoming partially
+	 * mapped. Once the split or migration finishes, every migration
+	 * entry becomes a mapping again.
 	 */
-	if (partially_mapped && folio_test_anon(folio) &&
+	if (may_defer_split && partially_mapped && folio_test_anon(folio) &&
 	    !folio_test_partially_mapped(folio) &&
 	    !folio_is_device_private(folio))
 		deferred_split_folio(folio, true);
@@ -1921,7 +1926,7 @@ static __always_inline void __folio_remove_rmap(struct folio *folio,
 void folio_remove_rmap_ptes(struct folio *folio, struct page *page,
 		int nr_pages, struct vm_area_struct *vma)
 {
-	__folio_remove_rmap(folio, page, nr_pages, vma, PGTABLE_LEVEL_PTE);
+	__folio_remove_rmap(folio, page, nr_pages, vma, PGTABLE_LEVEL_PTE, true);
 }
 
 /**
@@ -1938,7 +1943,8 @@ void folio_remove_rmap_pmd(struct folio *folio, struct page *page,
 		struct vm_area_struct *vma)
 {
 #ifdef CONFIG_TRANSPARENT_HUGEPAGE
-	__folio_remove_rmap(folio, page, HPAGE_PMD_NR, vma, PGTABLE_LEVEL_PMD);
+	__folio_remove_rmap(folio, page, HPAGE_PMD_NR, vma, PGTABLE_LEVEL_PMD,
+			    true);
 #else
 	WARN_ON_ONCE(true);
 #endif
@@ -1959,7 +1965,8 @@ void folio_remove_rmap_pud(struct folio *folio, struct page *page,
 {
 #if defined(CONFIG_TRANSPARENT_HUGEPAGE) && \
 	defined(CONFIG_HAVE_ARCH_TRANSPARENT_HUGEPAGE_PUD)
-	__folio_remove_rmap(folio, page, HPAGE_PUD_NR, vma, PGTABLE_LEVEL_PUD);
+	__folio_remove_rmap(folio, page, HPAGE_PUD_NR, vma, PGTABLE_LEVEL_PUD,
+			    true);
 #else
 	WARN_ON_ONCE(true);
 #endif
@@ -2516,6 +2523,8 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 	mmu_notifier_invalidate_range_start(&range);
 
 	while (page_vma_mapped_walk(&pvmw)) {
+		bool may_defer_split = true;
+
 		/* PMD-mapped THP migration entry */
 		if (!pvmw.pte) {
 			__maybe_unused unsigned long pfn;
@@ -2772,12 +2781,20 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
 			 * No need to invalidate here it will synchronize on
 			 * against the special swap migration pte.
 			 */
+
+			/*
+			 * The folio only looks partially unmapped while this PTE
+			 * holds a migration entry, which is later turned back
+			 * into a mapping. Don't queue it for deferred split.
+			 */
+			may_defer_split = false;
 		}
 
 		if (unlikely(folio_test_hugetlb(folio)))
 			hugetlb_remove_rmap(folio);
 		else
-			folio_remove_rmap_pte(folio, subpage, vma);
+			__folio_remove_rmap(folio, subpage, 1, vma,
+					    PGTABLE_LEVEL_PTE, may_defer_split);
 		if (vma_test_any_mask(vma, VMA_LOCKED_MASK))
 			mlock_drain_local();
 		folio_put(folio);

base-commit: b95f8bb86713c860ce8536d22a988b949b884d5f
-- 
2.52.0