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

Peter Zijlstra
2026-03-19T12:50:15Z

On Fri, Mar 13, 2026 at 02:30:10AM +0000, John Stultz wrote:
> +/*
> + * 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)
> {
> + 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);

This comment confuses the heck out of me. It seems to imply we need to
schedule before dropping rq->lock.

> +
> + 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.
> + */

It might be good to explain where these callbacks come from.

> + 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.
> + */

idem

> + zap_balance_callbacks(rq);
> + rq_unpin_lock(rq, rf);
> + raw_spin_rq_unlock(rq);

This is in fact the very same sequence as above.

> +
> + /*
> + * 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);
> + 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);
> +
> + /* 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);
> }

Hurm... how about something like so?

---
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6662,6 +6662,28 @@ static bool proxy_deactivate(struct rq *
return try_to_block_task(rq, donor, &state, true);
}

+static inline void proxy_release_rq_lock(struct rq *rq, struct rq_flags *rf)
+ __releases(__rq_lockp(rq))
+{
+ /*
+ * 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);
+}
+
+static inline void proxy_reacquire_rq_lock(struct rq *rq, struct rq_flags *rf)
+ __acquires(__rq_lockp(rq))
+{
+ raw_spin_rq_lock(rq);
+ rq_repin_lock(rq, rf);
+ update_rq_clock(rq);
+}
+
/*
* If the blocked-on relationship crosses CPUs, migrate @p to the
* owner's CPU.
@@ -6676,6 +6698,7 @@ static bool proxy_deactivate(struct rq *
*/
static void proxy_migrate_task(struct rq *rq, struct rq_flags *rf,
struct task_struct *p, int target_cpu)
+ __must_hold(__rq_lockp(rq))
{
struct rq *target_rq = cpu_rq(target_cpu);

@@ -6699,98 +6722,72 @@ static void proxy_migrate_task(struct rq
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);
+ proxy_release_rq_lock(rq, rf);

attach_one_task(target_rq, p);

- raw_spin_rq_lock(rq);
- rq_repin_lock(rq, rf);
- update_rq_clock(rq);
+ proxy_reacquire_rq_lock(rq, rf);
}

static void proxy_force_return(struct rq *rq, struct rq_flags *rf,
struct task_struct *p)
+ __must_hold(__rq_lockp(rq))
{
- struct rq *this_rq, *target_rq;
- struct rq_flags this_rf;
+ struct rq *task_rq, *target_rq = NULL;
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);
+ proxy_release_rq_lock(rq, rf);

/*
* 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);
+ scoped_guard (task_rq_lock, p) {
+ task_rq = scope.rq;

- /*
- * 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;
+ /*
+ * 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 (task_rq != rq)
+ break;
+
+ /* Similarly, if we've been dequeued, someone else will wake us */
+ if (!task_on_rq_queued(p))
+ break;
+
+ /*
+ * 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(task_rq, p) || task_on_cpu(task_rq, p)) {
+ WARN_ONCE(1, "%s rq: %i current/on_cpu task %s %d on_cpu: %i\n",
+ __func__, cpu_of(task_rq),
+ p->comm, p->pid, p->on_cpu);
+ break;
+ }
+
+ update_rq_clock(task_rq);
+ proxy_resched_idle(task_rq);
+ deactivate_task(task_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);
}

- update_rq_clock(this_rq);
- proxy_resched_idle(this_rq);
- 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);
+ if (target_rq)
+ attach_one_task(target_rq, p);

- attach_one_task(target_rq, p);
-
- /* 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);
+ proxy_reacquire_rq_lock(rq, rf);
}

/*
@@ -6811,6 +6808,7 @@ static void proxy_force_return(struct rq
*/
static struct task_struct *
find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
+ __must_hold(__rq_lockp(rq))
{
enum { FOUND, DEACTIVATE_DONOR, MIGRATE, NEEDS_RETURN } action = FOUND;
struct task_struct *owner = NULL;