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

[Xukai Wang]

Hi Prateek,
Thanks for taking a look.

On 2026/7/10 15:23, K Prateek Nayak wrote:
> 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.
I agree that, if the rq state does not materially change while rq_lock is
held, the retry can converge back to the same donor. I did not intend to
make
a performance claim with the RFC. The counters were meant to characterize
what happens in the proxy path after a blocked donor has already been
committed.

I added some temporary instrumentation on the baseline kernel to check
the retry behavior after find_proxy_task() returns NULL. The test was
run on top of tip/sched/core at commit
19b7bdc3a1550ab2550427c33395bec7caeaf72d, with only temporary debug
instrumentation added. A small misc driver was used to repeatedly takes
a kernel struct mutex and busy-waits while holding it. A userspace
stress program creates one low-priority holder and multiple
higher-priority waiters pinned to selected CPUs. The specific run below
used:

|holder CPU: 1 waiter CPU base: 0 waiter CPU span: 4 waiters: 8
duration: 60s holder hold time: 3000 us waiter hold time: 1 us waiter
policy: SCHED_RR prio 50|

The temporary counters mean:

|spec_commit_blocked: pick_next_task() selected a blocked donor, and in
the baseline code that donor had already gone through
put_prev_set_next_task() before proxy-chain resolution.
spec_commit_then_null: the blocked donor had already been committed, but
find_proxy_task() returned NULL and __schedule() retried.
spec_commit_then_idle: the blocked donor had already been committed, but
find_proxy_task() returned rq->idle. spec_commit_then_success: the
blocked donor had already been committed, and find_proxy_task()
successfully found a task to run. null_retry_same_donor: after
find_proxy_task() returned NULL, the next pick selected the same donor
again. null_retry_diff_donor: after find_proxy_task() returned NULL, the
next pick selected a different non-idle donor. null_retry_to_idle: after
find_proxy_task() returned NULL, the next pick selected idle.|

In that run, I got:

|find_call 5407 find_success 523 find_idle 1660 find_null 3224
spec_commit_blocked 5407 spec_commit_then_null 3224
spec_commit_then_idle 1660 spec_commit_then_success 523
null_deactivate_no_mutex 6 null_deactivate_no_owner 125
null_deactivate_owner_not_runnable 962 null_migrate 2131
null_retry_same_donor 0 null_retry_diff_donor 2258 null_retry_to_idle 966|

So in this workload, the selected blocked donor had already been committed
before proxy-chain resolution, but most resolutions ended in NULL or idle.
For the NULL cases, the retry did not pick the same donor again. Most of
those NULL cases came from deactivate/migrate paths, so this appears to be
the kind of case where proxy state changes before the retry.

>> 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@xxxxxxxxxxxx>
>> ---
>>  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);

Based on your comments above, I am thinking of restructuring the patch
like this.

First, move the final put_prev_set_next_task()/rq_set_donor() out of the
proxy/non-proxy branches. I think the common commit target still needs to be
the donor rather than next, because after find_proxy_task(), next may be the
proxy owner that will actually run, while donor is still the selected
scheduling donor/accounting context:

|struct task_struct *donor; ...|

pick_again:
assert_balance_callbacks_empty(rq);
next = pick_next_task(rq, &rf);
rq->next_class = next->sched_class;
donor = next;

|if (sched_proxy_exec()) { donor->blocked_donor = NULL; if
(unlikely(donor->is_blocked)) { next = find_proxy_task(rq, donor, &rf);
if (!next) goto pick_again; if (next == rq->idle) goto keep_resched; }
if (donor == rq->donor && prev != next) {
donor->sched_class->put_prev_task(rq, donor, donor);
donor->sched_class->set_next_task(rq, donor, true); } }
put_prev_set_next_task(rq, rq->donor, donor); rq_set_donor(rq, donor);|

Second, for zap_balance_callbacks(), I noticed that some NULL/idle paths can
still go through proxy_resched_idle(), which performs the idle put/set
transition before returning. So instead of keeping
zap_balance_callbacks() at
the outer !next / next == rq->idle sites, I am thinking of moving it into
proxy_resched_idle(), next to the idle commit itself:

|static inline struct task_struct *proxy_resched_idle(struct rq *rq) {
put_prev_set_next_task(rq, rq->donor, rq->idle); rq->next_class =
&idle_sched_class; rq_set_donor(rq, rq->idle);
set_tsk_need_resched(rq->idle); zap_balance_callbacks(rq); return
rq->idle; }|

Then the outer NULL/idle paths no longer need to zap unconditionally. If no
put/set happened, there should be no callbacks to clear; if the path went
through proxy_resched_idle(), the callbacks are cleared at the place
that did
the idle commit.

Does this look like the right direction?

If useful, I can also share the temporary instrumentation and the exact
proxy_mutex_test setup used to reproduce the numbers above. I can further
look into whether this has a measurable performance effect, but I would keep
that separate from the correctness/commit-boundary question unless you think
it is necessary for this change.

>> +
>> +               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,
Xukai