[PATCH v2] mm/page_alloc: fix shuffling and reporting for no-lock frees

[Karl Mehltretter <kmehltretter@gmail.com>]

free_pages_nolock() may run from NMI context. The BPF arena also calls
it while holding arena->spinlock through raw_res_spin_lock_irqsave().
When the zone trylock succeeds, __free_one_page() still reaches page
shuffling and page reporting.

When shuffle_pick_tail() exhausts its random-bit cache, it calls
get_random_u64(). Once the CRNG is ready, this takes an rtmutex-backed
local lock on PREEMPT_RT, which cannot be acquired while a raw spinlock
is held. A test-only forced refill under a raw spinlock reported an
invalid wait context at batched_entropy_u64.lock.

page_reporting_notify_free() may call queue_delayed_work(), which is not
NMI-safe.

Pass the no-lock constraint to shuffle_pick_tail(). For no-lock frees,
use a pre-seeded per-CPU prandom state, following the slab allocator.
Keep the cached get_random_u64() path for ordinary frees. Seed all
possible CPUs before memblock memory is released to the buddy allocator.

Skip page reporting for no-lock frees. The page remains on the buddy
list for a later reporting pass.

For stable backports, FPI_NOLOCK was named FPI_TRYLOCK through v7.2.

Fixes: 8c57b687e833 ("mm, bpf: Introduce free_pages_nolock()")
Cc: stable@vger.kernel.org # needs FPI_TRYLOCK through v7.2
Suggested-by: Vlastimil Babka (SUSE) <vbabka@kernel.org>
Assisted-by: LLM
Signed-off-by: Karl Mehltretter <kmehltretter@gmail.com>
---
Changes in v2:
- Preserve shuffling for no-lock frees with a pre-seeded per-CPU prandom
  state, following the slab allocator.
- Keep the cached get_random_u64() path for ordinary frees.
- Keep the page-reporting guard and explain its NMI constraint.
- Check is_shuffle_order() before the no-lock flag.

v1: https://lore.kernel.org/r/20261005063515.6312-1-kmehltretter@gmail.com/

Testing:

A temporary caller freed 65 maximum-order pages under a raw spinlock,
then repeated the workload through __free_pages(). SMP=1 and SMP=4 QEMU
TCG guests produced the same path counts:

                                  no-lock  ordinary
  shuffle decisions                    65        65
  prandom_u32_state() decisions        65         0
  cached random-bit decisions           0        65
  get_random_u64() calls                0         2
  reporting notifications               0        65

Both guests ran PREEMPT_RT with lockdep, raw-lock nesting checks, page
reporting, and page shuffling enabled. Both lanes selected the list head
and tail. Neither guest reported an invalid wait context, lockdep splat,
atomic-sleep warning, BUG, oops, or panic.

The test does not inject an NMI. The in-tree BPF arena caller establishes
the raw-lock path, while the forced refill checks the previously failing
transition.


 mm/mm_init.c    |  1 +
 mm/page_alloc.c |  6 +++---
 mm/shuffle.c    | 24 +++++++++++++++++++++++-
 mm/shuffle.h    |  9 +++++++--
 4 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/mm/mm_init.c b/mm/mm_init.c
index 1533aebafb688..01bccb23a89c2 100644
--- a/mm/mm_init.c
+++ b/mm/mm_init.c
@@ -2643,6 +2643,7 @@ void __init mm_core_init_early(void)
 void __init mm_core_init(void)
 {
 	arch_mm_preinit();
+	shuffle_init();
 	init_zero_page_pfn();
 
 	/* Initializations relying on SMP setup */
diff --git a/mm/page_alloc.c b/mm/page_alloc.c
index 12fac9084c483..f13993bda67e1 100644
--- a/mm/page_alloc.c
+++ b/mm/page_alloc.c
@@ -1014,14 +1014,14 @@ static inline void __free_one_page(struct page *page,
 	if (fpi_flags & FPI_TO_TAIL)
 		to_tail = true;
 	else if (is_shuffle_order(order))
-		to_tail = shuffle_pick_tail();
+		to_tail = shuffle_pick_tail(!(fpi_flags & FPI_NOLOCK));
 	else
 		to_tail = buddy_merge_likely(pfn, buddy_pfn, page, order);
 
 	__add_to_free_list(page, zone, order, migratetype, to_tail);
 
-	/* Notify page reporting subsystem of freed page */
-	if (!(fpi_flags & FPI_SKIP_REPORT_NOTIFY))
+	/* No-lock frees may run from NMI, where delayed work cannot be queued. */
+	if (!(fpi_flags & (FPI_SKIP_REPORT_NOTIFY | FPI_NOLOCK)))
 		page_reporting_notify_free(order);
 }
 
diff --git a/mm/shuffle.c b/mm/shuffle.c
index 82a2c7725a08a..2e1f76d6bf92b 100644
--- a/mm/shuffle.c
+++ b/mm/shuffle.c
@@ -5,6 +5,7 @@
 #include <linux/init.h>
 #include <linux/mmzone.h>
 #include <linux/random.h>
+#include <linux/prandom.h>
 #include <linux/moduleparam.h>
 #include "internal.h"
 #include "page_alloc.h"
@@ -13,6 +14,7 @@
 DEFINE_STATIC_KEY_FALSE(page_alloc_shuffle_key);
 
 static bool shuffle_param;
+static DEFINE_PER_CPU(struct rnd_state, shuffle_rnd_state);
 
 static __meminit int shuffle_param_set(const char *val,
 		const struct kernel_param *kp)
@@ -30,6 +32,12 @@ static const struct kernel_param_ops shuffle_param_ops = {
 };
 module_param_cb(shuffle, &shuffle_param_ops, &shuffle_param, 0400);
 
+void __init shuffle_init(void)
+{
+	if (static_branch_unlikely(&page_alloc_shuffle_key))
+		prandom_seed_full_state(&shuffle_rnd_state);
+}
+
 /*
  * For two pages to be swapped in the shuffle, they must be free (on a
  * 'free_area' lru), have the same order, and have the same migratetype.
@@ -159,12 +167,26 @@ void __meminit __shuffle_free_memory(pg_data_t *pgdat)
 		shuffle_zone(z);
 }
 
-bool shuffle_pick_tail(void)
+bool shuffle_pick_tail(bool allow_spin)
 {
 	static u64 rand;
 	static u8 rand_bits;
 	bool ret;
 
+	if (!allow_spin) {
+		struct rnd_state *state;
+
+		/*
+		 * An interrupt or NMI handler might interrupt and change
+		 * the state in the middle, but that's safe.
+		 */
+		state = &get_cpu_var(shuffle_rnd_state);
+		ret = prandom_u32_state(state) & 1;
+		put_cpu_var(shuffle_rnd_state);
+
+		return ret;
+	}
+
 	/*
 	 * The lack of locking is deliberate. If 2 threads race to
 	 * update the rand state it just adds to the entropy.
diff --git a/mm/shuffle.h b/mm/shuffle.h
index 11bec7521ab89..92f5e758388d9 100644
--- a/mm/shuffle.h
+++ b/mm/shuffle.h
@@ -11,7 +11,8 @@
 #ifdef CONFIG_SHUFFLE_PAGE_ALLOCATOR
 DECLARE_STATIC_KEY_FALSE(page_alloc_shuffle_key);
 extern void __shuffle_free_memory(pg_data_t *pgdat);
-extern bool shuffle_pick_tail(void);
+void __init shuffle_init(void);
+bool shuffle_pick_tail(bool allow_spin);
 static inline void __meminit shuffle_free_memory(pg_data_t *pgdat)
 {
 	if (!static_branch_unlikely(&page_alloc_shuffle_key))
@@ -34,7 +35,11 @@ static inline bool is_shuffle_order(int order)
 	return order >= SHUFFLE_ORDER;
 }
 #else
-static inline bool shuffle_pick_tail(void)
+static inline void __init shuffle_init(void)
+{
+}
+
+static inline bool shuffle_pick_tail(bool allow_spin)
 {
 	return false;
 }
-- 
2.39.5 (Apple Git-154)