Re: [PATCH RFC] sched/proxy: Defer donor commit until after proxy resolution

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

Hello Xukai,

On 7/7/2026 9:16 PM, Xukai Wang wrote:
> pick_next_task() currently commits the selected task into the scheduling
> class state before proxy execution has resolved whether that task should
> really become the committed donor.
> 
> With proxy execution enabled, the task returned by pick_next_task() may
> be blocked. __schedule() then calls find_proxy_task() to resolve the
> proxy chain. If find_proxy_task() returns NULL, __schedule() retries
> through pick_again, but the picked donor has already gone through
> put_prev_set_next_task(). In that case the class current state was
> speculatively switched to a donor which is not the final scheduling
> decision.

As long as rq_lock is held, the pick will always select the same
donor since there cannot be a wakeup / migration until the rq_lock
is released.

If there is a change in the proxy state, and the lock needs to be
dropped, we do a proxy_resched_idle() and transiently switch to the
idle task's accounting context.

If we don't need to drop the rq_lock, we just retry the pick after
dropping the blocked_lock and the pick will converge to the same
donor unless ->balance() manages to pull something.

Very often, it is the same donor task pick converges on. I don't
see how this transient state is a big enough problem to see any
noticeable overhead.

> 
> Make pick_next_task() return only the selected candidate task. Move the
> put_prev_set_next_task() call into __schedule(), next to rq_set_donor(),
> so the class commit happens at the schedule level.
> 
> For proxy execution, keep the picked donor separate from the actual task
> to run:
> 
>   donor = task selected by pick_next_task()
>   next  = task returned by find_proxy_task(), possibly the mutex owner
> 
> Only commit the donor after proxy-chain resolution succeeds:
> 
>   put_prev_set_next_task(rq, prev_donor, donor);
>   rq_set_donor(rq, donor);
> 
> If find_proxy_task() returns NULL, the candidate donor has not been
> committed and __schedule() can retry without first undoing a speculative
> set_next() on that donor.
> 
> A proxy candidate is not necessarily the committed rq->donor anymore.
> Update proxy_deactivate() accordingly. If the task being blocked is
> still the committed donor, proxy_resched_idle() is needed to drop
> rq/class current references before block_task() clears ->on_rq. If
> it is only an uncommitted proxy candidate, it is still queued and can
> be blocked directly.
> 
> proxy_migrate_task() still switches the rq to idle before dropping the
> rq lock. Migrating a task found in the proxy chain abandons the current
> proxy pick attempt, and switching the committed donor to idle leaves
> rq->donor and the class current state in a defined state while the chain
> is modified and the task is attached elsewhere.
> 
> This keeps the proxy donor accounting semantics intact: the scheduling
> class state is committed to the donor, while the actual context switch
> may still be to the proxy owner returned by find_proxy_task().
> 
> Signed-off-by: Xukai Wang <kingxukai@zohomail.com>
> ---
>  kernel/sched/core.c | 70 ++++++++++++++++++++++++++++++-----------------------
>  1 file changed, 40 insertions(+), 30 deletions(-)
> 
> diff --git a/kernel/sched/core.c b/kernel/sched/core.c
> index 2e7cde033a31..3c1a23a9c12c 100644
> --- a/kernel/sched/core.c
> +++ b/kernel/sched/core.c
> @@ -6146,7 +6146,6 @@ __pick_next_task(struct rq *rq, struct rq_flags *rf)
>                 if (!p)
>                         p = pick_task_idle(rq, rf);
> 
> -               put_prev_set_next_task(rq, rq->donor, p);
>                 return p;
>         }
> 
> @@ -6157,10 +6156,8 @@ __pick_next_task(struct rq *rq, struct rq_flags *rf)
>                 p = class->pick_task(rq, rf);
>                 if (unlikely(p == RETRY_TASK))
>                         goto restart;
> -               if (p) {
> -                       put_prev_set_next_task(rq, rq->donor, p);
> +               if (p)
>                         return p;
> -               }
>         }
> 
>         BUG(); /* The idle class should always have a runnable task. */
> @@ -6257,7 +6254,7 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
>                 rq->dl_server = rq->core_dl_server;
>                 rq->core_pick = NULL;
>                 rq->core_dl_server = NULL;
> -               goto out_set_next;
> +               goto out_return_next;
>         }
> 
>         prev_balance(rq, rf);
> @@ -6311,7 +6308,7 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
>                          */
>                         WARN_ON_ONCE(fi_before);
>                         task_vruntime_update(rq, next, false);
> -                       goto out_set_next;
> +                       goto out_return_next;
>                 }
>         }
> 
> @@ -6441,8 +6438,7 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
>                 resched_curr(rq_i);
>         }
> 
> -out_set_next:
> -       put_prev_set_next_task(rq, rq->donor, next);
> +out_return_next:
>         if (rq->core->core_forceidle_count && next == rq->idle)
>                 queue_core_balance(rq);
> 
> @@ -6755,15 +6751,23 @@ static void proxy_deactivate(struct rq *rq, struct task_struct *donor)
>         WARN_ON_ONCE(state == TASK_RUNNING);
>         WARN_ON_ONCE(donor->blocked_on);
>         /*
> -        * Because we got donor from pick_next_task(), it is *crucial*
> -        * that we call proxy_resched_idle() before we deactivate it.
> -        * As once we deactivate donor, donor->on_rq is set to zero,
> -        * which allows ttwu() to immediately try to wake the task on
> -        * another rq. So we cannot use *any* references to donor
> -        * after that point. So things like cfs_rq->curr or rq->donor
> -        * need to be changed from next *before* we deactivate.
> +        * A proxy candidate is not necessarily the committed rq->donor.
> +        * pick_next_task() only selected it; the class current state is
> +        * updated later, after proxy-chain resolution.
> +        *
> +        * If @donor is still the committed donor, the rq and the scheduling
> +        * class may hold current references to it, such as rq->donor or
> +        * cfs_rq->curr/h_curr. Drop those references before block_task(),
> +        * because block_task() clears donor->on_rq and a concurrent wakeup
> +        * may then move the task elsewhere.
> +        *
> +        * If @donor is only an uncommitted proxy candidate, it is still a
> +        * queued task, not the class current task, so it can be blocked
> +        * directly.
>          */
> -       proxy_resched_idle(rq);
> +       if (donor == rq->donor)
> +               proxy_resched_idle(rq);
> +
>         block_task(rq, donor, state);
>  }
> 
> @@ -6816,15 +6820,17 @@ static void proxy_migrate_task(struct rq *rq, struct rq_flags *rf,
>         lockdep_assert_rq_held(rq);
>         WARN_ON(p == rq->curr);
>         /*
> -        * Since we are migrating a blocked donor, it could be rq->donor,
> -        * and we want to make sure there aren't any references from this
> -        * rq to it before we drop the lock. This avoids another cpu
> -        * jumping in and grabbing the rq lock and referencing rq->donor
> -        * or cfs_rq->curr, etc after we have migrated it to another cpu,
> -        * and before we pick_again in __schedule.
> +        * Migrating a task found in the proxy chain abandons the current proxy
> +        * pick attempt and drops this rq's lock.
> +        *
> +        * The picked proxy candidate has not necessarily been committed, so
> +        * @p is not necessarily rq->donor. Switch the currently committed
> +        * donor to idle before dropping the lock, leaving rq->donor and the
> +        * class current state in a well-defined state while the chain is
> +        * modified and @p is attached elsewhere.
>          *
> -        * So call proxy_resched_idle() to drop the rq->donor references
> -        * before we release the lock.
> +        * If @p is rq->donor, this also drops the direct rq/class-current
> +        * references to @p before it leaves this rq.
>          */
>         proxy_resched_idle(rq);
> 
> @@ -7147,11 +7153,11 @@ static void __sched notrace __schedule(int sched_mode)
>         rq->next_class = next->sched_class;
>         if (sched_proxy_exec()) {
>                 struct task_struct *prev_donor = rq->donor;
> +               struct task_struct *donor = next;
> 
> -               rq_set_donor(rq, next);
> -               next->blocked_donor = NULL;
> -               if (unlikely(next->is_blocked)) {
> -                       next = find_proxy_task(rq, next, &rf);
> +               donor->blocked_donor = NULL;
> +               if (unlikely(donor->is_blocked)) {
> +                       next = find_proxy_task(rq, donor, &rf);
>                         if (!next) {
>                                 zap_balance_callbacks(rq);

You don't need to zap any balance callbacks if you don't do
a put_prev_set_next_task(). It should always be empty here.

>                                 goto pick_again;
> @@ -7161,8 +7167,11 @@ static void __sched notrace __schedule(int sched_mode)
>                                 goto keep_resched;
>                         }
>                 }
> -               if (rq->donor == prev_donor && prev != next) {
> -                       struct task_struct *donor = rq->donor;
> +
> +               put_prev_set_next_task(rq, prev_donor, donor);
> +               rq_set_donor(rq, donor);

Since the next condition only cares for donor == prev_donor, I
think you can do this after the next block and unify the
!sched_proxy_exec() case too like:

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index f13f53bd6d9d7..c6e37005a73c4 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -7171,9 +7171,6 @@ static void __sched notrace __schedule(int sched_mode)
 			}
 		}
 
-		put_prev_set_next_task(rq, prev_donor, donor);
-		rq_set_donor(rq, donor);
-
 		if (donor == prev_donor && prev != next) {
 			/*
 			 * When transitioning like:
@@ -7190,10 +7187,9 @@ static void __sched notrace __schedule(int sched_mode)
 			donor->sched_class->put_prev_task(rq, donor, donor);
 			donor->sched_class->set_next_task(rq, donor, true);
 		}
-	} else {
-		put_prev_set_next_task(rq, rq->donor, next);
-		rq_set_donor(rq, next);
 	}
+	put_prev_set_next_task(rq, rq->donor, next);
+	rq_set_donor(rq, next);
 
 picked:
 	clear_tsk_need_resched(prev);
---

> +
> +               if (donor == prev_donor && prev != next) {
>                         /*
>                          * When transitioning like:
>                          *
> @@ -7179,6 +7188,7 @@ static void __sched notrace __schedule(int sched_mode)
>                         donor->sched_class->set_next_task(rq, donor, true);
>                 }
>         } else {
> +               put_prev_set_next_task(rq, rq->donor, next);
>                 rq_set_donor(rq, next);
>         }
> 
> 
> --
> 2.34.1
> 

-- 
Thanks and Regards,
Prateek