Re: [PATCH v5 23/38] mm/mlock: eliminate weird VMA_IO_BIT abuse and simplify

["Lorenzo Stoakes (ARM)" <ljs@kernel.org>]

On Mon, Oct 12, 2026 at 08:23:15PM +0200, David Hildenbrand (Arm) wrote:
> On 10/9/26 20:04, Lorenzo Stoakes (ARM) wrote:
> > When performing mlock() or munlock() otherwise normal VMAs have VMA_IO_BIT
> > solely to fix a race with migration which might otherwise double-count
> > mlock VMAs.
> >
> > This is unnecessary - at the point of applying folio mlock state, whether
> > setting or clearing PG_mlocked, we know whether or not we are locking.
> >
> > Solve this in two ways - thread a boolean through the page table walk
> > indicating whether a lock or unlock is being performed, and run a locking
> > walk with VMA_LOCKONFAULT_BIT set and VMA_LOCKED_BIT cleared.
> >
> > This state never occurs otherwise, as VMA_LOCKONFAULT_BIT always implies
> > VMA_LOCKED_BIT. These are also always cleared together.
> >
> > Then, update folio_add_lru_vma() and mlock_folio() to check only for
> > VMA_LOCKED_BIT, and update try_to_unmap_one() to check for VMA_LOCKED_MASK
> > instead.
> >
> > Also remove the useless invocation of allow_mlock_munlock() which simply
> > returns true if unlocking and instead rename it to allow_mlock() and only
> > call it when locking.
> >
> > With the other mlock abuse of VMA_IO_BIT addressed, update
> > mlock_vma_folio() and folio_add_lru_vma() to simply test for
> > VMA_LOCKED_BIT.
> >
> > The munlock_vma_folio() function tests VMA_LOCKED_MASK instead, as an
> > unmap racing with the locking walk must still munlock folios the walk has
> > already counted.
> >
> > Additionally, in order to ensure a racing unmap/migration releases the
> > reference munlock_folio() took (whose put was batched), check
> > VMA_LOCKED_MASK before calling mlock_drain_local() in try_to_unmap_one()
> > and try_to_migrate_one().
> >
> > While here, also replace some deprecated VMA flag predicates.
> >
> > Reviewed-by: Zi Yan <ziy@nvidia.com>
> > Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
> > ---
> >  mm/folio.c    |  2 +-
> >  mm/internal.h | 10 +++++++---
> >  mm/mlock.c    | 51 +++++++++++++++++++--------------------------------
> >  mm/rmap.c     |  8 +++++---
> >  4 files changed, 32 insertions(+), 39 deletions(-)
> >
> > diff --git a/mm/folio.c b/mm/folio.c
> > index 47a437e0f7fd..35e242b48870 100644
> > --- a/mm/folio.c
> > +++ b/mm/folio.c
> > @@ -505,7 +505,7 @@ void folio_add_lru_vma(struct folio *folio, struct vm_area_struct *vma)
> >  {
> >  	VM_BUG_ON_FOLIO(folio_test_lru(folio), folio);
> >
> > -	if (unlikely((vma->vm_flags & (VM_LOCKED | VM_SPECIAL)) == VM_LOCKED))
> > +	if (vma_test(vma, VMA_LOCKED_BIT))
> >  		mlock_new_folio(folio);
> >  	else
> >  		folio_add_lru(folio);
> > diff --git a/mm/internal.h b/mm/internal.h
> > index 4ae08fedffff..ea8fa94ebe1b 100644
> > --- a/mm/internal.h
> > +++ b/mm/internal.h
> > @@ -970,8 +970,7 @@ void mlock_folio(struct folio *folio);
> >  static inline void mlock_vma_folio(struct folio *folio,
> >  				struct vm_area_struct *vma)
> >  {
> > -	/* The VM_IO check prevents migration from double-counting during mlock. */
> > -	if (unlikely((vma->vm_flags & (VM_LOCKED|VM_SPECIAL)) == VM_LOCKED))
> > +	if (vma_test(vma, VMA_LOCKED_BIT))
> >  		mlock_folio(folio);
> >  }
> >
> > @@ -988,7 +987,12 @@ static inline void munlock_vma_folio(struct folio *folio,
> >  	 * always munlock the folio and page reclaim will correct it
> >  	 * if it's wrong.
> >  	 */
> > -	if (unlikely(vma->vm_flags & VM_LOCKED))
> > +	/*
> > +	 * VMA_LOCKONFAULT_BIT alone marks an mlock walk in progress, see
> > +	 * mlock_vma_pages_range(). An unmap racing with the walk must still
> > +	 * munlock folios the walk has already counted.
> > +	 */
> > +	if (unlikely(vma_test_any_mask(vma, VMA_LOCKED_MASK)))
> >  		munlock_folio(folio);
> >  }
> >
> > diff --git a/mm/mlock.c b/mm/mlock.c
> > index 39215a3eab1f..4235a1518fc9 100644
> > --- a/mm/mlock.c
> > +++ b/mm/mlock.c
> > @@ -316,22 +316,10 @@ static inline unsigned int folio_mlock_step(struct folio *folio,
> >  	return folio_pte_batch(folio, pte, ptent, count);
> >  }
> >
> > -static inline bool allow_mlock_munlock(struct folio *folio,
> > +static inline bool allow_mlock(struct folio *folio,
> >  		struct vm_area_struct *vma, unsigned long start,
> >  		unsigned long end, unsigned int step)
> >  {
> > -	/*
> > -	 * For unlock, allow munlock large folio which is partially
> > -	 * mapped to VMA. As it's possible that large folio is
> > -	 * mlocked and VMA is split later.
> > -	 *
> > -	 * During memory pressure, such kind of large folio can
> > -	 * be split. And the pages are not in VM_LOCKed VMA
> > -	 * can be reclaimed.
> > -	 */
> > -	if (!vma_test(vma, VMA_LOCKED_BIT))
> > -		return true;
> > -
> >  	/* folio_within_range() cannot take KSM, but any small folio is OK */
> >  	if (!folio_test_large(folio))
> >  		return true;
> > @@ -352,6 +340,7 @@ static int mlock_pte_range(pmd_t *pmd, unsigned long addr,
> >
> >  {
> >  	struct vm_area_struct *vma = walk->vma;
> > +	const bool lock = walk->private;
> >  	spinlock_t *ptl;
> >  	pte_t *start_pte, *pte;
> >  	pte_t ptent;
> > @@ -368,7 +357,7 @@ static int mlock_pte_range(pmd_t *pmd, unsigned long addr,
> >  		folio = pmd_folio(*pmd);
> >  		if (folio_is_zone_device(folio))
> >  			goto out;
> > -		if (vma_test(vma, VMA_LOCKED_BIT))
> > +		if (lock)
> >  			mlock_folio(folio);
> >  		else
> >  			munlock_folio(folio);
> > @@ -390,10 +379,10 @@ static int mlock_pte_range(pmd_t *pmd, unsigned long addr,
> >  			continue;
> >
> >  		step = folio_mlock_step(folio, pte, addr, end);
> > -		if (!allow_mlock_munlock(folio, vma, start, end, step))
> > +		if (lock && !allow_mlock(folio, vma, start, end, step))
> >  			goto next_entry;
> >
> > -		if (vma_test(vma, VMA_LOCKED_BIT))
> > +		if (lock)
> >  			mlock_folio(folio);
> >  		else
> >  			munlock_folio(folio);
> > @@ -428,31 +417,29 @@ static void mlock_vma_pages_range(struct vm_area_struct *vma,
> >  		.pmd_entry = mlock_pte_range,
> >  		.walk_lock = PGWALK_WRLOCK_VERIFY,
> >  	};
> > +	const bool lock = vma_flags_test(new_vma_flags, VMA_LOCKED_BIT);
> > +	vma_flags_t walk_flags = *new_vma_flags;
> >
> >  	/*
> > -	 * There is a slight chance that concurrent page migration,
> > -	 * or page reclaim finding a page of this now-VMA_LOCKED_BIT vma,
> > -	 * will call mlock_vma_folio() and raise page's mlock_count:
> > -	 * double counting, leaving the page unevictable indefinitely.
> > -	 * Communicate this danger to mlock_vma_folio() with VMA_IO_BIT,
> > -	 * which is a VMA_SPECIAL_FLAGS flag not allowed on VMA_LOCKED_BIT vmas.
> > -	 * mmap_lock is held in write mode here, so this weird
> > -	 * combination should not be visible to other mmap_lock users;
> > -	 * but WRITE_ONCE so rmap walkers must see VMA_IO_BIT if VMA_LOCKED_BIT.
> > +	 * LOCKONFAULT without LOCKED never otherwise occurs: it marks a walk in
> > +	 * progress so that rmap-side callers, which test VMA_LOCKED_BIT, do not
> > +	 * count folios, while try_to_unmap_one(), which tests VMA_LOCKED_MASK,
> > +	 * still refuses to unmap them.
> >  	 */
> > -	if (vma_flags_test(new_vma_flags, VMA_LOCKED_BIT))
> > -		vma_flags_set(new_vma_flags, VMA_IO_BIT);
> > +	if (lock) {
> > +		vma_flags_clear(&walk_flags, VMA_LOCKED_BIT);
> > +		vma_flags_set(&walk_flags, VMA_LOCKONFAULT_BIT);
> > +	}
> > +
> >  	vma_start_write(vma);
> > -	vma_flags_reset_once(vma, new_vma_flags);
> > +	vma_flags_reset_once(vma, &walk_flags);
> >
> >  	lru_add_drain();
> > -	walk_page_range_vma(vma, start, end, &mlock_walk_ops, NULL);
> > +	walk_page_range_vma(vma, start, end, &mlock_walk_ops, (void *)lock);
> >  	lru_add_drain();
> >
> > -	if (vma_flags_test(new_vma_flags, VMA_IO_BIT)) {
> > -		vma_flags_clear(new_vma_flags, VMA_IO_BIT);
> > +	if (lock)
> >  		vma_flags_reset_once(vma, new_vma_flags);
> > -	}
> >  }
> >
> >  /*
> > diff --git a/mm/rmap.c b/mm/rmap.c
> > index 5332c52909be..b4a4b3929044 100644
> > --- a/mm/rmap.c
> > +++ b/mm/rmap.c
> > @@ -2239,9 +2239,11 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
> >
> >  		/*
> >  		 * If the folio is in an mlock()d vma, we must not swap it out.
> > +		 * VMA_LOCKONFAULT_BIT alone marks an mlock walk in progress, see
> > +		 * mlock_vma_pages_range().
> >  		 */
> >  		if (!(flags & TTU_IGNORE_MLOCK) &&
> > -		    (vma->vm_flags & VM_LOCKED)) {
> > +		    vma_test_any_mask(vma, VMA_LOCKED_MASK)) {
> >  			ptes++;
> >
> >  			/*
> > @@ -2395,7 +2397,7 @@ static bool try_to_unmap_one(struct folio *folio, struct vm_area_struct *vma,
> >  		}
> >  finish_unmap:
> >  		folio_remove_rmap_ptes(folio, page, nr_pages, vma);
> > -		if (vma->vm_flags & VM_LOCKED)
> > +		if (vma_test_any_mask(vma, VMA_LOCKED_MASK))
> >  			mlock_drain_local();
> >  		folio_put_refs(folio, nr_pages);
> >
> > @@ -2770,7 +2772,7 @@ static bool try_to_migrate_one(struct folio *folio, struct vm_area_struct *vma,
> >  			hugetlb_remove_rmap(folio);
> >  		else
> >  			folio_remove_rmap_pte(folio, subpage, vma);
> > -		if (vma->vm_flags & VM_LOCKED)
> > +		if (vma_test_any_mask(vma, VMA_LOCKED_MASK))
> >  			mlock_drain_local();
> >  		folio_put(folio);
> >  	}
> >
>
> The mlock.c changes are a bit over my head (mostly the old VM_IO handling ...),
> but it looks good from what I can tell (hoping Hugh could review ...)

If it helps, Lance tested thoroughly and found 1 place I forgot to update, now
addressed, and Zi has carefully gone through too.

Essentially it's doing the same thing VMA_IO_BIT did only with
VMA_LOCKONFAULT_BIT, which is still pretty gross but much much less so and makes
it all make a lot more sense (esp. with the locked mask).

(VMA_LOCKONFAULT_BIT when used for lock-on-fault is only ever set with
VMA_LOCKED_BIT set too, we track the folio operation by it being set alone so it
doesn't trigger lock-on-fault behaviour).

>
> Acked-by: David Hildenbrand (Arm) <david@kernel.org>

Thanks!

>
> --
> Cheers,
>
> David

--
Cheers, Lorenzo