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

[Xukai Wang]

pick_next_task() currently commits the selected task into the scheduling
class state before proxy execution has resolved whether that task should
really become the committed donor.

With proxy execution enabled, the task returned by pick_next_task() may
be blocked. __schedule() then calls find_proxy_task() to resolve the
proxy chain. If find_proxy_task() returns NULL, __schedule() retries
through pick_again, but the picked donor has already gone through
put_prev_set_next_task(). In that case the class current state was
speculatively switched to a donor which is not the final scheduling
decision.

Make pick_next_task() return only the selected candidate task. Move the
put_prev_set_next_task() call into __schedule(), next to rq_set_donor(),
so the scheduling class commit happens after proxy-chain resolution.

For proxy execution, keep the picked donor separate from the actual task
to run:

  donor = task selected by pick_next_task()
  next  = task returned by find_proxy_task(), possibly the proxy owner

Only commit the donor after proxy-chain resolution succeeds:

  put_prev_set_next_task(rq, rq->donor, donor);
  rq_set_donor(rq, donor);

If find_proxy_task() returns NULL, the candidate donor has not been
committed and __schedule() can retry without first undoing a speculative
set_next() on that donor.

A proxy candidate is not necessarily the committed rq->donor anymore.
Update proxy_deactivate() accordingly. If the task being blocked is
still the committed donor, proxy_resched_idle() is needed to drop
rq/class current references before block_task() clears ->on_rq. If
it is only an uncommitted proxy candidate, it is still queued and can
be blocked directly.

proxy_migrate_task() still switches the rq to idle before dropping the
rq lock. Migrating a task found in the proxy chain abandons the current
proxy pick attempt, and switching the committed donor to idle leaves
rq->donor and the class current state in a defined state while the chain
is modified and the task is attached elsewhere.

Move zap_balance_callbacks() into proxy_resched_idle(), next to the idle
commit. The outer NULL/idle proxy-resolution paths no longer zap
unconditionally: if no put/set was done, no callbacks should have been
generated there; if the path went through proxy_resched_idle(),
callbacks are cleared at the point that performed the idle commit.

Signed-off-by: Xukai Wang <kingxukai@xxxxxxxxxxxx>
---
 kernel/sched/core.c | 90 +++++++++++++++++++++++++++++------------------------
 1 file changed, 49 insertions(+), 41 deletions(-)

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 2e7cde033a31..f09f3f2bea13 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6146,7 +6146,6 @@ __pick_next_task(struct rq *rq, struct rq_flags *rf)
 		if (!p)
 			p = pick_task_idle(rq, rf);
 
-		put_prev_set_next_task(rq, rq->donor, p);
 		return p;
 	}
 
@@ -6157,10 +6156,8 @@ __pick_next_task(struct rq *rq, struct rq_flags *rf)
 		p = class->pick_task(rq, rf);
 		if (unlikely(p == RETRY_TASK))
 			goto restart;
-		if (p) {
-			put_prev_set_next_task(rq, rq->donor, p);
+		if (p)
 			return p;
-		}
 	}
 
 	BUG(); /* The idle class should always have a runnable task. */
@@ -6257,7 +6254,7 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
 		rq->dl_server = rq->core_dl_server;
 		rq->core_pick = NULL;
 		rq->core_dl_server = NULL;
-		goto out_set_next;
+		goto out_return_next;
 	}
 
 	prev_balance(rq, rf);
@@ -6311,7 +6308,7 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
 			 */
 			WARN_ON_ONCE(fi_before);
 			task_vruntime_update(rq, next, false);
-			goto out_set_next;
+			goto out_return_next;
 		}
 	}
 
@@ -6441,8 +6438,7 @@ pick_next_task(struct rq *rq, struct rq_flags *rf)
 		resched_curr(rq_i);
 	}
 
-out_set_next:
-	put_prev_set_next_task(rq, rq->donor, next);
+out_return_next:
 	if (rq->core->core_forceidle_count && next == rq->idle)
 		queue_core_balance(rq);
 
@@ -6745,6 +6741,14 @@ static inline struct task_struct *proxy_resched_idle(struct rq *rq)
 	rq->next_class = &idle_sched_class;
 	rq_set_donor(rq, rq->idle);
 	set_tsk_need_resched(rq->idle);
+
+	/*
+	 * This helper performs a real idle commit. Some callers return to
+	 * __schedule() without dropping rq->lock, so clear callbacks
+	 * generated by the idle commit here. Paths that later drop
+	 * rq->lock may zap again.
+	 */
+	zap_balance_callbacks(rq);
 	return rq->idle;
 }
 
@@ -6755,15 +6759,23 @@ static void proxy_deactivate(struct rq *rq, struct task_struct *donor)
 	WARN_ON_ONCE(state == TASK_RUNNING);
 	WARN_ON_ONCE(donor->blocked_on);
 	/*
-	 * Because we got donor from pick_next_task(), it is *crucial*
-	 * that we call proxy_resched_idle() before we deactivate it.
-	 * As once we deactivate donor, donor->on_rq is set to zero,
-	 * which allows ttwu() to immediately try to wake the task on
-	 * another rq. So we cannot use *any* references to donor
-	 * after that point. So things like cfs_rq->curr or rq->donor
-	 * need to be changed from next *before* we deactivate.
+	 * A proxy candidate is not necessarily the committed rq->donor.
+	 * pick_next_task() only selected it; the class current state is
+	 * updated later, after proxy-chain resolution.
+	 *
+	 * If @donor is still the committed donor, the rq and the scheduling
+	 * class may hold current references to it, such as rq->donor or
+	 * cfs_rq->curr/h_curr. Drop those references before block_task(),
+	 * because block_task() clears donor->on_rq and a concurrent wakeup
+	 * may then move the task elsewhere.
+	 *
+	 * If @donor is only an uncommitted proxy candidate, it is still a
+	 * queued task, not the class current task, so it can be blocked
+	 * directly.
 	 */
-	proxy_resched_idle(rq);
+	if (donor == rq->donor)
+		proxy_resched_idle(rq);
+
 	block_task(rq, donor, state);
 }
 
@@ -6816,15 +6828,17 @@ static void proxy_migrate_task(struct rq *rq, struct rq_flags *rf,
 	lockdep_assert_rq_held(rq);
 	WARN_ON(p == rq->curr);
 	/*
-	 * Since we are migrating a blocked donor, it could be rq->donor,
-	 * and we want to make sure there aren't any references from this
-	 * rq to it before we drop the lock. This avoids another cpu
-	 * jumping in and grabbing the rq lock and referencing rq->donor
-	 * or cfs_rq->curr, etc after we have migrated it to another cpu,
-	 * and before we pick_again in __schedule.
+	 * Migrating a task found in the proxy chain abandons the current proxy
+	 * pick attempt and drops this rq's lock.
 	 *
-	 * So call proxy_resched_idle() to drop the rq->donor references
-	 * before we release the lock.
+	 * The picked proxy candidate has not necessarily been committed, so
+	 * @p is not necessarily rq->donor. Switch the currently committed
+	 * donor to idle before dropping the lock, leaving rq->donor and the
+	 * class current state in a well-defined state while the chain is
+	 * modified and @p is attached elsewhere.
+	 *
+	 * If @p is rq->donor, this also drops the direct rq/class-current
+	 * references to @p before it leaves this rq.
 	 */
 	proxy_resched_idle(rq);
 
@@ -7057,7 +7071,7 @@ find_proxy_task(struct rq *rq, struct task_struct *donor, struct rq_flags *rf)
  */
 static void __sched notrace __schedule(int sched_mode)
 {
-	struct task_struct *prev, *next;
+	struct task_struct *prev, *next, *donor;
 	/*
 	 * On PREEMPT_RT kernel, SM_RTLOCK_WAIT is noted
 	 * as a preemption by schedule_debug() and RCU.
@@ -7145,24 +7159,18 @@ static void __sched notrace __schedule(int sched_mode)
 	assert_balance_callbacks_empty(rq);
 	next = pick_next_task(rq, &rf);
 	rq->next_class = next->sched_class;
+	donor = next;
 	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)) {
+			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;
+
+		if (donor == rq->donor && prev != next) {
 			/*
 			 * When transitioning like:
 			 *
@@ -7178,9 +7186,9 @@ static void __sched notrace __schedule(int sched_mode)
 			donor->sched_class->put_prev_task(rq, donor, donor);
 			donor->sched_class->set_next_task(rq, donor, true);
 		}
-	} else {
-		rq_set_donor(rq, next);
 	}
+	put_prev_set_next_task(rq, rq->donor, donor);
+	rq_set_donor(rq, donor);
 
 picked:
 	clear_tsk_need_resched(prev);

-- 
2.34.1