[PATCH sched_ext/for-7.1-fixes] sched_ext: Call wakeup_preempt() in local_dsq_post_enq()

[Kuba Piecuch <jpiecuch@google.com>]

There are several edge cases (see linked thread) where an IMMED task
can be left lingering on a local DSQ if an RT task swoops in at the
wrong time. All of these edge cases are due to rq->next_class being idle
even after dispatching a task to rq's local DSQ. We should bump
rq->next_class to &ext_sched_class as soon as we've inserted a task into
the local DSQ.

To optimize the common case of rq->next_class == &ext_sched_class,
only call wakeup_preempt() if rq->next_class is below EXT. If next_class
is EXT or above, wakeup_preempt() is a no-op anyway.

Link: https://lore.kernel.org/all/DHZPHUFXB4N3.2RY28MUEWBNYK@google.com/
Signed-off-by: Kuba Piecuch <jpiecuch@google.com>
---
 kernel/sched/ext.c | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 1f670028bf19..034df77e3af1 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -1393,7 +1393,6 @@ static void local_dsq_post_enq(struct scx_sched *sch, struct scx_dispatch_q *dsq
 			       struct task_struct *p, u64 enq_flags)
 {
 	struct rq *rq = container_of(dsq, struct rq, scx.local_dsq);
-	bool preempt = false;
 
 	call_task_dequeue(sch, rq, p, 0);
 
@@ -1408,11 +1407,19 @@ static void local_dsq_post_enq(struct scx_sched *sch, struct scx_dispatch_q *dsq
 	if ((enq_flags & SCX_ENQ_PREEMPT) && p != rq->curr &&
 	    rq->curr->sched_class == &ext_sched_class) {
 		rq->curr->scx.slice = 0;
-		preempt = true;
+		resched_curr(rq);
 	}
 
-	if (preempt || sched_class_above(&ext_sched_class, rq->curr->sched_class))
-		resched_curr(rq);
+	/*
+	 * If @rq->next_class is currently idle, we need to bump it
+	 * to &ext_sched_class using wakeup_preempt(). Otherwise, if we drop
+	 * the rq lock later in the pick and an RT task wakes up on @rq,
+	 * wakeup_preempt_idle() will be called during RT task wakeup and
+	 * SCX won't have an opportunity to re-enqueue IMMED tasks from @rq's
+	 * local DSQ.
+	 */
+	if (sched_class_above(&ext_sched_class, rq->next_class))
+		wakeup_preempt(rq, p, 0);
 }
 
 static void dispatch_enqueue(struct scx_sched *sch, struct rq *rq,
-- 
2.54.0.rc2.544.gc7ae2d5bb8-goog