Re: [RFC PATCH 04/16] sched/core: Activate blocked donor when no owner is found

[K Prateek Nayak <kprateek.nayak@amd.com>]

Hello Andrea,

On 8/26/2026 10:50 PM, K Prateek Nayak wrote:
>>> XXX: Is there a better way to handle this? If we can confirm a owner in
>>> find_proxy_task(), we don't need to do a spurious wakeup of every task
>>> observing !owner.
>>>
>>> proxy_resched_idle() until owner appears in an option but it will spin
>>> until next owner appears.
>>
>> IIUC, the owner can remain NULL until the waiter selected by mutex_unlock() gets
>> CPU time and acquires the mutex, so proxy_resched_idle() could spin for longer
>> than just the unlock critical section.
>>
>> Maybe we could instead force a handoff from mutex_unlock_slowpath() when proxy
>> execution is enabled and the mutex has waiters? This would keep the owner
>> identifiable and avoid waking every task that happens to observe !owner.
> 
> I think that negates some of the benefits of the optimistic spinning +
> mutex_try_lock(). I'll see if it makes any difference to the benchmark
> results if we always force a handoff for MUTEX_FLAG_WAITERS.
> 
> wait_lock should give enough guarantee that the waiter cannot simple
> disappear before the handoff after MUTEX_FLAG_PICKUP is set since
> waiter has to try at least one mutex_trylock() under wait_lock before
> checking for pending signals.

Below  are the results from few experiments. All diffs pasted below are
based on John's tree at:

  https://github.com/johnstultz-work/linux-dev.git proxy-exec-v31-7.2-rc4

at commit 06ac43db4d8e ("[ANNOTATION] === Needs confirmation of
functionality past this point ===") with CONFIG_SCHED_PROXY_EXEC=y.

All diffs are very experimental: virtme-ng or testing with a disposable
environment is recommended.

============================
Experiment 1: Simple Handoff
============================

If I do a simple handoff like below, sched-messaging goes pretty bad:

diff --git a/kernel/locking/mutex.c b/kernel/locking/mutex.c
index 8a85912d7ee6..da14a49e4fa2 100644
--- a/kernel/locking/mutex.c
+++ b/kernel/locking/mutex.c
@@ -1009,7 +1009,7 @@ static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigne
 		MUTEX_WARN_ON(__owner_task(owner) != current);
 		MUTEX_WARN_ON(owner & MUTEX_FLAG_PICKUP);
 
-		if (sched_proxy_exec() && current->blocked_donor) {
+		if (sched_proxy_exec() && (current->blocked_donor || (owner & MUTEX_FLAG_WAITERS))) {
 			/* force handoff if we have a blocked_donor */
 			owner = MUTEX_FLAG_HANDOFF;
 			break;
---

With just a simple handoff on waiters, we have:

    ==================================================================
    Test          : sched-messaging
    Units         : Normalized time in seconds
    Interpretation: Lower is better
    Statistic     : AMean
    ==================================================================
      Test:                   vanilla                 handoff
       1-groups:         3.12 (0.00 pct)         3.47 (-11.21 pct)  *
       2-groups:         3.43 (0.00 pct)         4.33 (-26.23 pct)  *
       4-groups:         4.05 (0.00 pct)         5.95 (-46.91 pct)
       8-groups:         4.29 (0.00 pct)         9.56 (-122.84 pct)
      16-groups:         5.89 (0.00 pct)        12.29 (-108.65 pct)

     * Data points have > 10% run to run variance on all versions


==================================================
Experiment 2: Allow steal until next task is found
==================================================

If we open the opportunity to allow stealing of mutex until the next waiter
is found, the results are ever so slightly slightly better:

diff --git a/kernel/locking/mutex.c b/kernel/locking/mutex.c
index 8a85912d7ee6..5ebb2624b331 100644
--- a/kernel/locking/mutex.c
+++ b/kernel/locking/mutex.c
@@ -92,7 +92,16 @@ static inline struct task_struct *__mutex_trylock_common(struct mutex *lock, boo
 		unsigned long task = owner & ~MUTEX_FLAGS;
 
 		if (task) {
-			if (flags & MUTEX_FLAG_PICKUP) {
+			if (sched_proxy_exec() && (flags & MUTEX_FLAG_STEAL)) {
+				/*
+				 * STEAL cannot be set after HANDOFF has been
+				 * initiated. If STEAL is set, clear it and
+				 * preserve other flags
+				 */
+				MUTEX_WARN_ON(flags & (MUTEX_FLAG_PICKUP));
+				flags &= ~MUTEX_FLAG_STEAL;
+				task = curr;
+			}else if (flags & MUTEX_FLAG_PICKUP) {
 				if (task != curr)
 					break;
 				flags &= ~MUTEX_FLAG_PICKUP;
@@ -104,7 +113,7 @@ static inline struct task_struct *__mutex_trylock_common(struct mutex *lock, boo
 				break;
 			}
 		} else {
-			MUTEX_WARN_ON(flags & (MUTEX_FLAG_HANDOFF | MUTEX_FLAG_PICKUP));
+			MUTEX_WARN_ON(flags & (MUTEX_FLAG_HANDOFF | MUTEX_FLAG_PICKUP | MUTEX_FLAG_STEAL));
 			task = curr;
 		}
 
@@ -274,7 +283,29 @@ static void __mutex_handoff(struct mutex *lock, struct task_struct *task)
 		new |= (unsigned long)task;
 		if (task)
 			new |= MUTEX_FLAG_PICKUP;
+		if (atomic_long_try_cmpxchg_release(&lock->owner, &owner, new))
+			break;
+	}
+}
+
+static void __mutex_steal(struct mutex *lock, struct task_struct *task)
+{
+	unsigned long owner = atomic_long_read(&lock->owner);
+
+	for (;;) {
+		unsigned long new;
 
+		/* Lock was successfully stolen. */
+		if (__owner_task(owner) != current)
+			break;
+
+		MUTEX_WARN_ON(!(__owner_flags(owner) & MUTEX_FLAG_STEAL));
+		MUTEX_WARN_ON(owner & MUTEX_FLAG_PICKUP);
+
+		new = (owner & MUTEX_FLAG_WAITERS);
+		new |= (unsigned long)task;
+		if (task)
+			new |= MUTEX_FLAG_PICKUP;
 		if (atomic_long_try_cmpxchg_release(&lock->owner, &owner, new))
 			break;
 	}
@@ -389,7 +420,17 @@ bool mutex_spin_on_owner(struct mutex *lock, struct task_struct *owner,
 
 	lockdep_assert_preemption_disabled();
 
-	while (__mutex_owner(lock) == owner) {
+	for (;;) {
+		unsigned long __owner = atomic_long_read(&lock->owner);
+
+		/* If the owner changed, break out. */
+		if (__owner_task(__owner) != owner)
+			break;
+
+		/* If lock can be stolen, break out. */
+		if (sched_proxy_exec() && (__owner_flags(__owner) & MUTEX_FLAG_STEAL))
+			break;
+
 		/*
 		 * Ensure we emit the owner->on_cpu, dereference _after_
 		 * checking lock->owner still matches owner. And we already
@@ -985,6 +1026,7 @@ static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigne
 	struct mutex_waiter *waiter;
 	unsigned long owner;
 	unsigned long flags;
+	bool steal;
 
 	mutex_release(&lock->dep_map, ip);
 	__release(lock);
@@ -1006,19 +1048,31 @@ static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigne
 	 */
 	owner = atomic_long_read(&lock->owner);
 	for (;;) {
+		unsigned long owner_flags;
+
 		MUTEX_WARN_ON(__owner_task(owner) != current);
 		MUTEX_WARN_ON(owner & MUTEX_FLAG_PICKUP);
 
-		if (sched_proxy_exec() && current->blocked_donor) {
-			/* force handoff if we have a blocked_donor */
-			owner = MUTEX_FLAG_HANDOFF;
-			break;
-		}
-
 		if (owner & MUTEX_FLAG_HANDOFF)
 			break;
 
-		if (atomic_long_try_cmpxchg_release(&lock->owner, &owner, __owner_flags(owner))) {
+		owner_flags = __owner_flags(owner);
+		if (sched_proxy_exec()) {
+			if (current->blocked_donor) {
+				/* force handoff if we have a blocked_donor */
+				owner = MUTEX_FLAG_HANDOFF;
+				break;
+			}
+
+			if (owner & MUTEX_FLAG_WAITERS)
+				owner_flags = owner | MUTEX_FLAG_STEAL;
+		}
+
+		if (atomic_long_try_cmpxchg_release(&lock->owner, &owner, owner_flags)) {
+			if (owner_flags & MUTEX_FLAG_STEAL) {
+				steal = true;
+				break;
+			}
 			if (owner & MUTEX_FLAG_WAITERS)
 				break;
 
@@ -1071,6 +1125,9 @@ static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigne
 	if (owner & MUTEX_FLAG_HANDOFF)
 		__mutex_handoff(lock, next);
 
+	if (sched_proxy_exec() && steal)
+		__mutex_steal(lock, next);
+
 	raw_spin_unlock(&current->blocked_lock);
 	raw_spin_unlock_irqrestore(&lock->wait_lock, flags);
 	if (next) {
diff --git a/kernel/locking/mutex.h b/kernel/locking/mutex.h
index 3e263e98e5fc..eb4180745da0 100644
--- a/kernel/locking/mutex.h
+++ b/kernel/locking/mutex.h
@@ -33,8 +33,9 @@ struct mutex_waiter {
 #define MUTEX_FLAG_WAITERS	0x01
 #define MUTEX_FLAG_HANDOFF	0x02
 #define MUTEX_FLAG_PICKUP	0x04
+#define MUTEX_FLAG_STEAL	0x08
 
-#define MUTEX_FLAGS		0x07
+#define MUTEX_FLAGS		0x0F
 
 /*
  * Internal helper function; C doesn't allow us to hide it :/
---

With that small steal opportunity, we have:

    ==================================================================
    Test          : sched-messaging
    Units         : Normalized time in seconds
    Interpretation: Lower is better
    Statistic     : AMean
    ==================================================================
      Test:                   vanilla                 handoff             steal + handoff
       1-groups:         3.12 (0.00 pct)         3.47 (-11.21 pct)       3.63 (-16.34 pct)  *
       2-groups:         3.43 (0.00 pct)         4.33 (-26.23 pct)       4.14 (-20.69 pct)  *
       4-groups:         4.05 (0.00 pct)         5.95 (-46.91 pct)       5.45 (-34.56 pct)
       8-groups:         4.29 (0.00 pct)         9.56 (-122.84 pct)      7.80 (-81.81 pct)
      16-groups:         5.89 (0.00 pct)        12.29 (-108.65 pct)     11.89 (-101.86 pct)

     * Data points have > 10% run to run variance on all versions


So, the opportunity must be extended to allow stealing until the the waiter
wakes up for !HANDOFF cases. Few complications with that are:

o We cannot continue to persist the old owner after mutex_unlock() since that
  owner can die, block on other mutex, etc. and that breaks queuing on owner
  since new waiters can go and block on a dead task / task blocked on an
  incorrect owner.

o We cannot allow steal after handoff to new owner because __mutex_lock() will
  resolve to new owner that hasn't woken up yet and waiters start queuing on
  it but a concurrent task can come steal the lock and break proxy. Not very
  intuitive; adds a lot of complexity.


============================================
Experiment 3: STEAL + Temporary swap to idle
============================================

the unlock will temporarily swap to rq->idle of the lock owner's CPU with
MUTEX_FLAG_STEAL set to allow grabbing the task until the the waiter wakes
up and manages to grab the task itself for !HANDOFF cases. With that,
numbers are very close:

diff --git a/kernel/locking/mutex.c b/kernel/locking/mutex.c
index 8a85912d7ee6..187f95544453 100644
--- a/kernel/locking/mutex.c
+++ b/kernel/locking/mutex.c
@@ -92,7 +92,16 @@ static inline struct task_struct *__mutex_trylock_common(struct mutex *lock, boo
 		unsigned long task = owner & ~MUTEX_FLAGS;
 
 		if (task) {
-			if (flags & MUTEX_FLAG_PICKUP) {
+			if (sched_proxy_exec() && (flags & MUTEX_FLAG_STEAL)) {
+				/*
+				 * STEAL cannot be set after HANDOFF has been
+				 * initiated. If STEAL is set, clear it and
+				 * preserve other flags
+				 */
+				MUTEX_WARN_ON(flags & (MUTEX_FLAG_PICKUP));
+				flags &= ~MUTEX_FLAG_STEAL;
+				task = curr;
+			}else if (flags & MUTEX_FLAG_PICKUP) {
 				if (task != curr)
 					break;
 				flags &= ~MUTEX_FLAG_PICKUP;
@@ -104,7 +113,7 @@ static inline struct task_struct *__mutex_trylock_common(struct mutex *lock, boo
 				break;
 			}
 		} else {
-			MUTEX_WARN_ON(flags & (MUTEX_FLAG_HANDOFF | MUTEX_FLAG_PICKUP));
+			MUTEX_WARN_ON(flags & (MUTEX_FLAG_HANDOFF | MUTEX_FLAG_PICKUP | MUTEX_FLAG_STEAL));
 			task = curr;
 		}
 
@@ -242,7 +251,41 @@ __mutex_remove_waiter(struct mutex *lock, struct mutex_waiter *waiter)
 	__must_hold(&lock->wait_lock)
 {
 	if (list_empty(&waiter->list)) {
-		__mutex_clear_flag(lock, MUTEX_FLAGS);
+		/*
+		 * The last waiter can be interrupted before the full
+		 * unlock with STEAL is done.
+		 *
+		 *  LOCK lock->wait_lock
+		 *
+		 *  __mutex_trylock()
+		 *    // Sees old owner			__mutex_unlock_slowpath()
+		 *    return owner;			  atomic_long_cmpxchg_release(&owner, idle | STEAL)
+		 *					    // Succeeds
+		 *  if (signal_pending())
+		 *    goto err;
+		 *
+		 * err:
+		 *  __mutex_remove_waiter()
+		 *    __mutex_clear_flag(MUTEX_FLAGS)
+		 *    lock->first_waiter = NULL;
+		 *
+		 *  UNLOCK lock->wait_lock		  LOCK lock->wait_lock
+		 *					  waiter = lock->first_waiter; // NULL
+		 *					  // No wakeup
+		 *
+		 *  !!! lock->owner stuck as rq->idle without STEAL set !!!
+		 *
+		 * Persis the STEAL flag to prevent an idle task to
+		 * linger as lock owner. __mutex_trylock_fast() will
+		 * fail temporarily for first contender but following
+		 * __mutex_trylock_common() will do the right thing.
+		 *
+		 * XXX: This can also be solved by doing a
+		 * atomic_try_cmpxchg() or__mutex_clear_flag() in
+		 * __mutex_unlock_slowpath() if steal is set but no
+		 * waiter is found under lock->wait_lock.
+		 */
+		__mutex_clear_flag(lock, MUTEX_FLAGS & ~MUTEX_FLAG_STEAL);
 		lock->first_waiter = NULL;
 	} else {
 		if (lock->first_waiter == waiter)
@@ -274,7 +317,6 @@ static void __mutex_handoff(struct mutex *lock, struct task_struct *task)
 		new |= (unsigned long)task;
 		if (task)
 			new |= MUTEX_FLAG_PICKUP;
-
 		if (atomic_long_try_cmpxchg_release(&lock->owner, &owner, new))
 			break;
 	}
@@ -389,7 +431,17 @@ bool mutex_spin_on_owner(struct mutex *lock, struct task_struct *owner,
 
 	lockdep_assert_preemption_disabled();
 
-	while (__mutex_owner(lock) == owner) {
+	for (;;) {
+		unsigned long __owner = atomic_long_read(&lock->owner);
+
+		/* If the owner changed, break out. */
+		if (__owner_task(__owner) != owner)
+			break;
+
+		/* If lock can be stolen, break out. */
+		if (sched_proxy_exec() && (__owner_flags(__owner) & MUTEX_FLAG_STEAL))
+			break;
+
 		/*
 		 * Ensure we emit the owner->on_cpu, dereference _after_
 		 * checking lock->owner still matches owner. And we already
@@ -1006,19 +1058,42 @@ static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigne
 	 */
 	owner = atomic_long_read(&lock->owner);
 	for (;;) {
+		unsigned long owner_flags;
+
 		MUTEX_WARN_ON(__owner_task(owner) != current);
 		MUTEX_WARN_ON(owner & MUTEX_FLAG_PICKUP);
 
-		if (sched_proxy_exec() && current->blocked_donor) {
-			/* force handoff if we have a blocked_donor */
-			owner = MUTEX_FLAG_HANDOFF;
-			break;
-		}
-
 		if (owner & MUTEX_FLAG_HANDOFF)
 			break;
 
-		if (atomic_long_try_cmpxchg_release(&lock->owner, &owner, __owner_flags(owner))) {
+		owner_flags = __owner_flags(owner);
+		if (sched_proxy_exec()) {
+			if (current->blocked_donor) {
+				/* force handoff if we have a blocked_donor */
+				owner = MUTEX_FLAG_HANDOFF;
+				break;
+			}
+
+			if (owner & MUTEX_FLAG_WAITERS) {
+				unsigned long idle;
+				/*
+				 * Swap the owner to current CPU's idle task
+				 * with a STEAL flag.
+				 *
+				 * The lock is free to be stolen and
+				 * __mutex_owner() will resolve to idle task
+				 * that is always ->on_rq on this CPU.
+				 *
+				 * Proxy donors will temporarily migrate here
+				 * before a wakeup or an optimistic spinner
+				 * can grab the lock.
+				 */
+				idle = (unsigned long)idle_task(raw_smp_processor_id());
+				owner_flags = idle | MUTEX_FLAG_STEAL | owner_flags;
+			}
+		}
+
+		if (atomic_long_try_cmpxchg_release(&lock->owner, &owner, owner_flags)) {
 			if (owner & MUTEX_FLAG_WAITERS)
 				break;
 
diff --git a/kernel/locking/mutex.h b/kernel/locking/mutex.h
index 3e263e98e5fc..eb4180745da0 100644
--- a/kernel/locking/mutex.h
+++ b/kernel/locking/mutex.h
@@ -33,8 +33,9 @@ struct mutex_waiter {
 #define MUTEX_FLAG_WAITERS	0x01
 #define MUTEX_FLAG_HANDOFF	0x02
 #define MUTEX_FLAG_PICKUP	0x04
+#define MUTEX_FLAG_STEAL	0x08
 
-#define MUTEX_FLAGS		0x07
+#define MUTEX_FLAGS		0x0F
 
 /*
  * Internal helper function; C doesn't allow us to hide it :/
---

The results with temporary switch to idle + STEAL are:

    ==================================================================
    Test          : sched-messaging
    Units         : Normalized time in seconds
    Interpretation: Lower is better
    Statistic     : AMean
    ==================================================================
      Test:                   vanilla                 handoff             STEAL + handoff          idle + STEAL
       1-groups:         3.12 (0.00 pct)         3.47 (-11.21 pct)       3.63 (-16.34 pct)       3.59 (-15.06 pct)
       2-groups:         3.43 (0.00 pct)         4.33 (-26.23 pct)       4.14 (-20.69 pct)       3.48 (-1.45 pct)
       4-groups:         4.05 (0.00 pct)         5.95 (-46.91 pct)       5.45 (-34.56 pct)       4.00 (1.23 pct)
       8-groups:         4.29 (0.00 pct)         9.56 (-122.84 pct)      7.80 (-81.81 pct)       4.31 (-0.46 pct)
      16-groups:         5.89 (0.00 pct)        12.29 (-108.65 pct)     11.89 (-101.86 pct)      5.91 (-0.33 pct)

      * Data points have > 10% run to run variance on all versions


My machine has held up for some time with Experiment 3 so I'm
fairly confident at the very least mutual exclusion is holding
up - I haven't seen any lockups / hung task either so hopefully
other bits are fine too :-)

-- 
Thanks and Regards,
Prateek