Re: [RFC PATCH v2] mm/thp: serialize huge-zero folio state transitions

[David Hildenbrand (Arm)]

On 7/27/26 15:18, Hengbin Zhang wrote:
> The nonpersistent huge-zero shrinker clears huge_zero_folio and then
> invalidates huge_zero_pfn. A concurrent fault can publish a replacement
> folio and PFN between those updates, after which the old shrinker
> invalidates the replacement generation's PFN identity.
> 
> A later partial mprotect() can then misidentify the live special PMD and
> enter the ordinary anonymous THP split path.
> 
> A writer-side spinlock keeps huge_zero_folio and huge_zero_pfn updates
> from different generations together. Lockless getters still use the
> refcount fast path. atomic_set_release() publishes an initialized new
> generation before a successful atomic_inc_not_zero() can admit a getter
> on weakly ordered architectures.
> 
> The race was reproduced using test-only instrumentation that widens the
> shrinker/allocation window; that instrumentation is not included here.
> 
> Suggested-by: David Hildenbrand <david@xxxxxxxxxx>
> Link: https://lore.kernel.org/r/20260724100509.2300200-1-uqbarz@xxxxxxxxx
> Signed-off-by: Hengbin Zhang <uqbarz@xxxxxxxxx>
> ---
>  mm/huge_memory.c | 36 +++++++++++++++++++++++-------------
>  1 file changed, 23 insertions(+), 13 deletions(-)
> 
> diff --git a/mm/huge_memory.c b/mm/huge_memory.c
> index b5d1e9d4463d..89771a9b282d 100644
> --- a/mm/huge_memory.c
> +++ b/mm/huge_memory.c
> @@ -78,6 +78,7 @@ static unsigned long deferred_split_scan(struct shrinker *shrink,
>  static bool split_underused_thp = true;
>  
>  static atomic_t huge_zero_refcount;
> +static DEFINE_SPINLOCK(huge_zero_lock);
>  struct folio *huge_zero_folio __read_mostly;
>  unsigned long huge_zero_pfn __read_mostly = ~0UL;
>  unsigned long huge_anon_orders_always __read_mostly;
> @@ -237,17 +238,19 @@ static bool get_huge_zero_folio(void)
>  	}
>  	/* Ensure zero folio won't have large_rmappable flag set. */
>  	folio_clear_large_rmappable(zero_folio);
> -	preempt_disable();
> -	if (cmpxchg(&huge_zero_folio, NULL, zero_folio)) {
> -		preempt_enable();
> +	spin_lock(&huge_zero_lock);
> +	if (READ_ONCE(huge_zero_folio)) {
> +		spin_unlock(&huge_zero_lock);
>  		folio_put(zero_folio);
>  		goto retry;
>  	}
>  	WRITE_ONCE(huge_zero_pfn, folio_pfn(zero_folio));
> +	WRITE_ONCE(huge_zero_folio, zero_folio);
> +
> +	/* Publish the identity before admitting lockless getters. */
> +	atomic_set_release(&huge_zero_refcount, 2);
> +	spin_unlock(&huge_zero_lock);

Can't we just keep the

atomic_set(&huge_zero_refcount, 2);

after the spin_unlock() ?

>  
> -	/* We take additional reference here. It will be put back by shrinker */
> -	atomic_set(&huge_zero_refcount, 2);
> -	preempt_enable();
>  	count_vm_event(THP_ZERO_PAGE_ALLOC);
>  	return true;
>  }
> @@ -297,15 +300,22 @@ static unsigned long shrink_huge_zero_folio_count(struct shrinker *shrink,
>  static unsigned long shrink_huge_zero_folio_scan(struct shrinker *shrink,
>  						 struct shrink_control *sc)
>  {
> -	if (atomic_cmpxchg(&huge_zero_refcount, 1, 0) == 1) {
> -		struct folio *zero_folio = xchg(&huge_zero_folio, NULL);
> -		BUG_ON(zero_folio == NULL);
> -		WRITE_ONCE(huge_zero_pfn, ~0UL);
> -		folio_put(zero_folio);
> -		return HPAGE_PMD_NR;
> +	struct folio *zero_folio;
> +
> +	spin_lock(&huge_zero_lock);
> +	if (atomic_cmpxchg(&huge_zero_refcount, 1, 0) != 1) {
> +		spin_unlock(&huge_zero_lock);
> +		return 0;
>  	}

Do we really have to move that refcount modifications under the lock?

I'd assume it's sufficient to only move setting huge_zero_pfn+huge_zero_folio
under the spinlock.

-- 
Cheers,

David