sched: Add blocked_donor link to task for smarter mutex handoffs

Add link to the task this task is proxying for, and use it so
the mutex owner can do an intelligent hand-off of the mutex to
the task that the owner is running on behalf.

[jstultz: This patch was split out from larger proxy patch]
Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
Signed-off-by: Juri Lelli <juri.lelli@redhat.com>
Signed-off-by: Valentin Schneider <valentin.schneider@arm.com>
Signed-off-by: Connor O'Brien <connoro@google.com>
Signed-off-by: John Stultz <jstultz@google.com>
Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
Link: https://patch.msgid.link/20260512025635.2840817-8-jstultz@google.com

diff --git a/include/linux/sched.h b/include/linux/sched.h
index ec170663f99b308e746af1bee5ffd9c20b4c6050..e2f127a7ca0d3e379f2a5581bb4ff9c5003513cb 100644 (file)
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -1250,6 +1250,13 @@ struct task_struct {
        struct mutex                    *blocked_on;    /* lock we're blocked on */
        raw_spinlock_t                  blocked_lock;
 
+       /*
+        * The task that is boosting this task; a back link for the current
+        * donor stack. Set in schedule() -> find_proxy_task() and only stable
+        * under preempt_disable().
+        */
+       struct task_struct              *blocked_donor;
+
 #ifdef CONFIG_DETECT_HUNG_TASK_BLOCKER
        /*
         * Encoded lock address causing task block (lower 2 bits = type from

diff --git a/init/init_task.c b/init/init_task.c
index 3ecd66fbd563b8e1977ce0bbd67ed96aa3e4ab01..674d174e2e6a746f884148f91265d49bff3ada9f 100644 (file)
--- a/init/init_task.c
+++ b/init/init_task.c
@@ -200,6 +200,7 @@ struct task_struct init_task __aligned(L1_CACHE_BYTES) = {
        .mems_allowed_seq = SEQCNT_SPINLOCK_ZERO(init_task.mems_allowed_seq,
                                                 &init_task.alloc_lock),
 #endif
+       .blocked_donor = NULL,
 #ifdef CONFIG_RT_MUTEXES
        .pi_waiters     = RB_ROOT_CACHED,
        .pi_top_task    = NULL,

diff --git a/kernel/fork.c b/kernel/fork.c
index a679b24482347f911a0a82180b74ba112147b7a1..6fcca1db0af3465cba34799ea04270dd5b184dc4 100644 (file)
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -2224,6 +2224,7 @@ __latent_entropy struct task_struct *copy_process(
        lockdep_init_task(p);
 
        p->blocked_on = NULL; /* not blocked yet */
+       p->blocked_donor = NULL; /* nobody is boosting p yet */
 
 #ifdef CONFIG_BCACHE
        p->sequential_io        = 0;

diff --git a/kernel/locking/mutex.c b/kernel/locking/mutex.c
index a93d4c6bee1a3b78e60c220ab1e0561c1cb844ae..28677165785f754b0f60dd94fb2d017a24f00bb1 100644 (file)
--- a/kernel/locking/mutex.c
+++ b/kernel/locking/mutex.c
@@ -981,15 +981,22 @@ EXPORT_SYMBOL_GPL(ww_mutex_lock_interruptible);
 static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigned long ip)
        __releases(lock)
 {
-       struct task_struct *next = NULL;
+       struct task_struct *donor, *next = NULL;
        struct mutex_waiter *waiter;
-       DEFINE_WAKE_Q(wake_q);
        unsigned long owner;
        unsigned long flags;
 
        mutex_release(&lock->dep_map, ip);
        __release(lock);
 
+       /*
+        * Ensures the proxy donor stack is stable across unlock and handoff.
+        * Specifically, it avoids the case where current->blocked_donor is
+        * NULL when it is inspected while doing the unlock, but a preemption
+        * before taking the wake_lock would make it set and a hand-off is
+        * missed.
+        */
+       guard(preempt)();
        /*
         * Release the lock before (potentially) taking the spinlock such that
         * other contenders can get on with things ASAP.
@@ -1002,6 +1009,12 @@ 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) {
+                       /* force handoff if we have a blocked_donor */
+                       owner = MUTEX_FLAG_HANDOFF;
+                       break;
+               }
+
                if (owner & MUTEX_FLAG_HANDOFF)
                        break;
 
@@ -1014,20 +1027,53 @@ static noinline void __sched __mutex_unlock_slowpath(struct mutex *lock, unsigne
        }
 
        raw_spin_lock_irqsave(&lock->wait_lock, flags);
+       raw_spin_lock(&current->blocked_lock);
        debug_mutex_unlock(lock);
+
+       if (sched_proxy_exec()) {
+               /*
+                * If we have a task boosting current, and that task was boosting
+                * current through this lock, hand the lock to that task, as that
+                * is the highest waiter, as selected by the scheduling function.
+                */
+               donor = current->blocked_donor;
+               if (donor) {
+                       struct mutex *next_lock;
+
+                       raw_spin_lock_nested(&donor->blocked_lock, SINGLE_DEPTH_NESTING);
+                       next_lock = __get_task_blocked_on(donor);
+                       if (next_lock == lock) {
+                               next = get_task_struct(donor);
+                               __set_task_blocked_on_waking(donor, next_lock);
+                               current->blocked_donor = NULL;
+                       }
+                       raw_spin_unlock(&donor->blocked_lock);
+               }
+       }
+
+       /*
+        * Failing that, pick first on the wait list.
+        */
        waiter = lock->first_waiter;
-       if (waiter) {
-               next = waiter->task;
+       if (!next && waiter) {
+               next = get_task_struct(waiter->task);
 
+               raw_spin_lock_nested(&next->blocked_lock, SINGLE_DEPTH_NESTING);
                debug_mutex_wake_waiter(lock, waiter);
-               set_task_blocked_on_waking(next, lock);
-               wake_q_add(&wake_q, next);
+               __set_task_blocked_on_waking(next, lock);
+               raw_spin_unlock(&next->blocked_lock);
+
        }
 
        if (owner & MUTEX_FLAG_HANDOFF)
                __mutex_handoff(lock, next);
 
-       raw_spin_unlock_irqrestore_wake(&lock->wait_lock, flags, &wake_q);
+       raw_spin_unlock(&current->blocked_lock);
+       raw_spin_unlock_irqrestore(&lock->wait_lock, flags);
+       if (next) {
+               wake_up_process(next);
+               put_task_struct(next);
+       }
 }
 
 #ifndef CONFIG_DEBUG_LOCK_ALLOC

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index c7552869d5c4019ccdee6a663ade359b571a7d4d..4c6ceff3855e5ccd7ead9a48d7ac180d62e234f0 100644 (file)
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6827,7 +6827,17 @@ static void proxy_migrate_task(struct rq *rq, struct rq_flags *rf,
  * Find runnable lock owner to proxy for mutex blocked donor
  *
  * Follow the blocked-on relation:
- *   task->blocked_on -> mutex->owner -> task...
+ *
+ *                ,-> task
+ *                |     | blocked-on
+ *                |     v
+ *  blocked_donor |   mutex
+ *                |     | owner
+ *                |     v
+ *                `-- task
+ *
+ * and set the blocked_donor relation, this latter is used by the mutex
+ * code to find which (blocked) task to hand-off to.
  *
  * Lock order:
  *
@@ -6969,6 +6979,7 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
                 * rq, therefore holding @rq->lock is sufficient to
                 * guarantee its existence, as per ttwu_remote().
                 */
+               owner->blocked_donor = p;
        }
        WARN_ON_ONCE(owner && !owner->on_rq);
        return owner;
@@ -7125,6 +7136,7 @@ pick_again:
                        clear_task_blocked_on(prev, NULL);
 
                rq_set_donor(rq, next);
+               next->blocked_donor = NULL;
                if (unlikely(next->is_blocked && next->blocked_on)) {
                        next = find_proxy_task(rq, next, &rf);
                        if (!next) {