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

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

Hello Xukai,

On 10/12/2026 1:09 PM, Xukai Wang wrote:
> 
> On 2026/10/12 12:05, K Prateek Nayak wrote:
>> On 8/27/2026 2:48 PM, Xukai Wang wrote:
>>> @@ -7143,45 +7159,49 @@ static void __sched notrace __schedule(int sched_mode)
>>>  
>>>  pick_again:
>>>  	assert_balance_callbacks_empty(rq);
>>> -	next = pick_next_task(rq, &rf);
>>> -	rq->next_class = next->sched_class;
>>> +	donor = pick_next_task(rq, &rf);
>>> +	rq->next_class = donor->sched_class;
>>> +	next = donor;
>>>  	if (sched_proxy_exec()) {
>>> -		struct task_struct *prev_donor = rq->donor;
>>> -
>>> -		rq_set_donor(rq, next);
>>> -		next->blocked_donor = NULL;
>>> -		if (unlikely(next->is_blocked)) {
>>> -			next = find_proxy_task(rq, next, &rf);
>>> -			if (!next) {
>>> -				zap_balance_callbacks(rq);
>>> +		donor->blocked_donor = NULL;
>>> +		if (unlikely(donor->is_blocked)) {
>>> +			struct sched_dl_entity *donor_dl_server = rq->dl_server;
>>> +
>>> +			rq->dl_server = NULL;
>> I'm not able to understand this special dl_server handling.
>>
>> In case of "pick_again", pick_next_task() will clear the rq->dl_server
>> before the pick and for "keep_resched" with idle task, I don't think the
>> dl_server from pick makes any difference since we'll come back into
>> schedule right after the switch.
>>
>> For rest of find_proxy_task() we don't really care about dl_server afaict
>> unless I'm missing something. did you trip something that required this
>> change?
>>
>>> +
>>> +			next = find_proxy_task(rq, donor, &rf);
>>> +			if (!next)
>>>  				goto pick_again;
>>> -			}
>>> -			if (next == rq->idle) {
>>> -				zap_balance_callbacks(rq);
>>> +			if (next == rq->idle)
>>>  				goto keep_resched;
>>> -			}
>>> -		}
>>> -		if (rq->donor == prev_donor && prev != next) {
>>> -			struct task_struct *donor = rq->donor;
>>> -			/*
>>> -			 * When transitioning like:
>>> -			 *
>>> -			 *         prev         next
>>> -			 * donor:    B            B
>>> -			 * curr:     A          B or C
>>> -			 *
>>> -			 * then put_prev_set_next_task() will not have done
>>> -			 * anything, since B == B. However, A might have
>>> -			 * missed a RT/DL balance opportunity due to being
>>> -			 * on_cpu.
>>> -			 */
>>> -			donor->sched_class->put_prev_task(rq, donor, donor);
>>> -			donor->sched_class->set_next_task(rq, donor, true);
>>> +
>>> +			rq->dl_server = donor_dl_server;
>>>  		}
>>> -	} else {
>>> -		rq_set_donor(rq, next);
>>> +
>>>  	}
>>>  
>>> +	put_prev_set_next_task(rq, rq->donor, donor);
>>> +
>>> +	if (sched_proxy_exec() &&
>>> +	    donor == rq->donor && prev != next) {
> 
> Hi Prateek,
> 
> The case I had in mind is when pick_next_task() selects the candidate
> through a DL server, and find_proxy_task() later goes through
> proxy_resched_idle().
> 
> 1. Before this patch, the sequence is effectively:
> 
> |-> pick_next_task()
> 
> |---|-> rq->dl_server = server;
> 
> |---|-> put_prev_set_next_task(rq, rq->donor, donor);
> 
> |---|---|-> donor->dl_server = rq->dl_server;
> 
> |---|---|-> rq->dl_server = NULL;
> 
> |-> find_proxy_task(...)
> 
> |---|-> proxy_resched_idle(rq);
> 
> |---|---|->put_prev_set_next_task(rq, rq->donor, rq->idle);
> 
> |---|---|---|-> rq->idle->dl_server = rq->dl_server; /* NULL */
> 
> |---|---|---|-> rq->dl_server = NULL;
> 
> 
> So find_proxy_task() is entered with rq->dl_server already consumed by
> the donor commit, i.e. rq->dl_server == NULL.
> 
> 2. After deferring put_prev_set_next_task(), without the save/clear, the
> sequence can instead become:
> 
> |-> pick_next_task()
> 
> |---|-> rq->dl_server = server;
> 
> |-> next = find_proxy_task(...)
> 
> |---|-> proxy_resched_idle(rq);
> 
> |---|---|->put_prev_set_next_task(rq, rq->donor, rq->idle);
> 
> |---|---|---|-> rq->idle->dl_server = rq->dl_server; /* server */
> 
> |---|---|---|-> rq->dl_server = NULL;
> 
> |->if (next == rq->idle) goto keep_resched;
> 
> ...
> 
> 
> In that case the DL server selected for the uncommitted donor gets
> attached to rq->idle.

I see. I guess, proxy_reshed_idle() would be a better spot to
consolidate it rather than juggling it in __schedule().

> 
> The next pick clearing rq->dl_server does not undo that assignment,
> since the pointer has already moved to rq->idle->dl_server. It is only
> cleared later when idle becomes the prev task of another
> put_prev_set_next_task() transition.

Which should be the very next __schedule() since we retain NEED_RESCHED
except if all the tasks migrate away and we genuinely select rq->idle,
then the dl_server lingers around.

> 
> 3. That is what the save/clear/restore in v3 was intended to avoid:
> 
> pick_again:
> 
> |-> pick_next_task()
> 
> |---|-> rq->dl_server = server;
> 
> |->donor_dl_server = rq->dl_server;
> 
> |-> rq->dl_server = NULL;
> 
> |-> next = find_proxy_task(...)
> 
> |---|-> proxy_resched_idle(rq);
> 
> |---|---|->put_prev_set_next_task(rq, rq->donor, rq->idle);
> 
> |---|---|---|-> rq->idle->dl_server = rq->dl_server; /* NULL */
> 
> |---|---|---|-> rq->dl_server = NULL;
> 
> |-> if (!next) goto pick_again;
> 
> |-> if (next == rq->idle) goto keep_resched;
> 
> |-> rq->dl_server = donor_dl_server;
> 
> ...
> 
> |-> put_prev_set_next_task(rq, rq->donor, donor);
> 
> 
> If proxy resolution returns NULL or idle, the saved value is not
> restored into rq->dl_server for a donor commit that will not happen.
> 
> So the intent is to keep the DL-server selection made for the candidate
> out of rq->dl_server while proxy resolution is in progress, and restore
> it only if that donor is actually going to be committed, while keeping
> the state seen by find_proxy_task() consistent with the pre-patch path.
> 
> Do you think rq->idle->dl_server being set in the second sequence above
> is harmless, or is keeping the save/clear/restore the right way to avoid
> that state?

I assume it may take the dl_defer_idle path but I'm not super sure. I
personally prefer resetting state in proxy_resched_idle() since it is
all in one place to refer at any point.

> 
> I’ll look at the sched-ext interaction and the __schedule() structure
> separately while rebasing onto current tip.

Thank you but I'd wait for Peter and John to chime in regarding the added
code complexity given the visible benefits of this optimization isn't
too much.

-- 
Thanks and Regards,
Prateek