[RFC PATCH 53/57] selftests/mm: cover collapse of mlocked ranges

[Kiryl Shutsemau <kirill@shutemov.name>]

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

Collapsing an mlocked range makes the teardown do something it does
nowhere else: the sources have to be munlocked while the destination
arrives already mlocked, and munlocking takes a reference.  So a teardown
that reaches a source before it is unfrozen fails on a refcount that is
not allowed to move.  That is the one ordering constraint in the putback
with no other way to be caught.

An mlocked range was collapsible before, as long as the whole VMA was
locked.  This case is the other shape.  mlock() over part of a VMA splits
it, leaving the locked part smaller than a PMD, which khugepaged passed
over for as long as its coverage was rooted at PMD-aligned spans.  A
partially mlocked region therefore went uncollapsed however long it lived.

Cover it deterministically: mlock a window, collapse it, check the
contents survive.  Drive it under contention too, with a thread mlocking
and munlocking random spans across the race harness's region, since the
ordering only breaks when a teardown and an mlock overlap.  The plain
racers never touch VM_LOCKED at all.

Assisted-by: Claude-Code:claude-opus-5
Signed-off-by: Kiryl Shutsemau (Meta) <kas@kernel.org>
---
 tools/testing/selftests/mm/khugepaged.c      | 38 ++++++++++++++++++++
 tools/testing/selftests/mm/khugepaged_race.c | 29 +++++++++++++--
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/tools/testing/selftests/mm/khugepaged.c b/tools/testing/selftests/mm/khugepaged.c
index b61e32566d47..208300ecb344 100644
--- a/tools/testing/selftests/mm/khugepaged.c
+++ b/tools/testing/selftests/mm/khugepaged.c
@@ -1599,6 +1599,43 @@ static void collapse_order_sub_pmd_range(struct collapse_context *c,
 	__collapse_order_sub_pmd_vma(c, ops, nr_windows, __func__);
 }
 
+/*
+ * Collapse of an mlocked window: source teardown munlocks the old
+ * pages while the new folio arrives mlocked via folio_add_lru_vma().
+ * A teardown that touches the sources while they are still frozen
+ * blows up exactly here (munlock_folio() takes a reference).
+ */
+static void collapse_order_mlocked(struct collapse_context *c,
+				   struct mem_ops *ops)
+{
+	size_t window = mthp_window_size();
+	void *p;
+
+	mthp_push_target_order();
+
+	p = ops->setup_area(1);
+	ops->fault(p, 0, window);
+	if (mlock(p, window))
+		ksft_exit_fail_perror("mlock()");
+	if (!window_not_collapsed(p, hpage_pmd_size))
+		ksft_exit_fail_msg("Unexpected large folio after fault\n");
+
+	madvise(p, hpage_pmd_size, MADV_HUGEPAGE);
+	ksft_print_msg("Collapse fully populated mlocked window...");
+	if (!khugepaged_wait_full_pass())
+		fail("Timeout");
+	else if (window_collapsed(p, window))
+		success("OK");
+	else
+		fail("Fail");
+
+	validate_memory(p, 0, window);
+	munlock(p, window);
+	ops->cleanup_area(p, hpage_pmd_size);
+	thp_pop_settings();
+	ksft_test_result_report(exit_status, "%s\n", __func__);
+}
+
 /*
  * A partially populated window in a sub-PMD VMA: population and
  * sub-PMD eligibility at once. The unfaulted slots must come back
@@ -1938,6 +1975,7 @@ int main(int argc, char **argv)
 		TEST(collapse_order_sub_pmd_vma, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_sub_pmd_range, mthp_khugepaged_context, anon_ops);
 		TEST(collapse_order_sub_pmd_holes, mthp_khugepaged_context, anon_ops);
+		TEST(collapse_order_mlocked, mthp_khugepaged_context, anon_ops);
 	}
 
 	TEST(collapse_full, madvise_context, anon_ops);
diff --git a/tools/testing/selftests/mm/khugepaged_race.c b/tools/testing/selftests/mm/khugepaged_race.c
index 6682bbae0a8f..a4710130aabf 100644
--- a/tools/testing/selftests/mm/khugepaged_race.c
+++ b/tools/testing/selftests/mm/khugepaged_race.c
@@ -219,6 +219,29 @@ static void *forker_fn(void *arg)
 	return NULL;
 }
 
+/*
+ * mlock/munlock cycling over the shared areas: collapse of an mlocked
+ * range munlocks the sources at teardown and mlocks the new folio --
+ * the interaction the fuzzer caught (munlock on a frozen source) and
+ * the plain racers never drove.
+ */
+static void *mlocker_fn(void *arg)
+{
+	unsigned int seed = (unsigned long)arg;
+
+	while (!stop) {
+		unsigned long page_idx = rand_page(&seed);
+		unsigned long nr = 1UL << (rand_r(&seed) % 8);	/* 1..128 pages */
+
+		if (rand_r(&seed) & 1)
+			mlock(region + page_idx * page_size, nr * page_size);
+		else
+			munlock(region + page_idx * page_size, nr * page_size);
+		usleep(rand_r(&seed) % 1000);
+	}
+	return NULL;
+}
+
 static void *mremapper_fn(void *arg)
 {
 	unsigned int seed = (unsigned long)arg;
@@ -328,14 +351,14 @@ int main(int argc, char **argv)
 {
 	static const char * const thread_names[] = {
 		"faulter", "faulter2", "dontneed", "pinner", "forker",
-		"mremapper", "pageout", "compactor",
+		"mremapper", "mlocker", "pageout", "compactor",
 	};
 	void *(*const thread_fns[])(void *) = {
 		faulter_fn, faulter_fn, dontneed_fn, pinner_fn, forker_fn,
-		mremapper_fn, pageout_fn, compactor_fn,
+		mremapper_fn, mlocker_fn, pageout_fn, compactor_fn,
 	};
 	enum { T_FAULTER, T_FAULTER2, T_DONTNEED, T_PINNER, T_FORKER,
-	       T_MREMAPPER, T_PAGEOUT, T_COMPACTOR };
+	       T_MREMAPPER, T_MLOCKER, T_PAGEOUT, T_COMPACTOR };
 	const unsigned long pageout_bit = 1UL << T_PAGEOUT;
 	const unsigned long compactor_bit = 1UL << T_COMPACTOR;
 	const int nr_threads = ARRAY_SIZE(thread_names);
-- 
2.54.0