[PATCH] mm/migrate: fix mlock state loss when migrating PTE-mapped large folios

[PATCH] mm/migrate: fix mlock state loss when migrating PTE-mapped large folios

[Wandun Chen]

From: Wandun Chen <chenwandun@lixiang.com>

try_to_migrate() clears PG_mlocked when unmapping the source folio.
For PTE-mapped large folios, remove_migration_pte() restores one PTE at
a time, but the rmap helpers only restore mlock for full-folio mappings.
This leaves the destination PTE-mapped large folio without PG_mlocked.

Restore mlock once all PTEs of the destination folio are restored under
a single page table lock.

Fixes: dc68badcede4 ("mm: mlock: update mlock_pte_range to handle large folio")
Signed-off-by: Wandun Chen <chenwandun@lixiang.com>
---
 mm/migrate.c | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/mm/migrate.c b/mm/migrate.c
index 94033c222e04..140e291d4587 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -353,6 +353,7 @@ static bool remove_migration_pte(struct folio *folio,
 {
 	struct rmap_walk_arg *rmap_walk_arg = arg;
 	DEFINE_FOLIO_VMA_WALK(pvmw, rmap_walk_arg->folio, vma, addr, PVMW_SYNC | PVMW_MIGRATION);
+	unsigned long nr_ptes = 0;
 
 	while (page_vma_mapped_walk(&pvmw)) {
 		rmap_t rmap_flags = RMAP_NONE;
@@ -446,9 +447,16 @@ static bool remove_migration_pte(struct folio *folio,
 			else
 				folio_add_file_rmap_pte(folio, new, vma);
 			set_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
+
+			nr_ptes++;
 		}
-		if (READ_ONCE(vma->vm_flags) & VM_LOCKED)
+		if (READ_ONCE(vma->vm_flags) & VM_LOCKED) {
+			if (folio_test_large(folio) &&
+			    nr_ptes == folio_nr_pages(folio) &&
+			    !(pvmw.flags & PVMW_PGTABLE_CROSSED))
+				mlock_vma_folio(folio, vma);
 			mlock_drain_local();
+		}
 
 		trace_remove_migration_pte(pvmw.address, pte_val(pte),
 					   compound_order(new));
-- 
2.43.0

Re: [PATCH] mm/migrate: fix mlock state loss when migrating PTE-mapped large folios

[David Hildenbrand (Arm)]

On 11/9/26 10:08, Wandun Chen wrote:
> From: Wandun Chen <chenwandun@lixiang.com>
> 
> try_to_migrate() clears PG_mlocked when unmapping the source folio.
> For PTE-mapped large folios, remove_migration_pte() restores one PTE at
> a time, but the rmap helpers only restore mlock for full-folio mappings.
> This leaves the destination PTE-mapped large folio without PG_mlocked.
> 
> Restore mlock once all PTEs of the destination folio are restored under
> a single page table lock.
> 
> Fixes: dc68badcede4 ("mm: mlock: update mlock_pte_range to handle large folio")
> Signed-off-by: Wandun Chen <chenwandun@lixiang.com>
> ---
>  mm/migrate.c | 10 +++++++++-
>  1 file changed, 9 insertions(+), 1 deletion(-)
> 
> diff --git a/mm/migrate.c b/mm/migrate.c
> index 94033c222e04..140e291d4587 100644
> --- a/mm/migrate.c
> +++ b/mm/migrate.c
> @@ -353,6 +353,7 @@ static bool remove_migration_pte(struct folio *folio,
>  {
>  	struct rmap_walk_arg *rmap_walk_arg = arg;
>  	DEFINE_FOLIO_VMA_WALK(pvmw, rmap_walk_arg->folio, vma, addr, PVMW_SYNC | PVMW_MIGRATION);
> +	unsigned long nr_ptes = 0;
>  
>  	while (page_vma_mapped_walk(&pvmw)) {
>  		rmap_t rmap_flags = RMAP_NONE;
> @@ -446,9 +447,16 @@ static bool remove_migration_pte(struct folio *folio,
>  			else
>  				folio_add_file_rmap_pte(folio, new, vma);
>  			set_pte_at(vma->vm_mm, pvmw.address, pvmw.pte, pte);
> +
> +			nr_ptes++;
>  		}
> -		if (READ_ONCE(vma->vm_flags) & VM_LOCKED)
> +		if (READ_ONCE(vma->vm_flags) & VM_LOCKED) {
> +			if (folio_test_large(folio) &&
> +			    nr_ptes == folio_nr_pages(folio) &&
> +			    !(pvmw.flags & PVMW_PGTABLE_CROSSED))
> +				mlock_vma_folio(folio, vma);
>  			mlock_drain_local();
> +		}

I have the recollection that the next try_to_unmap() would fix that up? Or is
that not the case?

The "nr_ptes == folio_nr_pages(folio)" case is shaky either way for folios that
are mapped into multiple VMAs, or what's the logic behind that, such that is is
fine to only focus on large folios that getting fully pte-mapped?

-- 
Cheers,

David