[RFC PATCH 3/5] sched_ext: Release the CPU when its running task leaves SCX

[Kuba Piecuch <jpiecuch@google.com>]

When the task running on a CPU switches from SCX to another sched_class
through sched_setscheduler() or a PI boost, the CPU is taken away from
the BPF scheduler without a context switch. sched_change_begin() calls
put_prev_task() with a NULL @next, so put_prev_task_scx() can't tell
where the task is going, and the new class's set_next_task() doesn't
notify SCX. The BPF scheduler is never told that the CPU has been
released.

Release the CPU from switched_from_scx() if @p is the CPU's scheduling
context. Every class @p can switch to is above SCX, so the CPU is no
longer available. As the new class isn't known yet, report the new
SCX_CPU_PREEMPT_CLASS_CHANGE reason, with @args->task pointing to @p.
The release happens before ops.disable() for @p.

switched_from_scx() is the only reliable place for this:
put_prev_task_scx() with a NULL @next is also called for changes which
keep @p in SCX, and the wakeup_preempt() call at the end of
sched_change_end() is skipped if rq->next_class has already been raised.

The reverse direction needs nothing new: when a running task switches to
SCX, sched_change_end() calls set_next_task_scx(), which acquires the
CPU.

sched_setscheduler() and PI boosting can target a task running on
another CPU, so ops.cpu_release() can now be invoked for a remote CPU,
with that CPU's rq locked. The legacy scx_bpf_reenqueue_local() was
already switched to the locked rq to handle this.

Signed-off-by: Kuba Piecuch <jpiecuch@google.com>
---
 kernel/sched/ext/ext.c      | 11 +++++++++++
 kernel/sched/ext/internal.h |  2 ++
 2 files changed, 13 insertions(+)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 4d02d2d6e6d0..d56c692f3dfa 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -4433,6 +4433,17 @@ static void switching_to_scx(struct rq *rq, struct task_struct *p)
 
 static void switched_from_scx(struct rq *rq, struct task_struct *p)
 {
+	/*
+	 * The running task changing class (sched_setscheduler(), PI boost)
+	 * keeps the CPU without a context switch, so neither
+	 * put_prev_task_scx() nor the idle class sees it. Every class @p can
+	 * switch to is above SCX, release the CPU. The new class isn't known
+	 * yet. The scheduler being disabled is excluded by the static key.
+	 */
+	if (static_branch_unlikely(&__scx_cpu_avail_enabled) &&
+	    task_current_donor(rq, p))
+		scx_cpu_release(rq, SCX_CPU_PREEMPT_CLASS_CHANGE, p);
+
 	if (task_dead_and_done(p))
 		return;
 
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 2d8c789c1e76..65b6ec7dc2f3 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -300,6 +300,8 @@ enum scx_cpu_preempt_reason {
 	SCX_CPU_PREEMPT_UNKNOWN,
 	/* next task is being scheduled by &fair_sched_class (partial mode) */
 	SCX_CPU_PREEMPT_FAIR,
+	/* the running task switched from SCX to another sched_class */
+	SCX_CPU_PREEMPT_CLASS_CHANGE,
 };
 
 /*
-- 
2.56.0.176.g3575e65682-goog