[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

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

[Wandun]

On 11/9/26 21:51, David Hildenbrand (Arm) wrote:
> 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?

It is the case, try_to_unmap() can restore it. However, until reclaim scans
the folio, Mlocked in /proc/meminfo can remain undercounted, giving users
a misleading impression that less memory is locked.

A similar fix was made in commit
	ab521b4142aa ("mm/rmap: improve mlock tracking for large folios").

and handles cases where an entire folio is added in one rmap operation.

However, when a PTE-mapped large folio is migrated, remove_migration_pte()
restores PTEs individually, passing nr_pages = 1 to the rmap helpers,
as a result, PG_mlocked is not set on the destination folio.

This patch handles that case once the full mapping has been restored,
rather than waiting for try_to_unmap() to restore the mlock state.

> 
> 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?

The logic follows the mlock restoration in try_to_unmap_one(), IIUC
only mlock a large folio when it is fully mapped in a VM_LOCKED VMA.
nr_ptes is counted per VMA. A complete mapping in any one locked VMA
is sufficient, regardless of other mappings.

Best regards,
Wandun

>