[PATCH 0/7] Harden preferred CPUs management

[PATCH 0/7] Harden preferred CPUs management

[Yury Norov]

On top of tip sched/core.

This series hardens preferred CPU mask management against CPU hotplug and
prevents competing governors from overriding each other’s selections.

This is partially a resend of [1], excluding the write access, and
partially the new material.

The write access to the preferred cpumask is a subject of the ongoing
discussion. I'll submit it as a follow up shortly.

[1] https://lore.kernel.org/all/20261101192203.6979-1-ynorov@nvidia.com/

Yury Norov (7):
  cpumask: add reset_preferred_cpumask() helper
  sched: restore preferred CPUs when hotplug removes the last one
  virt: steal_governor: remove the minimum elapsed-time check
  virt: steal_governor: serialize sampling against CPU hotplug
  virt: steal_governor: warn on unexpected preferred mask states
  virt: steal_governor: restart sampling immediately after CPU hotplug
  cpu: make preferred CPU policy writers mutually exclusive

 Documentation/scheduler/sched-paravirt.rst |  2 +-
 drivers/virt/steal_governor.c              | 82 +++++++++++++++-------
 include/linux/cpumask.h                    | 30 ++++++++
 kernel/sched/core.c                        |  8 +++
 lib/cpumask.c                              | 67 ++++++++++++++++++
 5 files changed, 164 insertions(+), 25 deletions(-)

-- 
2.53.0

[PATCH 1/7] cpumask: add reset_preferred_cpumask() helper

[Yury Norov]

This series adds more flexibility to control the 'preferred' status of
CPUs. As a preparation, move the steal governor's preferred mask reset
into cpumask.h to make it globally available.

While there, rename it to reset_preferred_cpumask().

Signed-off-by: Yury Norov <ynorov@nvidia.com>
---
 drivers/virt/steal_governor.c | 16 +++++-----------
 include/linux/cpumask.h       | 13 +++++++++++++
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/drivers/virt/steal_governor.c b/drivers/virt/steal_governor.c
index 81e5089ba824..99bebd3dee8c 100644
--- a/drivers/virt/steal_governor.c
+++ b/drivers/virt/steal_governor.c
@@ -50,15 +50,6 @@ static struct steal_governor sg_ctx = {
 	.low_threshold	=	200,	/* 2% */
 };
 
-static void restore_preferred_to_active(void)
-{
-	int cpu;
-
-	guard(cpus_read_lock)();
-	for_each_cpu(cpu, cpu_active_mask)
-		set_cpu_preferred(cpu, true);
-}
-
 static int param_set_interval_ms(const char *val, const struct kernel_param *kp)
 {
 	unsigned int interval;
@@ -252,7 +243,9 @@ static void steal_governor_loop(struct work_struct *work)
 	 * module restores to active if CPU hotplug breaks those assumptions.
 	 */
 	if (!preferred_cpus_valid()) {
-		restore_preferred_to_active();
+		guard(cpus_read_lock)();
+
+		reset_preferred_cpumask();
 		return;
 	}
 
@@ -288,7 +281,8 @@ static int __init steal_governor_init(void)
 static void __exit steal_governor_exit(void)
 {
 	disable_delayed_work_sync(&sg_ctx.work);
-	restore_preferred_to_active();
+	guard(cpus_read_lock)();
+	reset_preferred_cpumask();
 	pr_info("disabled\n");
 }
 
diff --git a/include/linux/cpumask.h b/include/linux/cpumask.h
index bf89bb3f30f6..b13d099474bc 100644
--- a/include/linux/cpumask.h
+++ b/include/linux/cpumask.h
@@ -1195,6 +1195,19 @@ void init_cpu_possible(const struct cpumask *src);
 #define set_cpu_preferred(cpu, preferred) do { } while (0)
 #endif
 
+/**
+ * reset_preferred_cpumask - reset preferred CPUs to the active mask.
+ *
+ * Callers must serialize against CPU hotplug and other preferred cpumask
+ * writers.
+ */
+static inline void reset_preferred_cpumask(void)
+{
+#ifdef CONFIG_PREFERRED_CPU
+	cpumask_copy(&__cpu_preferred_mask, cpu_active_mask);
+#endif
+}
+
 void set_cpu_online(unsigned int cpu, bool online);
 void set_cpu_possible(unsigned int cpu, bool possible);
 
-- 
2.53.0

[PATCH 2/7] sched: restore preferred CPUs when hotplug removes the last one

[Yury Norov]

Offlining the last preferred CPU leaves cpu_preferred_mask empty. Normal
fair load balancing stops, potentially leaving CPUs idle while having
queued tasks.

Currently, steal governor is the only module that may modify the mask,
so it takes care of it by resetting the mask to mirror the active cpumask.

The design allows alternative preferred cpumask managers; relying on
loadable modules to maintain consistency of the core data structure
is not worthwhile in general.

Restore all remaining active CPUs as preferred in sched_cpu_deactivate()
path, without relying on the steal governor to recover.

Signed-off-by: Yury Norov <ynorov@nvidia.com>
---
 kernel/sched/core.c | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 175b52a7ea79..b5c158c4efad 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -8723,6 +8723,14 @@ int sched_cpu_deactivate(unsigned int cpu)
 
 	set_cpu_active(cpu, false);
 
+#ifdef CONFIG_PREFERRED_CPU
+	if (cpumask_empty(cpu_preferred_mask)) {
+		pr_notice("sched: last preferred CPU%u offlined; make all active CPUs preferred\n",
+			  cpu);
+		reset_preferred_cpumask();
+	}
+#endif
+
 	/*
 	 * From this point forward, this CPU will refuse to run any task that
 	 * is not: migrate_disable() or KTHREAD_IS_PER_CPU, and will actively
-- 
2.53.0

[PATCH 3/7] virt: steal_governor: remove the minimum elapsed-time check

[Yury Norov]

The governor queues its worker with a sampling interval of at least
100ms, so its one-millisecond elapsed-time check is redundant. The
existing denominator clamp already prevents division by zero when
calculating the steal ratio.

This series adds a hook on hotplug events. It may trigger this error
path pointlessly. Just drop it.

Signed-off-by: Yury Norov <ynorov@nvidia.com>
---
 drivers/virt/steal_governor.c | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/drivers/virt/steal_governor.c b/drivers/virt/steal_governor.c
index 99bebd3dee8c..26b6bcca27ed 100644
--- a/drivers/virt/steal_governor.c
+++ b/drivers/virt/steal_governor.c
@@ -216,11 +216,6 @@ static void steal_governor_loop(struct work_struct *work)
 	now = ktime_get();
 	delta_ns = ktime_to_ns(ktime_sub(now, sg_ctx.time));
 
-	if (unlikely(delta_ns < NSEC_PER_MSEC)) {
-		pr_err_ratelimited("work scheduled too soon delta_ns: %llu\n", delta_ns);
-		goto requeue_work;
-	}
-
 	curr_steal = get_system_steal_time();
 	delta_steal = curr_steal > sg_ctx.steal ? curr_steal - sg_ctx.steal : 0;
 	sg_ctx.steal = curr_steal;
@@ -249,7 +244,6 @@ static void steal_governor_loop(struct work_struct *work)
 		return;
 	}
 
-requeue_work:
 	schedule_delayed_work(&sg_ctx.work, sg_ctx.delay);
 }
 
-- 
2.53.0

[PATCH 4/7] virt: steal_governor: serialize sampling against CPU hotplug

[Yury Norov]

Steal governor is not safe against a concurrent CPU hotplug operation
because it can change active and preferred cpumasks during the governor
sampling.

To address the race, the governor would take the cpus_write_lock(). But
that may be impractical for the loadable module.

Instead, implement the following locking policy:

 - if the code takes cpus_write_lock(), it's enough to write-access
   the preferred cpumask;
 - if the code takes cpus_read_lock(), it must also take
   cpu_preferred_mutex in order to modify the preferred mask.

The core scheduler and hotplug paths already modify the preferred cpumask
under the cpus_write_lock(), so no change in core logic is required.

This provides a common update protocol for multiple policy drivers
without requiring writers to take the broader CPU hotplug write lock.

Fixes: 27d47ebce4d6 ("virt/steal_governor: Implement steal_governor policy loop")
Signed-off-by: Yury Norov <ynorov@nvidia.com>
---
 drivers/virt/steal_governor.c | 17 +++++++++++++----
 include/linux/cpumask.h       | 11 ++++++++++-
 lib/cpumask.c                 |  5 +++++
 3 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/drivers/virt/steal_governor.c b/drivers/virt/steal_governor.c
index 26b6bcca27ed..8cd5c0f3e3ef 100644
--- a/drivers/virt/steal_governor.c
+++ b/drivers/virt/steal_governor.c
@@ -23,6 +23,7 @@
 #include <linux/ktime.h>
 #include <linux/math64.h>
 #include <linux/module.h>
+#include <linux/mutex.h>
 #include <linux/sched/isolation.h>
 #include <linux/topology.h>
 #include <linux/types.h>
@@ -139,7 +140,8 @@ static void decrease_preferred_cpus(void)
 	int target_cpu = nr_cpu_ids;
 	int cpu;
 
-	guard(cpus_read_lock)();
+	lockdep_assert_cpus_held();
+	lockdep_assert_held(&cpu_preferred_mutex);
 	cpu = cpumask_first_and(housekeeping_cpumask(HK_TYPE_KERNEL_NOISE),
 				cpu_preferred_mask);
 	if (cpu >= nr_cpu_ids)
@@ -179,7 +181,8 @@ static void increase_preferred_cpus(void)
 {
 	int first_cpu, cpu;
 
-	guard(cpus_read_lock)();
+	lockdep_assert_cpus_held();
+	lockdep_assert_held(&cpu_preferred_mutex);
 	first_cpu = cpumask_first_andnot(cpu_active_mask, cpu_preferred_mask);
 
 	/* All CPUs are preferred. Nothing to increase further */
@@ -193,6 +196,9 @@ static void increase_preferred_cpus(void)
 
 static bool preferred_cpus_valid(void)
 {
+	lockdep_assert_cpus_held();
+	lockdep_assert_held(&cpu_preferred_mutex);
+
 	if (cpumask_empty(cpu_preferred_mask)) {
 		pr_err("empty preferred mask. stopping\n");
 		return false;
@@ -213,6 +219,9 @@ static void steal_governor_loop(struct work_struct *work)
 	u64 curr_steal, delta_steal, delta_ns, steal_ratio;
 	ktime_t now;
 
+	guard(cpus_read_lock)();
+	guard(mutex)(&cpu_preferred_mutex);
+
 	now = ktime_get();
 	delta_ns = ktime_to_ns(ktime_sub(now, sg_ctx.time));
 
@@ -238,8 +247,6 @@ static void steal_governor_loop(struct work_struct *work)
 	 * module restores to active if CPU hotplug breaks those assumptions.
 	 */
 	if (!preferred_cpus_valid()) {
-		guard(cpus_read_lock)();
-
 		reset_preferred_cpumask();
 		return;
 	}
@@ -276,6 +283,8 @@ static void __exit steal_governor_exit(void)
 {
 	disable_delayed_work_sync(&sg_ctx.work);
 	guard(cpus_read_lock)();
+	guard(mutex)(&cpu_preferred_mutex);
+
 	reset_preferred_cpumask();
 	pr_info("disabled\n");
 }
diff --git a/include/linux/cpumask.h b/include/linux/cpumask.h
index b13d099474bc..3716213d0730 100644
--- a/include/linux/cpumask.h
+++ b/include/linux/cpumask.h
@@ -124,6 +124,14 @@ extern struct cpumask __cpu_dying_mask;
 
 #ifdef CONFIG_PREFERRED_CPU
 extern struct cpumask __cpu_preferred_mask;
+
+/*
+ * Drivers must take cpus_read_lock(), then cpu_preferred_mutex and
+ * hold both across mask validation and updates. If cpus_write_lock()
+ * is held, it guarantees the exclusive access to the mask, thus
+ * holding cpu_preferred_mutex isn't necessary.
+ */
+extern struct mutex cpu_preferred_mutex;
 #else
 #define __cpu_preferred_mask __cpu_active_mask
 #endif
@@ -1199,7 +1207,8 @@ void init_cpu_possible(const struct cpumask *src);
  * reset_preferred_cpumask - reset preferred CPUs to the active mask.
  *
  * Callers must serialize against CPU hotplug and other preferred cpumask
- * writers.
+ * writers. Requires cpus_read_lock() and cpu_preferred_mutex; or
+ * cpus_write_lock().
  */
 static inline void reset_preferred_cpumask(void)
 {
diff --git a/lib/cpumask.c b/lib/cpumask.c
index 5adb9874fbd0..2e16e19a2bdb 100644
--- a/lib/cpumask.c
+++ b/lib/cpumask.c
@@ -166,3 +166,8 @@ unsigned int cpumask_any_distribute(const struct cpumask *srcp)
 	return next;
 }
 EXPORT_SYMBOL(cpumask_any_distribute);
+
+#ifdef CONFIG_PREFERRED_CPU
+DEFINE_MUTEX(cpu_preferred_mutex);
+EXPORT_SYMBOL_GPL(cpu_preferred_mutex);
+#endif
-- 
2.53.0

[PATCH 5/7] virt: steal_governor: warn on unexpected preferred mask states

[Yury Norov]

Now that preferred cpumask access is properly serialized, and CPU
hotplug restores all remaining active CPUs as preferred when the
last preferred CPU gets offlined, the corresponding logic in the
driver is redundant.

Instead of removing it completely, add WARN_ON()s in the corresponding
paths to harden the preferred CPUs mask management and highlight that
this condition is not expected during normal operation.

Signed-off-by: Yury Norov <ynorov@nvidia.com>
---
 drivers/virt/steal_governor.c | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/drivers/virt/steal_governor.c b/drivers/virt/steal_governor.c
index 8cd5c0f3e3ef..c34d91267ec1 100644
--- a/drivers/virt/steal_governor.c
+++ b/drivers/virt/steal_governor.c
@@ -142,6 +142,11 @@ static void decrease_preferred_cpus(void)
 
 	lockdep_assert_cpus_held();
 	lockdep_assert_held(&cpu_preferred_mutex);
+
+	/* We can't end up with no preferred CPUs under normal conditions. */
+	if (WARN_ON(cpumask_empty(cpu_preferred_mask)))
+		return;
+
 	cpu = cpumask_first_and(housekeeping_cpumask(HK_TYPE_KERNEL_NOISE),
 				cpu_preferred_mask);
 	if (cpu >= nr_cpu_ids)
@@ -150,8 +155,6 @@ static void decrease_preferred_cpus(void)
 	/* Always leave first housekeeping core as preferred. */
 	first_hk_core = topology_sibling_cpumask(cpu);
 	cpu = cpumask_last(cpu_preferred_mask);
-	if (cpu >= nr_cpu_ids)
-		return;
 
 	/* Find the last CPU which doesn't belong to that first hk_core. */
 	if (!cpumask_test_cpu(cpu, first_hk_core)) {
@@ -246,7 +249,7 @@ static void steal_governor_loop(struct work_struct *work)
 	 * else: steal ratio is within bounds. Still do design checks so that
 	 * module restores to active if CPU hotplug breaks those assumptions.
 	 */
-	if (!preferred_cpus_valid()) {
+	if (WARN_ON(!preferred_cpus_valid())) {
 		reset_preferred_cpumask();
 		return;
 	}
-- 
2.53.0

[PATCH 6/7] virt: steal_governor: restart sampling immediately after CPU hotplug

[Yury Norov]

CPU hotplug between samples can distort the calculated steal ratio because
the governor divides accumulated steal time by the current active CPU
count, even when that count changed during the sampling interval.

For example, assume a constant 3% steal rate on each CPU, with four active
CPUs for 900ms followed by two active CPUs for 100ms. The accumulated
steal time is:

    4 * 900ms * 3% + 2 * 100ms * 3% = 108ms + 6ms = 114ms

The elapsed time is 1000ms. Using the final count of two active CPUs, the
governor calculates:

    114ms / (2 * 1000ms) = 5.7%

This exceeds the default 5% high threshold, even though the actual steal
rate stayed at 3%. Using the CPU counts over the whole interval gives:

    114ms / (4 * 900ms + 2 * 100ms) = 3%

Checking num_online_cpus() doesn't account for cases where CPUs go online
and offline such that the total number of online CPUs doesn't change. So,
implement the sampling restart on hotplug event.

Fixes: 27d47ebce4d6 ("virt/steal_governor: Implement steal_governor policy loop")
Signed-off-by: Yury Norov <ynorov@nvidia.com>
---
 drivers/virt/steal_governor.c | 37 ++++++++++++++++++++++++++++++++---
 1 file changed, 34 insertions(+), 3 deletions(-)

diff --git a/drivers/virt/steal_governor.c b/drivers/virt/steal_governor.c
index c34d91267ec1..f71dbf27e4a9 100644
--- a/drivers/virt/steal_governor.c
+++ b/drivers/virt/steal_governor.c
@@ -14,6 +14,7 @@
 #define pr_fmt(fmt) KBUILD_MODNAME ": " fmt
 
 #include <linux/cleanup.h>
+#include <linux/cpuhotplug.h>
 #include <linux/cpuhplock.h>
 #include <linux/cpumask.h>
 #include <linux/init.h>
@@ -42,6 +43,8 @@ struct steal_governor {
 	unsigned long		delay;
 	unsigned int		high_threshold;
 	unsigned int		low_threshold;
+	int			hp_state;
+	bool			sample_invalidated;
 	struct delayed_work	work;
 };
 
@@ -217,6 +220,13 @@ static bool preferred_cpus_valid(void)
 	return true;
 }
 
+static int hotplug_callback(unsigned int cpu)
+{
+	sg_ctx.sample_invalidated = true;
+	mod_delayed_work(system_wq, &sg_ctx.work, 0);
+	return 0;
+}
+
 static void steal_governor_loop(struct work_struct *work)
 {
 	u64 curr_steal, delta_steal, delta_ns, steal_ratio;
@@ -233,6 +243,11 @@ static void steal_governor_loop(struct work_struct *work)
 	sg_ctx.steal = curr_steal;
 	sg_ctx.time = now;
 
+	if (sg_ctx.sample_invalidated) {
+		sg_ctx.sample_invalidated = false;
+		goto requeue_work;
+	}
+
 	/*
 	 * steal_ratio = (delta_steal * 100*100)/(delta_ns * num_cpus())
 	 * To avoid possible overflow, divide the denominator early.
@@ -254,11 +269,14 @@ static void steal_governor_loop(struct work_struct *work)
 		return;
 	}
 
-	schedule_delayed_work(&sg_ctx.work, sg_ctx.delay);
+requeue_work:
+	mod_delayed_work(system_wq, &sg_ctx.work, sg_ctx.delay);
 }
 
 static int __init steal_governor_init(void)
 {
+	int ret;
+
 #ifdef CONFIG_XEN
 	if (xen_initial_domain()) {
 		pr_err("Cannot load in Xen Dom0 (Host OS). Driver is for guests only.\n");
@@ -273,8 +291,20 @@ static int __init steal_governor_init(void)
 	}
 
 	INIT_DELAYED_WORK(&sg_ctx.work, steal_governor_loop);
-	sg_ctx.steal = get_system_steal_time();
-	sg_ctx.time = ktime_get();
+	/* Lock to prevent unobserved hotplug event during setup. */
+	scoped_guard(cpus_read_lock) {
+		sg_ctx.steal = get_system_steal_time();
+		sg_ctx.time = ktime_get();
+		sg_ctx.sample_invalidated = false;
+		ret = cpuhp_setup_state_nocalls_cpuslocked(CPUHP_AP_ONLINE_DYN,
+							 "virt/steal_governor:online",
+							 hotplug_callback,
+							 hotplug_callback);
+	}
+	if (ret < 0)
+		return ret;
+	sg_ctx.hp_state = ret;
+
 	schedule_delayed_work(&sg_ctx.work, sg_ctx.delay);
 	pr_info("enabled. interval: %ums, high_threshold: %u, low_threshold: %u\n",
 		jiffies_to_msecs(sg_ctx.delay), sg_ctx.high_threshold, sg_ctx.low_threshold);
@@ -285,6 +315,7 @@ static int __init steal_governor_init(void)
 static void __exit steal_governor_exit(void)
 {
 	disable_delayed_work_sync(&sg_ctx.work);
+	cpuhp_remove_state_nocalls(sg_ctx.hp_state);
 	guard(cpus_read_lock)();
 	guard(mutex)(&cpu_preferred_mutex);
 
-- 
2.53.0

[PATCH 7/7] cpu: make preferred CPU policy writers mutually exclusive

[Yury Norov]

Serializing preferred mask updates still lets competing policy drivers
repeatedly override each other's selections.

Disallow multiple policy drivers running at the same time with the
exclusive ownership rule.

Suggested-by: Peter Zijlstra (Intel) <peterz@infradead.org>
Signed-off-by: Yury Norov <ynorov@nvidia.com>
---
 Documentation/scheduler/sched-paravirt.rst |  2 +-
 drivers/virt/steal_governor.c              | 13 +++--
 include/linux/cpumask.h                    |  8 +++
 lib/cpumask.c                              | 62 ++++++++++++++++++++++
 4 files changed, 79 insertions(+), 6 deletions(-)

diff --git a/Documentation/scheduler/sched-paravirt.rst b/Documentation/scheduler/sched-paravirt.rst
index 311cdbf2722b..d8f446c3ab6e 100644
--- a/Documentation/scheduler/sched-paravirt.rst
+++ b/Documentation/scheduler/sched-paravirt.rst
@@ -48,7 +48,7 @@ Notes:
    makes decisions should enable it. For example, steal_governor driver
    (``CONFIG_STEAL_GOVERNOR``). On enabling the driver, CPU preferred state
    can change based on steal time. Without the driver, preferred CPUs is
-   same as active CPUs.
+   same as active CPUs. Only one governor may be enabled at a time.
 
 2. This feature works for the FAIR class only.
 
diff --git a/drivers/virt/steal_governor.c b/drivers/virt/steal_governor.c
index f71dbf27e4a9..b456b03714e4 100644
--- a/drivers/virt/steal_governor.c
+++ b/drivers/virt/steal_governor.c
@@ -291,6 +291,10 @@ static int __init steal_governor_init(void)
 	}
 
 	INIT_DELAYED_WORK(&sg_ctx.work, steal_governor_loop);
+	ret = claim_preferred_cpumask(&sg_ctx);
+	if (ret)
+		return ret;
+
 	/* Lock to prevent unobserved hotplug event during setup. */
 	scoped_guard(cpus_read_lock) {
 		sg_ctx.steal = get_system_steal_time();
@@ -301,8 +305,10 @@ static int __init steal_governor_init(void)
 							 hotplug_callback,
 							 hotplug_callback);
 	}
-	if (ret < 0)
+	if (ret < 0) {
+		release_preferred_cpumask(&sg_ctx);
 		return ret;
+	}
 	sg_ctx.hp_state = ret;
 
 	schedule_delayed_work(&sg_ctx.work, sg_ctx.delay);
@@ -316,10 +322,7 @@ static void __exit steal_governor_exit(void)
 {
 	disable_delayed_work_sync(&sg_ctx.work);
 	cpuhp_remove_state_nocalls(sg_ctx.hp_state);
-	guard(cpus_read_lock)();
-	guard(mutex)(&cpu_preferred_mutex);
-
-	reset_preferred_cpumask();
+	release_preferred_cpumask(&sg_ctx);
 	pr_info("disabled\n");
 }
 
diff --git a/include/linux/cpumask.h b/include/linux/cpumask.h
index 3716213d0730..1f4d93adb917 100644
--- a/include/linux/cpumask.h
+++ b/include/linux/cpumask.h
@@ -132,6 +132,14 @@ extern struct cpumask __cpu_preferred_mask;
  * holding cpu_preferred_mutex isn't necessary.
  */
 extern struct mutex cpu_preferred_mutex;
+
+/*
+ * Drivers must claim exclusive ownership of cpu_preferred_mask
+ * for the whole period of operation.
+ */
+int claim_preferred_cpumask(const void *owner);
+void release_preferred_cpumask(const void *owner);
+bool preferred_cpumask_owned_by(const void *owner);
 #else
 #define __cpu_preferred_mask __cpu_active_mask
 #endif
diff --git a/lib/cpumask.c b/lib/cpumask.c
index 2e16e19a2bdb..ee710c3c8920 100644
--- a/lib/cpumask.c
+++ b/lib/cpumask.c
@@ -3,8 +3,10 @@
 #include <linux/kernel.h>
 #include <linux/bitops.h>
 #include <linux/cpumask.h>
+#include <linux/cpuhplock.h>
 #include <linux/export.h>
 #include <linux/memblock.h>
+#include <linux/mutex.h>
 #include <linux/numa.h>
 
 /* These are not inline because of header tangles. */
@@ -170,4 +172,64 @@ EXPORT_SYMBOL(cpumask_any_distribute);
 #ifdef CONFIG_PREFERRED_CPU
 DEFINE_MUTEX(cpu_preferred_mutex);
 EXPORT_SYMBOL_GPL(cpu_preferred_mutex);
+
+static const void *preferred_cpumask_owner;
+
+/**
+ * preferred_cpumask_owned_by - test preferred CPU policy ownership
+ * @owner: policy identity
+ *
+ * Return: true if @owner holds the policy, false otherwise.
+ */
+bool preferred_cpumask_owned_by(const void *owner)
+{
+	lockdep_assert_held(&cpu_preferred_mutex);
+
+	return preferred_cpumask_owner == owner;
+}
+EXPORT_SYMBOL_GPL(preferred_cpumask_owned_by);
+
+/**
+ * claim_preferred_cpumask - acquire exclusive preferred CPU policy ownership
+ * @owner: non-NULL identity that remains valid until ownership is released
+ *
+ * Call before starting policy work, normally on module init.
+ *
+ * Return: 0 on success, -EBUSY if claimed, or -EINVAL for a NULL identity.
+ */
+int claim_preferred_cpumask(const void *owner)
+{
+	if (!owner)
+		return -EINVAL;
+
+	guard(cpus_read_lock)();
+	guard(mutex)(&cpu_preferred_mutex);
+
+	if (preferred_cpumask_owner)
+		return -EBUSY;
+
+	preferred_cpumask_owner = owner;
+	reset_preferred_cpumask();
+	return 0;
+}
+EXPORT_SYMBOL_GPL(claim_preferred_cpumask);
+
+/**
+ * release_preferred_cpumask - release exclusive preferred CPU policy ownership
+ * @owner: identity used to claim ownership
+ *
+ * Call after stopping policy work, normally on module exit.
+ */
+void release_preferred_cpumask(const void *owner)
+{
+	guard(cpus_read_lock)();
+	guard(mutex)(&cpu_preferred_mutex);
+
+	if (WARN_ON_ONCE(!owner || !preferred_cpumask_owned_by(owner)))
+		return;
+
+	reset_preferred_cpumask();
+	preferred_cpumask_owner = NULL;
+}
+EXPORT_SYMBOL_GPL(release_preferred_cpumask);
 #endif
-- 
2.53.0