Re: [PATCH v23 7/9] sched: Have try_to_wake_up() handle return-migration for PROXY_WAKING case

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

Hello John,

On 10/30/2025 5:48 AM, John Stultz wrote:
> +/*
> + * Checks to see if task p has been proxy-migrated to another rq
> + * and needs to be returned. If so, we deactivate the task here
> + * so that it can be properly woken up on the p->wake_cpu
> + * (or whichever cpu select_task_rq() picks at the bottom of
> + * try_to_wake_up()
> + */
> +static inline bool proxy_needs_return(struct rq *rq, struct task_struct *p)
> +{
> +	bool ret = false;
> +
> +	if (!sched_proxy_exec())
> +		return false;
> +
> +	raw_spin_lock(&p->blocked_lock);
> +	if (p->blocked_on == PROXY_WAKING) {
> +		if (!task_current(rq, p) && p->wake_cpu != cpu_of(rq)) {
> +			if (task_current_donor(rq, p))
> +				proxy_resched_idle(rq);
> +
> +			deactivate_task(rq, p, DEQUEUE_NOCLOCK);
> +			ret = true;
> +		}
> +		__clear_task_blocked_on(p, PROXY_WAKING);
> +		resched_curr(rq);

Do we need to resched_curr() if we've simply dequeued a waiting
owner who is neither the current, nor the donor? I would have
preferred if this function was similar to ttwu_queue_cond() with
each step explaining the the intent - something like:

static inline bool
proxy_needs_return(struct rq *rq, struct task_struct *p, int wake_flags)
{
	if (!sched_proxy_exec())
		return false;

	guard(raw_spinlock)(&p->blocked_lock);

	/* Task has been woken up / is running. */
	if (p->blocked_on != PROXY_WAKING)
		return false;

	__clear_task_blocked_on(p, PROXY_WAKING);

	/* Task is running, allow schedule() to reevaluate. */
	if (task_current(rq, p)) {
		resched_curr(rq);
		return false;
	}

	/*
	 * Task belongs to the same CPU.
	 * Check if it should be run now.
	 */
	if (p->wake_cpu == cpu_of(rq)) {
		wakeup_preempt(rq, p, wake_flags);
		return false;
	}

	/*
	 * Task was migrated from a different CPU for proxy.
	 * Block the task, and do full wakeup. If the task is
	 * the donor, ensure we call put_prev_task() before
	 * proceeding
	 */
	if (task_current_donor(rq, p))
		proxy_resched_idle(rq);

	/*
	 * (ab)Use DEQUEUE_SPECIAL to ensure task is always
	 * blocked here.
	 */
	block_task(rq, p, DEQUEUE_NOCLOCK | DEQUEUE_SPECIAL);
	return true;
}

I would wait for Peter to comment before changing all this up in case
I'm terribly wrong and am missing some subtleties. Also reason why we
may want to switch to block_task() is explained below.

> +	}
> +	raw_spin_unlock(&p->blocked_lock);
> +	return ret;
> +}
> +#else /* !CONFIG_SCHED_PROXY_EXEC */
> +static bool proxy_task_runnable_but_waking(struct task_struct *p)
> +{
> +	return false;
> +}
> +
> +static inline bool proxy_needs_return(struct rq *rq, struct task_struct *p)
> +{
> +	return false;
> +}
>  #endif /* CONFIG_SCHED_PROXY_EXEC */
>  
>  static void
> @@ -3784,6 +3834,8 @@ static int ttwu_runnable(struct task_struct *p, int wake_flags)
>  		update_rq_clock(rq);
>  		if (p->se.sched_delayed)
>  			enqueue_task(rq, p, ENQUEUE_NOCLOCK | ENQUEUE_DELAYED);
> +		if (proxy_needs_return(rq, p))
> +			goto out;
>  		if (!task_on_cpu(rq, p)) {
>  			/*
>  			 * When on_rq && !on_cpu the task is preempted, see if
> @@ -3794,6 +3846,7 @@ static int ttwu_runnable(struct task_struct *p, int wake_flags)
>  		ttwu_do_wakeup(p);
>  		ret = 1;
>  	}
> +out:
>  	__task_rq_unlock(rq, &rf);
>  
>  	return ret;
> @@ -3924,6 +3977,14 @@ static inline bool ttwu_queue_cond(struct task_struct *p, int cpu)
>  		return false;
>  #endif
>  
> +	/*
> +	 * If we're PROXY_WAKING, we have deactivated on this cpu, so we should
> +	 * activate it here as well, to avoid IPI'ing a cpu that is stuck in
> +	 * task_rq_lock() spinning on p->on_rq, deadlocking that cpu.
> +	 */

Sounds like block_task() would be better than deactivate_task() above
in that case. Anything that is waiting on the task's state change takes
the pi_lock afaik and the wakeup is always done with pi_lock held so
blocking the task shouldn't cause any problems based on my reading.

> +	if (task_on_rq_migrating(p))
> +		return false;
> +
>  	/*
>  	 * Do not complicate things with the async wake_list while the CPU is
>  	 * in hotplug state.
> @@ -4181,6 +4242,8 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
>  		 *    it disabling IRQs (this allows not taking ->pi_lock).
>  		 */
>  		WARN_ON_ONCE(p->se.sched_delayed);
> +		/* If p is current, we know we can run here, so clear blocked_on */
> +		clear_task_blocked_on(p, NULL);
>  		if (!ttwu_state_match(p, state, &success))
>  			goto out;
>  
> @@ -4197,8 +4260,15 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
>  	 */
>  	scoped_guard (raw_spinlock_irqsave, &p->pi_lock) {
>  		smp_mb__after_spinlock();
> -		if (!ttwu_state_match(p, state, &success))
> -			break;
> +		if (!ttwu_state_match(p, state, &success)) {
> +			/*
> +			 * If we're already TASK_RUNNING, and PROXY_WAKING
> +			 * continue on to ttwu_runnable check to force
> +			 * proxy_needs_return evaluation
> +			 */
> +			if (!proxy_task_runnable_but_waking(p))
> +				break;
> +		}
>  
>  		trace_sched_waking(p);
>  

-- 
Thanks and Regards,
Prateek