[RFC PATCH 44/57] mm/khugepaged: hold the address space open across a scan

[Kiryl Shutsemau]

From: "Kiryl Shutsemau (Meta)" <kas@xxxxxxxxxx>

Preparation for taking a per-VMA read lock instead of mmap_lock.

What tells khugepaged an address space is going away is the barrier in
__khugepaged_exit(): it runs before exit_mmap() and takes mmap_lock for
writing, which waits for a scan holding it for reading.  A scan under a
per-VMA lock holds no mmap_lock, so nothing waits for it and exit_mmap()
frees the page tables it is walking.

Take a reference on mm_users for the pass instead.  __mmput() cannot start
while one is held, so neither can exit_mmap(), whatever lock the pass
uses.

Drop it with mmput_async(), so the last reference does not tear an address
space down inside khugepaged.  Drop it before the exiting mm is judged,
too: that judgement needs the true count to release the slot.

The reference is also what the exiting-mm checks were reading, so an
address space whose owner has gone now shows as one reference rather than
none.  The three checks inside the pass ask collapse_test_exit_mmref()
instead; the slot-release judgement keeps the old test, running after the
reference is dropped.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@xxxxxxxxxx>
---
 mm/collapse.h   | 19 +++++++++++++++++--
 mm/khugepaged.c | 31 +++++++++++++++++++++++++++----
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/mm/collapse.h b/mm/collapse.h
index 9e2cec1f250b..74e513c5c76c 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -170,6 +170,11 @@ struct collapse_control {
 	pte_t *saved_ptes;
 };
 
+static inline int collapse_disabled(struct mm_struct *mm)
+{
+	return mm_flags_test(MMF_DISABLE_THP_COMPLETELY, mm);
+}
+
 static inline int collapse_test_exit(struct mm_struct *mm)
 {
 	return atomic_read(&mm->mm_users) == 0;
@@ -177,8 +182,18 @@ static inline int collapse_test_exit(struct mm_struct *mm)
 
 static inline int collapse_test_exit_or_disable(struct mm_struct *mm)
 {
-	return collapse_test_exit(mm) ||
-		mm_flags_test(MMF_DISABLE_THP_COMPLETELY, mm);
+	return collapse_test_exit(mm) || collapse_disabled(mm);
+}
+
+/* The owner has gone: the caller's own reference is the only one left */
+static inline int collapse_test_exit_mmref(struct mm_struct *mm)
+{
+	return atomic_read(&mm->mm_users) == 1;
+}
+
+static inline int collapse_test_exit_or_disable_mmref(struct mm_struct *mm)
+{
+	return collapse_test_exit_mmref(mm) || collapse_disabled(mm);
 }
 
 /*
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index cc5ff429d811..f3ea1846990e 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -531,16 +531,31 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 	spin_unlock(&khugepaged_mm_lock);
 
 	mm = slot->mm;
+	vma = NULL;
+
+	/*
+	 * A reference on mm_users for as long as the pass works on this address
+	 * space.  __mmput() cannot start while one is held, so neither can
+	 * exit_mmap(), and the VMAs and page tables stay where they are.
+	 *
+	 * Once per pass, not once per table: the reference is what makes the
+	 * address space safe to work on, and a pass is how long that is wanted
+	 * for.  Nothing else in mm takes it per unit of work -- DAMON takes one
+	 * per target and walks every region under it, swapoff one per mm across
+	 * the whole address space, userfaultfd one per call.
+	 */
+	if (!mmget_not_zero(mm))
+		goto breakouterloop_no_mmput;
+
 	/*
 	 * Don't wait for semaphore (to avoid long wait times).  Just move to
 	 * the next mm on the list.
 	 */
-	vma = NULL;
 	if (unlikely(!mmap_read_trylock(mm)))
 		goto breakouterloop_mmap_lock;
 
 	cc->progress++;
-	if (unlikely(collapse_test_exit_or_disable(mm)))
+	if (unlikely(collapse_test_exit_or_disable_mmref(mm)))
 		goto breakouterloop;
 
 	vma_iter_init(&vmi, mm, khugepaged_scan.address);
@@ -549,7 +564,7 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 		unsigned long orders;
 
 		cond_resched();
-		if (unlikely(collapse_test_exit_or_disable(mm))) {
+		if (unlikely(collapse_test_exit_or_disable_mmref(mm))) {
 			cc->progress++;
 			break;
 		}
@@ -595,7 +610,7 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 			range_end = min(hend, pmd_addr + HPAGE_PMD_SIZE);
 
 			cond_resched();
-			if (unlikely(collapse_test_exit_or_disable(mm)) ||
+			if (unlikely(collapse_test_exit_or_disable_mmref(mm)) ||
 			    cc->progress >= progress_max)
 				goto breakouterloop;
 
@@ -622,6 +637,14 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 breakouterloop:
 	mmap_read_unlock(mm); /* exit_mmap will destroy ptes after this */
 breakouterloop_mmap_lock:
+	/*
+	 * Not mmput(): the last reference would run exit_mmap() here, and
+	 * khugepaged is not the thread that should tear an address space down.
+	 * Dropped before the exiting mm is judged below, so that judgement still
+	 * sees the true count.
+	 */
+	mmput_async(mm);
+breakouterloop_no_mmput:
 
 	spin_lock(&khugepaged_mm_lock);
 	VM_BUG_ON(khugepaged_scan.mm_slot != slot);
-- 
2.54.0