[PATCH v5 01/12] mm/khugepaged: deposit a newly allocated page table on collapse

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

collapse_huge_page() deposits a PTE page table on PMD collapse in order
that it can be utilised for subsequent split operations, meaning that those
operations do not need to perform an allocation (as they are in a context
where it might be unwise).

However the PTE page table which is deposited is the one which is currently
mapped by the PMD entry that is in the process of being collapsed.

Once deposited, the PTE page table may be used in a split of any other
unrelated PMD entry.

This is currently not an issue as this operation is performed with VMA/mmap
write lock + anon rmap locks held, so ordinary page table walkers will
never accidentally end up walking the wrong thing, and GUP-fast is
protected by an IPI via tlb_remove_table_sync_one().

However, the series to which this commit belongs implements RCU-safe page
table traversal, at which point this becomes problematic.

This can be resolved by using pte_offset_map_lock() which gates on a PTE
PTL and a pmd_same() check, but lockless walks are unsafe as things stand.

Resolve this by simply allocating a new, zeroed, PTE page table to deposit
at the point of collapse.

This path is already costly and an allocation has already been performed
for the huge folio, so this allocation is statistical noise in terms of
performance and memory usage at this point.

With this PTE page table deposited, RCU-free the existing PTE page table
so it is safe for page table walkers to traverse within a grace period.

This also brings this deposit case in line with all other page table
deposit logic which deposit a fresh page table.

Additionally, this was the only place in the kernel that displaced a page
table like this, so eliminating it also helps consistency.

Since khugepaged runs as a kernel thread, do a little dance in
alloc_deposit_pte_table() to correctly charge the allocation.

This is already done for the folio allocation via alloc_charge_folio() but
no such wrapper exists for a page table allocation.

Reviewed-by: Baolin Wang <baolin.wang@linux.alibaba.com>
Acked-by: David Hildenbrand (Arm) <david@kernel.org>
Reviewed-by: Lance Yang <lance.yang@linux.dev>
Signed-off-by: Lorenzo Stoakes (ARM) <ljs@kernel.org>
---
 mm/khugepaged.c | 34 ++++++++++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index f49a6710933b..639029cc4f5e 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -1278,6 +1278,23 @@ static enum scan_result alloc_charge_folio(struct folio **foliop, struct mm_stru
 	return SCAN_SUCCEED;
 }
 
+static pgtable_t alloc_deposit_pte_table(struct mm_struct *mm)
+{
+	/*
+	 * khugepaged is run from a kernel thread, so need to manually set the
+	 * correct memcg so the allocation gets charged correctly.
+	 */
+	struct mem_cgroup *memcg = get_mem_cgroup_from_mm(mm);
+	struct mem_cgroup *old_memcg = set_active_memcg(memcg);
+	pgtable_t pgtable;
+
+	pgtable = pte_alloc_one(mm);
+
+	set_active_memcg(old_memcg);
+	mem_cgroup_put(memcg);
+	return pgtable;
+}
+
 /*
  * collapse_huge_page() expects the mmap_lock to be unlocked before entering and
  * will always return with the lock unlocked, to avoid holding the mmap_lock
@@ -1293,7 +1310,7 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long s
 	LIST_HEAD(compound_pagelist);
 	pmd_t *pmd, _pmd;
 	pte_t *pte = NULL;
-	pgtable_t pgtable;
+	pgtable_t pgtable = NULL;
 	struct folio *folio;
 	spinlock_t *pmd_ptl, *pte_ptl;
 	enum scan_result result = SCAN_FAIL;
@@ -1310,6 +1327,14 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long s
 		goto out_nolock;
 	}
 
+	if (is_pmd_order(order)) {
+		pgtable = alloc_deposit_pte_table(mm);
+		if (!pgtable) {
+			result = SCAN_ALLOC_HUGE_PAGE_FAIL;
+			goto out_nolock;
+		}
+	}
+
 	mmap_read_lock(mm);
 	result = hugepage_vma_revalidate(mm, pmd_addr, /*expect_anon=*/ true,
 					 &vma, cc, order);
@@ -1433,8 +1458,8 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long s
 	spin_lock(pmd_ptl);
 	VM_WARN_ON_ONCE(!pmd_none(*pmd));
 	if (is_pmd_order(order)) {
-		pgtable = pmd_pgtable(_pmd);
 		pgtable_trans_huge_deposit(mm, pmd, pgtable);
+		pgtable = NULL;
 		map_anon_folio_pmd_nopf(folio, pmd, vma, pmd_addr);
 	} else {
 		/*
@@ -1453,6 +1478,9 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long s
 	}
 	spin_unlock(pmd_ptl);
 
+	if (is_pmd_order(order))
+		pte_free_defer(mm, pmd_pgtable(_pmd));
+
 	folio = NULL;
 
 	result = SCAN_SUCCEED;
@@ -1463,6 +1491,8 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long s
 		anon_vma_unlock_write(vma->anon_vma);
 	mmap_write_unlock(mm);
 out_nolock:
+	if (pgtable)
+		pte_free(mm, pgtable);
 	if (folio)
 		folio_put(folio);
 	trace_mm_collapse_huge_page(mm, result == SCAN_SUCCEED, result, order);

-- 
2.55.0