From: "Kiryl Shutsemau (Meta)" <[email protected]>

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) <[email protected]>
---
 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


Reply via email to