Re: [PATCH v5 12/17] mm/huge_memory: move anon_vma handling into the anon split helper

[Yeoreum Yun <yeoreum.yun@arm.com>]

LGTM.

Reviewed-by: Yeoreum Yun <yeoreum.yun@arm.com>

> From: Kairui Song <kasong@tencent.com>
> 
> Only anon split needs the anon_vma, and it only needs it to unmap and
> remap.  Move the folio_get_anon_vma()/anon_vma_lock_write() pair out of
> __folio_split() into the anon helper next to the folio_mapped()
> check that already gates unmap_folio().
> 
> This makes the anon_vma conditional on folio_mapped(), which is a
> behaviour change but should be fine.  folio_get_anon_vma() returns
> NULL whenever !folio_mapped(), so an anon folio with
> folio_mapcount() == 0 used to get -EBUSY from split_huge_page() and is
> now split instead.  A realistic case is a THP that has been fully
> swapped out and is still in the swap cache: swap PTEs do not contribute
> mapcount, so it is !folio_mapped() but still alive.
> 
> That should be safe and right to have because:
> 
> - folio_ref_freeze() below still rejects a folio that picked up any
>   reference, a mapping or a GUP pin, in the meantime.
> 
> - A parallel split is excluded by the folio lock.  The anon_vma
>   write lock was added to serialize split in commit 062f1af2170a
>   ("mm: thp: acquire the anon_vma rwsem for write during split"), when
>   split_huge_page() did not hold the folio lock throughout. commit
>   e9b61f19858a ("thp: reintroduce split_huge_page()") later made the
>   folio lock a caller requirement and added the folio_ref_freeze()
>   scheme, so that has been covered ever since.
> 
> - Unmapped path is already exercised by folio_split_unmapped(),
>   and the swap cache split already runs well for a partially
>   swapped-out mapped THP.
> 
> - A !folio_mapped() folio cannot become mapped meanwhile: mapping it
>   requires the folio lock.
> 
> For mapped folios the anon_vma write lock is now released before
> __folio_split() unlocks the after-split sub-folios, where previously it
> was held across that loop; that window is harmless as the sub-folios
> stay folio-locked and referenced until it.
> 
> Reviewed-by: Zi Yan <ziy@nvidia.com>
> Signed-off-by: Kairui Song <kasong@tencent.com>
> ---
>  mm/huge_memory.c | 58 ++++++++++++++++++++++++++++++--------------------------
>  1 file changed, 31 insertions(+), 27 deletions(-)
> 
> diff --git a/mm/huge_memory.c b/mm/huge_memory.c
> index 8fc1dd5b8354..a4c2e38f4a6d 100644
> --- a/mm/huge_memory.c
> +++ b/mm/huge_memory.c
> @@ -4029,16 +4029,38 @@ static int __folio_freeze_split_anon(struct folio *folio,
>  	struct swap_cluster_info *ci = NULL;
>  	int old_order = folio_order(folio);
>  	struct folio *new_folio, *next;
> +	struct anon_vma *anon_vma = NULL;
>  	enum ttu_flags ttu_flags = 0;
>  	struct lruvec *lruvec;
> -	bool need_remap = false;
>  	int ret = 0;
>  
> +	/*
> +	 * Unmap/remap needs the anon_vma, so we first take a reference on
> +	 * it to prevent it from disappearing, and lock it for write here,
> +	 * letting unmap_folio() walk the rmap with TTU_RMAP_LOCKED.
> +	 *
> +	 * folio_mapped() is not stable here, but it can only change in
> +	 * one direction while the folio is locked. The mapcount can drop
> +	 * to zero at any time, zap_pte_range() takes no folio lock. It
> +	 * cannot go up: swapin, migration and uffd move all lock the folio
> +	 * before mapping it, and fork only copies PTEs that already exist.
> +	 *
> +	 * So if we see the folio mapped, the worst case is an empty rmap
> +	 * walk. If we see it unmapped, it stays unmapped and needs neither
> +	 * the reference nor the lock. Anything else needs a reference
> +	 * first and folio_ref_freeze() below catches it.
> +	 *
> +	 * Note a swapped-out THP counts as unmapped here as swap PTEs do
> +	 * not contribute mapcount, and they are splittable.
> +	 */
>  	if (folio_mapped(folio)) {
> -		need_remap = true;
> +		anon_vma = folio_get_anon_vma(folio);
> +		if (!anon_vma)
> +			return -EBUSY;
> +		anon_vma_lock_write(anon_vma);
>  		ret = unmap_folio(folio);
>  		if (ret)
> -			return ret;
> +			goto out_unlock;
>  	}
>  
>  	local_irq_disable();
> @@ -4096,11 +4118,16 @@ static int __folio_freeze_split_anon(struct folio *folio,
>  		swap_cluster_unlock(ci);
>  out_no_split:
>  	local_irq_enable();
> -	if (need_remap) {
> +	if (anon_vma) {
>  		if (!ret && !folio_is_device_private(folio))
>  			ttu_flags = TTU_USE_SHARED_ZEROPAGE;
>  		remap_anon_folio(folio, 1 << old_order, ttu_flags);
>  	}
> +out_unlock:
> +	if (anon_vma) {
> +		anon_vma_unlock_write(anon_vma);
> +		put_anon_vma(anon_vma);
> +	}
>  
>  	return ret;
>  }
> @@ -4295,7 +4322,6 @@ static int __folio_split(struct folio *folio, unsigned int new_order,
>  	struct folio *end_folio = folio_next(folio);
>  	bool is_anon = folio_test_anon(folio);
>  	struct mem_cgroup *memcg, *old_memcg;
> -	struct anon_vma *anon_vma = NULL;
>  	int old_order = folio_order(folio);
>  	struct folio *new_folio, *next;
>  	int ret;
> @@ -4326,23 +4352,6 @@ static int __folio_split(struct folio *folio, unsigned int new_order,
>  	memcg = get_mem_cgroup_from_folio(folio);
>  	old_memcg = set_active_memcg(memcg);
>  
> -	if (is_anon) {
> -		/*
> -		 * The caller does not necessarily hold an mmap_lock that would
> -		 * prevent the anon_vma disappearing so we first we take a
> -		 * reference to it and then lock the anon_vma for write. This
> -		 * is similar to folio_lock_anon_vma_read except the write lock
> -		 * is taken to serialise against parallel split or collapse
> -		 * operations.
> -		 */
> -		anon_vma = folio_get_anon_vma(folio);
> -		if (!anon_vma) {
> -			ret = -EBUSY;
> -			goto out;
> -		}
> -		anon_vma_lock_write(anon_vma);
> -	}
> -
>  	if (is_anon)
>  		ret = __folio_freeze_split_anon(folio, new_order, split_at,
>  						true, list, split_type);
> @@ -4369,11 +4378,6 @@ static int __folio_split(struct folio *folio, unsigned int new_order,
>  		free_folio_and_swap_cache(new_folio);
>  	}
>  
> -	if (anon_vma) {
> -		anon_vma_unlock_write(anon_vma);
> -		put_anon_vma(anon_vma);
> -	}
> -out:
>  	/* restore to caller's old_memcg */
>  	set_active_memcg(old_memcg);
>  	mem_cgroup_put(memcg);
> 
> -- 
> 2.55.0
> 
> 

-- 
Sincerely,
Yeoreum Yun