[RFC PATCH 13/57] mm/collapse: revalidate a round against the VMA

[Kiryl Shutsemau <kirill@shutemov.name>]

From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>

Fill in the pass that re-establishes what the round is working on.
Selection ran under mmap_lock and the allocation ran without it, so by
the time the round takes the lock back the address space may have changed
underneath it.

Check that the mm is not exiting and has not had THP disabled, that the
VMA the round looked up is still anonymous with an anon_vma, and find the
PTE table again in case it became a huge PMD or went away.

Then re-check each candidate on its own.  A VMA that shrank, or was
replaced by a smaller one, may no longer hold a window that fitted when
it was selected, and per-size enablement may have been turned off for its
order since.  Such a candidate is dropped and the rest of the round goes
on without it.

The check is per candidate rather than over the batch because a window is
aligned to its own order: thp_vma_suitable_order() on each one is the
containment check in full.

What the walk leaves is what the round goes on to freeze, so it also
settles the batch's span, in cc->batch_start and cc->batch_end, for the
one invalidate the round issues.  A candidate the walk dropped is not in
the span, and a round left with no candidates has no span and nothing to
run.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 mm/collapse.c   | 49 ++++++++++++++++++++++++++++++++++++++++++++-----
 mm/collapse.h   | 11 +++++++++++
 mm/khugepaged.c | 11 -----------
 3 files changed, 55 insertions(+), 16 deletions(-)

diff --git a/mm/collapse.c b/mm/collapse.c
index 58c8d83f3468..1367ade721f7 100644
--- a/mm/collapse.c
+++ b/mm/collapse.c
@@ -177,18 +177,36 @@ int collapse_control_init(struct collapse_control *cc)
 
 /*
  * The scan and the allocation both dropped mmap_lock, so nothing seen before it
- * can be trusted: find the VMA and the PTE table again, and check they still
- * allow every provisioned candidate.
+ * can be trusted: check the VMA the round just looked up and the PTE table
+ * again, and that they still allow every provisioned candidate.
  *
- * This is also where the batch's span is settled, for the invalidate the round
- * issues over it.
+ * The VMA was found by address, so it need not be the one the scan saw, nor
+ * still cover everything the round collected -- thp_vma_suitable_order() asks
+ * that of each candidate, since a window is aligned to its own order.  A VMA
+ * that shrank under a candidate therefore refuses that candidate and no more,
+ * like every other pass.
+ *
+ * What survives is what the round goes on to freeze, so this is also where the
+ * batch's span is settled, for the invalidate the round issues over it.
  */
 static enum scan_result collapse_revalidate(struct vm_area_struct *vma,
 					    unsigned long pmd_addr,
 					    struct collapse_control *cc,
 					    pmd_t **pmdp)
 {
-	unsigned int i;
+	struct mm_struct *mm = vma->vm_mm;
+	enum scan_result result;
+	unsigned int i, nr_live = 0;
+
+	if (unlikely(collapse_test_exit_or_disable(mm)))
+		return SCAN_ANY_PROCESS;
+
+	if (!vma->anon_vma || !vma_is_anonymous(vma))
+		return SCAN_PAGE_ANON;
+
+	result = find_pmd_or_thp_or_none(mm, pmd_addr, pmdp);
+	if (result != SCAN_SUCCEED)
+		return result;
 
 	cc->batch_start = ULONG_MAX;
 	cc->batch_end = 0;
@@ -196,10 +214,31 @@ static enum scan_result collapse_revalidate(struct vm_area_struct *vma,
 	for (i = 0; i < cc->nr_candidates; i++) {
 		struct collapse_candidate *cand = &cc->candidates[i];
 
+		if (cand->state != CAND_SELECTED)
+			continue;
+
+		/*
+		 * The window has to still fit the VMA, which may have shrunk or
+		 * been replaced, and its order to still be one the VMA allows.
+		 */
+		if (!thp_vma_suitable_order(vma, cand->addr, cand->order) ||
+		    !thp_vma_allowable_orders(vma, vma->vm_flags,
+					      cc->policy.tva_type,
+					      BIT(cand->order))) {
+			cand->state = CAND_SKIPPED;
+			cand->result = SCAN_VMA_CHECK;
+			continue;
+		}
+
 		cc->batch_start = min(cc->batch_start, candidate_start(cand));
 		cc->batch_end = max(cc->batch_end, candidate_end(cand));
+		nr_live++;
 	}
 
+	/* Nothing the VMA still allows: no span to invalidate, nothing to run */
+	if (!nr_live)
+		return SCAN_VMA_CHECK;
+
 	return SCAN_SUCCEED;
 }
 
diff --git a/mm/collapse.h b/mm/collapse.h
index feb2e0d57339..0d6f77a7233b 100644
--- a/mm/collapse.h
+++ b/mm/collapse.h
@@ -138,6 +138,17 @@ struct collapse_control {
 	unsigned long batch_end;
 };
 
+static inline int collapse_test_exit(struct mm_struct *mm)
+{
+	return atomic_read(&mm->mm_users) == 0;
+}
+
+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);
+}
+
 int collapse_control_init(struct collapse_control *cc);
 void collapse_control_release(struct collapse_control *cc);
 
diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 50b520961b9b..1244e161beae 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -421,17 +421,6 @@ void __init khugepaged_destroy(void)
 	kmem_cache_destroy(mm_slot_cache);
 }
 
-static inline int collapse_test_exit(struct mm_struct *mm)
-{
-	return atomic_read(&mm->mm_users) == 0;
-}
-
-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);
-}
-
 static inline bool anon_hpage_enabled(void)
 {
 	if (READ_ONCE(huge_anon_orders_always))
-- 
2.54.0