Simple Donor Migration for Proxy Execution

kernel/sched/deadline.c +4 -2

index d08b004293234..4e746f4de6529 100644
@@ -2801,12 +2801,24 @@ static int find_later_rq(struct task_struct *task)
 
 static struct task_struct *pick_next_pushable_dl_task(struct rq *rq)
 {
-	struct task_struct *p;
+	struct task_struct *p = NULL;
+	struct rb_node *next_node;
 
 	if (!has_pushable_dl_tasks(rq))
 		return NULL;
 
-	p = __node_2_pdl(rb_first_cached(&rq->dl.pushable_dl_tasks_root));
+	next_node = rb_first_cached(&rq->dl.pushable_dl_tasks_root);
+	while (next_node) {

Hello,

On 13/03/26 02:30, John Stultz wrote:

...

Quoted 15 lines

kernel/sched/deadline.c +2 -1

index d08b004293234..4e746f4de6529 100644
@@ -2801,12 +2801,24 @@ static int find_later_rq(struct task_struct *task)
 
 static struct task_struct *pick_next_pushable_dl_task(struct rq *rq)
 {
-     struct task_struct *p;
+     struct task_struct *p = NULL;
+     struct rb_node *next_node;
 
      if (!has_pushable_dl_tasks(rq))
              return NULL;
 

kernel/sched/deadline.c static int find_later_rq(struct task_struct *task) · line 2821

Quoted 13 lines

kernel/sched/deadline.c +12 -1

-     p = __node_2_pdl(rb_first_cached(&rq->dl.pushable_dl_tasks_root));
+     next_node = rb_first_cached(&rq->dl.pushable_dl_tasks_root);
+     while (next_node) {
+             p = __node_2_pdl(next_node);
+             /* make sure task isn't on_cpu (possible with proxy-exec) */
+             if (!task_on_cpu(rq, p))
+                     break;
+
+             next_node = rb_next(next_node);
+     }
+
+     if (!p)
+             return NULL;

Can't this return an on_cpu task if we hit the corner case where all pushable tasks are on_cpu?

kernel/sched/core.c +0 -17

index b7f77c165a6e0..d86d648a75a4b 100644
@@ -6702,23 +6702,6 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
 }
 #endif /* SCHED_PROXY_EXEC */
 
-static inline void proxy_tag_curr(struct rq *rq, struct task_struct *owner)
-{
-	if (!sched_proxy_exec())
-		return;
-	/*
-	 * pick_next_task() calls set_next_task() on the chosen task
-	 * at some point, which ensures it is not push/pullable.
-	 * However, the chosen/donor task *and* the mutex owner form an
-	 * atomic pair wrt push/pull.
-	 *
-	 * Make sure owner we run is not pushable. Unfortunately we can
-	 * only deal with that by means of a dequeue/enqueue cycle. :-/
-	 */
-	dequeue_task(rq, owner, DEQUEUE_NOCLOCK | DEQUEUE_SAVE);
-	enqueue_task(rq, owner, ENQUEUE_NOCLOCK | ENQUEUE_RESTORE);
-}
-
 /*
  * __schedule() is the main scheduler function.
  *

kernel/sched/core.c +0 -4

 		/*
 		 * The membarrier system call requires each architecture
 		 * to have a full memory barrier after updating

@@ -6907,10 +6887,6 @@ static void __sched notrace __schedule(int sched_mode)
 		/* Also unlocks the rq: */
 		rq = context_switch(rq, prev, next, &rf);
 	} else {
-		/* In case next was already curr but just got blocked_donor */
-		if (!task_current_donor(rq, next))
-			proxy_tag_curr(rq, next);
-
 		rq_unpin_lock(rq, &rf);
 		__balance_callbacks(rq, NULL);
 		raw_spin_rq_unlock_irq(rq);

Back to my concern with the queuing of the balance_callback, and the deadline and RT folks can keep me honest here, consider the following:

    CPU0
    ====

======> Task A (prio: 80)
...

mutex_lock(Mutex0)
... /* Executing critical section. */

    =====> Interrupt: Wakes up Task B (prio: 50); B->blocked_on = Mutex0;
      resched_curr()
    <===== Interrupt return
  preempt_schedule_irq()
    schedule()
      put_prev_set_next_Task(A, B)
      rq->donor = B
      if (task_is_blocked(B)
        next = find_proxy_task() /* Return Task A */
      rq->curr = A
      queue_balance_callback()
    do_balance_callbacks()
      /* Finds A as task_on_cpu(); Does nothing. */

... /* returns from schedule */
... /* continues with critical section */

mutex_unlock(Mutex0)

    mutex_handoff(B /* Task B */)
    preempt_disable()
      try_to_wake_up()
        resched_curr()
    preempt_enable()
      preempt_schedule()
        proxy_force_return()
          /* Returns to same CPU */

        /*
         * put_prev_set_next_task() is skipped since
         * rq->donor context is same. no balance
         * callbacks are queued. Task A still on the
         * push list.
         */
        rq->donor = B
        rq->curr = B

=======> sched_out: Task A

!!! No balance callback; Task A still on push list. !!!

<======= sched_in: Task B

So what I'm getting to is, if we find that rq->donor has not changed with sched_proxy_exec() but rq->curr has changed during schedule(), we should forcefully do a:

  prev->sched_class->put_prev_task(rq, rq->donor, rq->donor /* or rq->idle / NULL ? */);
  next->sched_class->set_next_task(rq, rq->donor, true /* to queue balance callback. */);

That way, when we do set_nex_task(), we see if we potentially have tasks in the push list and queue a balance callback since the task_on_cpu() condition may no longer apply to the tasks left behind on the list.

Thoughts?

Quoted 50 lines · K Prateek Nayak

Back to my concern with the queuing of the balance_callback, and the deadline and RT folks can keep me honest here, consider the following:

    CPU0
    ====

======> Task A (prio: 80)
...

mutex_lock(Mutex0)
... /* Executing critical section. */

    =====> Interrupt: Wakes up Task B (prio: 50); B->blocked_on = Mutex0;
      resched_curr()
    <===== Interrupt return
  preempt_schedule_irq()
    schedule()
      put_prev_set_next_Task(A, B)
      rq->donor = B
      if (task_is_blocked(B)
        next = find_proxy_task() /* Return Task A */
      rq->curr = A
      queue_balance_callback()
    do_balance_callbacks()
      /* Finds A as task_on_cpu(); Does nothing. */

... /* returns from schedule */
... /* continues with critical section */

mutex_unlock(Mutex0)

    mutex_handoff(B /* Task B */)
    preempt_disable()
      try_to_wake_up()
        resched_curr()
    preempt_enable()
      preempt_schedule()
        proxy_force_return()
          /* Returns to same CPU */

        /*
         * put_prev_set_next_task() is skipped since
         * rq->donor context is same. no balance
         * callbacks are queued. Task A still on the
         * push list.
         */
        rq->donor = B
        rq->curr = B

=======> sched_out: Task A

Quoted 3 lines

!!! No balance callback; Task A still on push list. !!!

<======= sched_in: Task B

Hrm. I'm feeling like I'm a little lost here, specifically after proxy_force_return(), since it doesn't exist yet at this point in the patch series.

But assuming we're at the "Handle blocked-waiter migration" point in the series, I'd think it would be something like:

rq->donor= B
rq->curr = A
<< task A >>
mutex_unlock(Mutex0)
mutex_handoff(B /* Task B */)
preempt_disable()

     try_to_wake_up()
       resched_curr()
    preempt_enable()
      preempt_schedule()
        __schedule()
          find_proxy_task()
             proxy_force_return()
             return NULL
        pick_again:
          next = pick_next_task()
                      __pick_next_task() /* Returns B */
          rq->donor =B
          rq->curr = B
          context_switch()

<<switch to B >>
finish_task_switch()
finish_lock_switch()
__balance_callbacks()

Your point "put_prev_set_next_task() is skipped since rq->donor context is same" wasn't initially obvious to me, as the fair scheduler does have a (p == prev) check, but it doesn't enqueue balance callbacks. And for RT/DL/SCX we should be using the pick_task() method, which calls put_prev_set_next_task() in __pick_next_task(). But indeed, *inside* of put_prev_set_next_task() we return early if (next == prev).

So I see your concern and agree.

Quoted 7 lines

So what I'm getting to is, if we find that rq->donor has not changed with sched_proxy_exec() but rq->curr has changed during schedule(), we should forcefully do a:

  prev->sched_class->put_prev_task(rq, rq->donor, rq->donor /* or rq->idle / NULL ? */);
  next->sched_class->set_next_task(rq, rq->donor, true /* to queue balance callback. */);

Quoted 6 lines

That way, when we do set_nex_task(), we see if we potentially have tasks in the push list and queue a balance callback since the task_on_cpu() condition may no longer apply to the tasks left behind on the list.

Thoughts?

Yeah. I wonder if we can express this inside of put_prev_set_next_task(). Reworking the shortcut to maybe:

    if (next == prev && next != rq->curr)

I probably need to think on this tomorrow, as I suspect the above has some holes, but it seems like it would catch the cases that would matter

(maybe the issue is it catches too much - we'd probably also trip it if we A boosted B, and then we hit schedule and again chose A to boost B, which we probably could have skipped).

I guess adding a new helper function to manually do the put_prev/set_next could be added to the top level __schedule() logic in the (prev != next) case, though we'll have to preserve the prev_donor on the stack probably.

Quoted 5 lines · John Stultz

I guess adding a new helper function to manually do the put_prev/set_next could be added to the top level __schedule() logic in the (prev != next) case, though we'll have to preserve the prev_donor on the stack probably.

That seems like the best option to me too.

Also, deadline, RT, fair, and idle don't really care about the "next" argument of put_prev_task() and the only one that does care is put_prev_task_scx() to call switch_class() callback so putting it as either NULL or "rq->donor" should be safe.

Quoted 9 lines · K Prateek Nayak

Quoted 5 lines · John Stultz

I guess adding a new helper function to manually do the put_prev/set_next could be added to the top level __schedule() logic in the (prev != next) case, though we'll have to preserve the prev_donor on the stack probably.

That seems like the best option to me too.

Quoted 4 lines

Also, deadline, RT, fair, and idle don't really care about the "next" argument of put_prev_task() and the only one that does care is put_prev_task_scx() to call switch_class() callback so putting it as either NULL or "rq->donor" should be safe.

Ack.
Here's the change I'm testing tonight (against 6.18):
https://github.com/johnstultz-work/linux-dev/commit/0cc72a4923143f496e33711cbcc1afdf6d861ca6

Feel free to suggest a better name for the helper function. It feels a little clunky (and sort of sad right after getting rid of the clunky proxy_tag_curr(), to re-add something so similar).

Hello John,

Quoted 10 lines · John Stultz

Quoted 9 lines · K Prateek Nayak

Quoted 5 lines · John Stultz

I guess adding a new helper function to manually do the put_prev/set_next could be added to the top level __schedule() logic in the (prev != next) case, though we'll have to preserve the prev_donor on the stack probably.

That seems like the best option to me too.

Quoted 8 lines

Quoted 4 lines

Also, deadline, RT, fair, and idle don't really care about the "next" argument of put_prev_task() and the only one that does care is put_prev_task_scx() to call switch_class() callback so putting it as either NULL or "rq->donor" should be safe.

Ack.
Here's the change I'm testing tonight (against 6.18):
https://github.com/johnstultz-work/linux-dev/commit/0cc72a4923143f496e33711cbcc1afdf6d861ca6

Thanks a ton for pushing out to WIP! Only nit. would be a:

s/rq->donor/prev_donor/

on the lines with {put_prev/set_next}_task() to save on the additional dereference since they are both the same (but maybe the complier figures that out on its own?)

Also would a sched_proxy_exec() check within that function make sense to skip evaluation of that branch entirely when proxy exec is disabled via cmdline?

Quoted 3 lines

Feel free to suggest a better name for the helper function. It feels a little clunky (and sort of sad right after getting rid of the clunky proxy_tag_curr(), to re-add something so similar).

I know and I'm sorry about that (T_T)

Quoted 7 lines · K Prateek Nayak

Quoted 5 lines · John Stultz

I guess adding a new helper function to manually do the put_prev/set_next could be added to the top level __schedule() logic in the (prev != next) case, though we'll have to preserve the prev_donor on the stack probably.

Quoted 6 lines

That seems like the best option to me too.

Also, deadline, RT, fair, and idle don't really care about the "next" argument of put_prev_task() and the only one that does care is put_prev_task_scx() to call switch_class() callback so putting it as either NULL or "rq->donor" should be safe.

Ack.
Here's the change I'm testing tonight (against 6.18):
https://github.com/johnstultz-work/linux-dev/commit/0cc72a4923143f496e33711cbcc1afdf6d861ca6

Feel free to suggest a better name for the helper function. It feels a little clunky (and sort of sad right after getting rid of the clunky proxy_tag_curr(), to re-add something so similar).

Quoted 15 lines · John Stultz

Quoted 7 lines · K Prateek Nayak

Quoted 5 lines · John Stultz

I guess adding a new helper function to manually do the put_prev/set_next could be added to the top level __schedule() logic in the (prev != next) case, though we'll have to preserve the prev_donor on the stack probably.

Quoted 6 lines

That seems like the best option to me too.

Also, deadline, RT, fair, and idle don't really care about the "next" argument of put_prev_task() and the only one that does care is put_prev_task_scx() to call switch_class() callback so putting it as either NULL or "rq->donor" should be safe.

Quoted 7 lines

Ack.
Here's the change I'm testing tonight (against 6.18):
https://github.com/johnstultz-work/linux-dev/commit/0cc72a4923143f496e33711cbcc1afdf6d861ca6

Feel free to suggest a better name for the helper function. It feels a little clunky (and sort of sad right after getting rid of the clunky proxy_tag_curr(), to re-add something so similar).

Does this capture it?

kernel/sched/core.c +6 -1

@@ -7100,9 +7103,11 @@ static void __sched notrace __schedule(i
 pick_again:
 	assert_balance_callbacks_empty(rq);
 	next = pick_next_task(rq, rq->donor, &rf);
-	rq_set_donor(rq, next);
 	rq->next_class = next->sched_class;
 	if (sched_proxy_exec()) {
+		struct task_struct *prev_donor = rq->donor;
+
+		rq_set_donor(rq, next);
 		if (unlikely(next->blocked_on)) {
 			next = find_proxy_task(rq, next, &rf);
 			if (!next) {

@@ -7114,6 +7119,24 @@ static void __sched notrace __schedule(i
 				goto keep_resched;
 			}
 		}
+
+		/*
+		 * When transitioning like:

kernel/sched/core.c +5 -0

+			next->sched_class->put_prev_task(rq, next, next);
+			next->sched_class->set_next_task(rq, next, true);
+		}
+	} else {
+		rq_set_donor(rq, next);
 	}
 picked:
 	clear_tsk_need_resched(prev);

kernel/sched/core.c static void __sched notrace __schedule(i · line 7128

Quoted 11 lines

kernel/sched/core.c +7 -0

@@ -7114,6 +7119,24 @@ static void __sched notrace __schedule(i
 				goto keep_resched;
 			}
 		}
+
+		/*
+		 * When transitioning like:
+		 *
+		 *	   prev		next
+		 * donor:    B		  B
+		 * curr:     A		  B

kernel/sched/core.c static void __sched notrace __schedule(i · line 7134

Quoted 6 lines

kernel/sched/core.c +6 -0

+		 *
+		 * 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.
+		 */
+		if (next == rq->donor && next == prev_donor) {

	&& next != prev

Quoted 8 lines

kernel/sched/core.c +5 -0

+			next->sched_class->put_prev_task(rq, next, next);
+			next->sched_class->set_next_task(rq, next, true);
+		}
+	} else {
+		rq_set_donor(rq, next);
 	}
 picked:
 	clear_tsk_need_resched(prev);

Quoted 19 lines · John Stultz

Quoted 7 lines · K Prateek Nayak

Quoted 5 lines · John Stultz

I guess adding a new helper function to manually do the put_prev/set_next could be added to the top level __schedule() logic in the (prev != next) case, though we'll have to preserve the prev_donor on the stack probably.

Quoted 6 lines

That seems like the best option to me too.

Also, deadline, RT, fair, and idle don't really care about the "next" argument of put_prev_task() and the only one that does care is put_prev_task_scx() to call switch_class() callback so putting it as either NULL or "rq->donor" should be safe.

Ack.
Here's the change I'm testing tonight (against 6.18):
https://github.com/johnstultz-work/linux-dev/commit/0cc72a4923143f496e33711cbcc1afdf6d861ca6

Quoted 3 lines

Feel free to suggest a better name for the helper function. It feels a little clunky (and sort of sad right after getting rid of the clunky proxy_tag_curr(), to re-add something so similar).

Does this capture it?

Yeah, this looks equivalent to what I had in my patch linked above. I'll go ahead and take your version though.

The only tweak I might consider is setting the prev_donor along with prev, so we don't have to move the rq_set_donor() call to be on both sides of the sched_proxy_exec() conditional.

Feels much simpler to read, and I'd hope the compiler would optimize the extra assignment away if it were unused at compile time. But maybe there is a concern you have I'm missing?

Quoted 10 lines · John Stultz

Yeah, this looks equivalent to what I had in my patch linked above. I'll go ahead and take your version though.

The only tweak I might consider is setting the prev_donor along with prev, so we don't have to move the rq_set_donor() call to be on both sides of the sched_proxy_exec() conditional.

Feels much simpler to read, and I'd hope the compiler would optimize the extra assignment away if it were unused at compile time. But maybe there is a concern you have I'm missing?

So my goal was to capture as much of the PE specific stuff inside that sched_proxy_exec() branch.

Having to push that rq_set_donor() into the else branch was a little unfortunate, but keeping it all inside that branch makes that we can easily break that out into its own function.

Also, deadline, RT, fair, and idle don't really care about the "next" argument of put_prev_task() and the only one that does care is put_prev_task_scx() to call switch_class() callback so putting it as either NULL or "rq->donor" should be safe.

Quoted 4 lines · K Prateek Nayak

Also, deadline, RT, fair, and idle don't really care about the "next" argument of put_prev_task() and the only one that does care is put_prev_task_scx() to call switch_class() callback so putting it as either NULL or "rq->donor" should be safe.

https://lkml.kernel.org/r/[email protected]

Makes fair care about the @next argument to put_prev_task_fair().

Peter noted: Compilers are really bad (as in they utterly refuse) optimizing (even when marked with __pure) the static branch things, and will happily emit multiple identical in a row.

So pull out the one obvious sched_proxy_exec() branch in __schedule() and remove some of the 'implicit' ones in that path.

Suggested-by Peter Zijlstra <[email protected]>

Signed-off-by John Stultz <[email protected]>

kernel/sched/core.c +1 -0

+	for (p = donor; (mutex = p->blocked_on); p = owner) {
 		/*
 		 * By taking mutex->wait_lock we hold off concurrent mutex_unlock()
 		 * and ensure @owner sticks around.

kernel/sched/sched.h +5 -0

+	update_rq_clock(rq);
+	attach_task(rq, p);
+	rq_unlock(rq, &rf);
+}
+

kernel/sched/sched.h extern void deactivate_task(struct rq *rq, struct task_struct *p, int flags); · line 3031

Quoted 9 lines · John Stultz

kernel/sched/sched.h +9 -0

+/*
+ * attach_one_task() -- attaches the task returned from detach_one_task() to
+ * its new rq.
+ */
+static inline void attach_one_task(struct rq *rq, struct task_struct *p)
+{
+     struct rq_flags rf;
+
+     rq_lock(rq, &rf);

nit. We can now use guard(rq_lock)(rq) and save on needing to declare a "rf". Apart from that, feel free to include:

Quoted 13 lines · K Prateek Nayak

kernel/sched/sched.h extern void deactivate_task(struct rq *rq, struct task_struct *p, int flags); · line 3031

Quoted 9 lines · John Stultz

kernel/sched/sched.h +9 -0

+/*
+ * attach_one_task() -- attaches the task returned from detach_one_task() to
+ * its new rq.
+ */
+static inline void attach_one_task(struct rq *rq, struct task_struct *p)
+{
+     struct rq_flags rf;
+
+     rq_lock(rq, &rf);

nit. We can now use guard(rq_lock)(rq) and save on needing to declare a "rf". Apart from that, feel free to include:

I actually did this in a later patch in the full series, as it seemed more clear I wasn't modifying logic when moving the code:
https://github.com/johnstultz-work/linux-dev/commit/8d8a12278d81ce81af6b0dfd051750f4ce2ec0e5

But, given your feedback, I'll go ahead and fold that fix down to this change and add a note in the commit message.

kernel/sched/core.c +17 -0

index af497b8c72dce..fe20204cf51cc 100644
@@ -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)

kernel/sched/core.c +100 -9

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

kernel/sched/core.c +12 -0

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

Quoted 35 lines · John Stultz

kernel/sched/core.c +17 -3

index af497b8c72dce..fe20204cf51cc 100644
@@ -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

kernel/sched/core.c int try_to_wake_up(struct task_struct *p, unsigned int state, int wake_flags) · line 4251

Quoted 4 lines

kernel/sched/core.c +0 -4

-      * 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?

Quoted 52 lines

kernel/sched/core.c +51 -0

+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;
      }

kernel/sched/core.c static bool __proxy_deactivate(struct rq *rq, struct task_struct *donor) · line 6716

Quoted 4 lines

kernel/sched/core.c +3 -1

-     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?).

kernel/sched/core.c static bool __proxy_deactivate(struct rq *rq, struct task_struct *donor) · line 6724

Quoted 9 lines

kernel/sched/core.c +8 -0

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

kernel/sched/core.c +25 -7

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

Quoted 4 lines

kernel/sched/core.c +4 -0

+	deactivate_task(rq, p, DEQUEUE_NOCLOCK);
+	proxy_set_task_cpu(p, target_cpu);
+
+	/*

kernel/sched/core.c static bool __proxy_deactivate(struct rq *rq, struct task_struct *donor) · line 6655

Quoted 10 lines

kernel/sched/core.c +10 -0

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

We release rq lock between deactivate and attach (and we don't hold neither wait_lock nor blocked_lock as they are out of scope at this point). Can't something like the following happen?

- Task A: blocked on mutex M, queued on CPU 0
- Task B: owns mutex M, running on CPU 1

CPU 0 (migrating A→CPU 1) CPU 1 (B finishes critical section)
------------------------- ------------------------------------
find_proxy_task(donor=A):
owner = B, owner_cpu = 1
action = MIGRATE
// guard releases wait_lock

proxy_migrate_task(A, cpu=1):

    deactivate_task(rq0, A)
      → A->on_rq = 0

    proxy_set_task_cpu(A, 1)
      → A->cpu = 1
    raw_spin_rq_unlock(rq0)
      → RQ0 LOCK RELEASED
                                   // Task B running
                                   mutex_unlock(M):
                                     lock(&M->wait_lock)   // ← Can grab it
                                     A->blocked_on = PROXY_WAKING
                                     unlock(&M->wait_lock)
                                     wake_up_q():
                                       try_to_wake_up(A):

                                         sees A->on_rq == 0
                                         cpu = select_task_rq(A)
                                           → returns CPU 2
                                         set_task_cpu(A, 2)
                                         ttwu_queue(A, 2)
                                           → A enqueued on CPU 2
                                           → A->on_rq = 1, A->cpu = 2

attach_one_task(rq1, A):

    attach_task(rq1, A):
      WARN_ON_ONCE(task_rq(A) != rq1)
        → Fires! task_rq(A) = rq2
      activate_task(rq1, A)
        → Double-enqueue! A->on_rq already = 1

Hello Juri,

Quoted 29 lines · Juri Lelli

Quoted 4 lines

kernel/sched/core.c +4 -0

+	deactivate_task(rq, p, DEQUEUE_NOCLOCK);
+	proxy_set_task_cpu(p, target_cpu);
+
+	/*

kernel/sched/core.c static bool __proxy_deactivate(struct rq *rq, struct task_struct *donor) · line 6655

Quoted 10 lines

kernel/sched/core.c +10 -0

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

We release rq lock between deactivate and attach (and we don't hold neither wait_lock nor blocked_lock as they are out of scope at this point). Can't something like the following happen?

- Task A: blocked on mutex M, queued on CPU 0
- Task B: owns mutex M, running on CPU 1

CPU 0 (migrating A→CPU 1) CPU 1 (B finishes critical section)
------------------------- ------------------------------------
find_proxy_task(donor=A):
owner = B, owner_cpu = 1
action = MIGRATE
// guard releases wait_lock

Quoted 3 lines

proxy_migrate_task(A, cpu=1):

    deactivate_task(rq0, A)
      → A->on_rq = 0

This sets TASK_ON_RQ_MIGRATING
before dequeuing.

Add logic to handle migrating a blocked waiter to a remote cpu where the lock owner is runnable.

Additionally, as the blocked task may not be able to run on the remote cpu, add logic to handle return migration once the waiting task is given the mutex.

Because tasks may get migrated to where they cannot run, also modify the scheduling classes to avoid sched class migrations on mutex blocked tasks, leaving find_proxy_task() and related logic to do the migrations and return migrations.

This was split out from the larger proxy patch, and significantly reworked.

Credits for the original patch go to:
Peter Zijlstra (Intel) <[email protected]>
Juri Lelli <[email protected]>
Valentin Schneider <[email protected]>
Connor O'Brien <[email protected]>

Signed-off-by John Stultz <[email protected]>

(quoted diff) +20 -0

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

(quoted diff) +17 -0

+	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;

(quoted diff) +46 -1

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

Quoted 26 lines · John Stultz

(quoted diff) +25 -0

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

Review

Quoted 5 lines

(quoted diff) +5 -0

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

Quoted 6 lines

kernel/sched/core.c +6 -0

+
+     WARN_ON(p == rq->curr);
+
+     deactivate_task(rq, p, DEQUEUE_NOCLOCK);
+     proxy_set_task_cpu(p, target_cpu);
+