Re: [PATCH v25 9/9] sched: Handle blocked-waiter migration (and return migration)

K Prateek Nayak
2026-03-15T17:38:36Z

Hello John,

On 3/13/2026 8:00 AM, John Stultz wrote:
> diff --git a/kernel/sched/core.c b/kernel/sched/core.c
> index af497b8c72dce..fe20204cf51cc 100644
> --- a/kernel/sched/core.c
> +++ b/kernel/sched/core.c
> @@ -3643,6 +3643,23 @@ void update_rq_avg_idle(struct rq *rq)
> rq->idle_stamp = 0;
> }
>
> +#ifdef CONFIG_SCHED_PROXY_EXEC
> +static inline void proxy_set_task_cpu(struct task_struct *p, int cpu)
> +{
> + unsigned int wake_cpu;
> +
> + /*
> + * Since we are enqueuing a blocked task on a cpu it may
> + * not be able to run on, preserve wake_cpu when we
> + * __set_task_cpu so we can return the task to where it
> + * was previously runnable.
> + */
> + wake_cpu = p->wake_cpu;
> + __set_task_cpu(p, cpu);
> + p->wake_cpu = wake_cpu;
> +}
> +#endif /* CONFIG_SCHED_PROXY_EXEC */
> +
> static void
> ttwu_do_activate(struct rq *rq, struct task_struct *p, int wake_flags,
> struct rq_flags *rf)
> @@ -4242,13 +4259,6 @@ int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags)
> ttwu_queue(p, cpu, wake_flags);
> }
> out:
> - /*
> - * For now, if we've been woken up, clear the task->blocked_on
> - * regardless if it was set to a mutex or PROXY_WAKING so the
> - * task can run. We will need to be more careful later when
> - * properly handling proxy migration
> - */
> - clear_task_blocked_on(p, NULL);

So, for this bit, there are mutex variants that are interruptible and
killable which probably benefits from clearing the blocked_on
relation.

For potential proxy task that are still queued, we'll hit the
ttwu_runnable() path and resched out of there so it makes sense to
mark them as PROXY_WAKING so schedule() can return migrate them, they
run and hit the signal_pending_state() check in __mutex_lock_common()
loop, and return -EINTR.

Otherwise, if they need a full wakeup, they may be blocked on a
sleeping owner, in which case it is beneficial to clear blocked_on, do
a full wakeup. and let them run to evaluate the pending signal.

ttwu_state_match() should filter out any spurious signals. Thoughts?

> if (success)
> ttwu_stat(p, task_cpu(p), wake_flags);
>
> @@ -6575,7 +6585,7 @@ static inline struct task_struct *proxy_resched_idle(struct rq *rq)
> return rq->idle;
> }
>
> -static bool __proxy_deactivate(struct rq *rq, struct task_struct *donor)
> +static bool proxy_deactivate(struct rq *rq, struct task_struct *donor)
> {
> unsigned long state = READ_ONCE(donor->__state);
>
> @@ -6595,17 +6605,135 @@ static bool __proxy_deactivate(struct rq *rq, struct task_struct *donor)
> return try_to_block_task(rq, donor, &state, true);
> }
>
> -static struct task_struct *proxy_deactivate(struct rq *rq, struct task_struct *donor)
> +/*
> + * If the blocked-on relationship crosses CPUs, migrate @p to the
> + * owner's CPU.
> + *
> + * This is because we must respect the CPU affinity of execution
> + * contexts (owner) but we can ignore affinity for scheduling
> + * contexts (@p). So we have to move scheduling contexts towards
> + * potential execution contexts.
> + *
> + * Note: The owner can disappear, but simply migrate to @target_cpu
> + * and leave that CPU to sort things out.
> + */
> +static void proxy_migrate_task(struct rq *rq, struct rq_flags *rf,
> + struct task_struct *p, int target_cpu)
> {
> - if (!__proxy_deactivate(rq, donor)) {
> - /*
> - * XXX: For now, if deactivation failed, set donor
> - * as unblocked, as we aren't doing proxy-migrations
> - * yet (more logic will be needed then).
> - */
> - clear_task_blocked_on(donor, NULL);
> + struct rq *target_rq = cpu_rq(target_cpu);
> +
> + lockdep_assert_rq_held(rq);
> +
> + /*
> + * Since we're going to drop @rq, we have to put(@rq->donor) first,
> + * otherwise we have a reference that no longer belongs to us.
> + *
> + * Additionally, as we put_prev_task(prev) earlier, its possible that
> + * prev will migrate away as soon as we drop the rq lock, however we
> + * still have it marked as rq->curr, as we've not yet switched tasks.
> + *
> + * So call proxy_resched_idle() to let go of the references before
> + * we release the lock.
> + */
> + proxy_resched_idle(rq);
> +
> + WARN_ON(p == rq->curr);
> +
> + deactivate_task(rq, p, DEQUEUE_NOCLOCK);
> + proxy_set_task_cpu(p, target_cpu);
> +
> + /*
> + * We have to zap callbacks before unlocking the rq
> + * as another CPU may jump in and call sched_balance_rq
> + * which can trip the warning in rq_pin_lock() if we
> + * leave callbacks set.
> + */
> + zap_balance_callbacks(rq);
> + rq_unpin_lock(rq, rf);
> + raw_spin_rq_unlock(rq);
> +
> + attach_one_task(target_rq, p);
> +
> + raw_spin_rq_lock(rq);
> + rq_repin_lock(rq, rf);
> + update_rq_clock(rq);
> +}
> +
> +static void proxy_force_return(struct rq *rq, struct rq_flags *rf,
> + struct task_struct *p)
> +{
> + struct rq *this_rq, *target_rq;
> + struct rq_flags this_rf;
> + int cpu, wake_flag = WF_TTWU;
> +
> + lockdep_assert_rq_held(rq);
> + WARN_ON(p == rq->curr);
> +
> + /*
> + * We have to zap callbacks before unlocking the rq
> + * as another CPU may jump in and call sched_balance_rq
> + * which can trip the warning in rq_pin_lock() if we
> + * leave callbacks set.
> + */
> + zap_balance_callbacks(rq);
> + rq_unpin_lock(rq, rf);
> + raw_spin_rq_unlock(rq);
> +
> + /*
> + * We drop the rq lock, and re-grab task_rq_lock to get
> + * the pi_lock (needed for select_task_rq) as well.
> + */
> + this_rq = task_rq_lock(p, &this_rf);
> +
> + /*
> + * Since we let go of the rq lock, the task may have been
> + * woken or migrated to another rq before we got the
> + * task_rq_lock. So re-check we're on the same RQ. If
> + * not, the task has already been migrated and that CPU
> + * will handle any futher migrations.
> + */
> + if (this_rq != rq)
> + goto err_out;
> +
> + /* Similarly, if we've been dequeued, someone else will wake us */
> + if (!task_on_rq_queued(p))
> + goto err_out;
> +
> + /*
> + * Since we should only be calling here from __schedule()
> + * -> find_proxy_task(), no one else should have
> + * assigned current out from under us. But check and warn
> + * if we see this, then bail.
> + */
> + if (task_current(this_rq, p) || task_on_cpu(this_rq, p)) {
> + WARN_ONCE(1, "%s rq: %i current/on_cpu task %s %d on_cpu: %i\n",
> + __func__, cpu_of(this_rq),
> + p->comm, p->pid, p->on_cpu);
> + goto err_out;
> }
> - return NULL;
> +
> + update_rq_clock(this_rq);
> + proxy_resched_idle(this_rq);

I still think this is too late, and only required if we are moving the
donor. Can we do this before we drop the rq_lock so that a remote
wakeup doesn't need to clear the this? (although I think we don't have
that bit in the ttwu path anymore and we rely on the schedule() bits
completely for return migration on this version - any particular
reason?).

> + deactivate_task(this_rq, p, DEQUEUE_NOCLOCK);
> + cpu = select_task_rq(p, p->wake_cpu, &wake_flag);
> + set_task_cpu(p, cpu);
> + target_rq = cpu_rq(cpu);
> + clear_task_blocked_on(p, NULL);
> + task_rq_unlock(this_rq, p, &this_rf);
> +
> + attach_one_task(target_rq, p);

I'm still having a hard time believing we cannot use wake_up_process()
but let me look more into that tomorrow when the sun rises.

> +
> + /* Finally, re-grab the origianl rq lock and return to pick-again */
> + raw_spin_rq_lock(rq);
> + rq_repin_lock(rq, rf);
> + update_rq_clock(rq);
> + return;
> +
> +err_out:
> + task_rq_unlock(this_rq, p, &this_rf);
> + raw_spin_rq_lock(rq);
> + rq_repin_lock(rq, rf);
> + update_rq_clock(rq);
> }
>
> /*
--
Thanks and Regards,
Prateek