[PATCH v5 sched_ext/for-7.3 19/33] sched_ext: Add SCX_ENQ_IGNORE_CAPS for in-place restore

[Tejun Heo <tj@kernel.org>]

A SAVE/RESTORE requeue re-inserts a running task in place and is immediately
followed by set_next_task_scx(). It is not a real scheduling event: the task
is already admitted to its cid and must return to the local DSQ
unconditionally.

scx_caps_for_enq() maps an enqueue to the cap its local-DSQ insert requires.
Add SCX_ENQ_IGNORE_CAPS, set it on the RESTORE-in-place branch of
enqueue_task_scx(), and have scx_caps_for_enq() require no caps for it, so
the cid admission gate never diverts an in-place restore to the reject DSQ.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c      | 7 +++++--
 kernel/sched/ext/internal.h | 1 +
 kernel/sched/ext/sub.h      | 3 +++
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 3e1829126fee..80c8971638f9 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -1878,10 +1878,13 @@ static void enqueue_task_scx(struct rq *rq, struct task_struct *p, int core_enq_
 	 * Restoring a running task will be immediately followed by
 	 * set_next_task_scx() which expects the task to not be on the BPF
 	 * scheduler as tasks can only start running through local DSQs. Force
-	 * direct-dispatch into the local DSQ by setting the sticky_cpu.
+	 * direct-dispatch into the local DSQ by setting the sticky_cpu. Mark
+	 * IGNORE_CAPS to force entry into the local DSQ.
 	 */
-	if (unlikely(enq_flags & ENQUEUE_RESTORE) && task_current(rq, p))
+	if (unlikely(enq_flags & ENQUEUE_RESTORE) && task_current(rq, p)) {
 		sticky_cpu = cpu_of(rq);
+		enq_flags |= SCX_ENQ_IGNORE_CAPS;
+	}
 
 	if (p->scx.flags & SCX_TASK_QUEUED) {
 		WARN_ON_ONCE(!task_runnable(p));
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index a360d312702b..079ee941b343 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -1568,6 +1568,7 @@ enum scx_enq_flags {
 	SCX_ENQ_DSQ_PRIQ	= 1LLU << 57,
 	SCX_ENQ_NESTED		= 1LLU << 58,
 	SCX_ENQ_GDSQ_FALLBACK	= 1LLU << 59,	/* fell back to global DSQ */
+	SCX_ENQ_IGNORE_CAPS	= 1LLU << 60,	/* admit to local DSQ ignoring caps */
 };
 
 enum scx_deq_flags {
diff --git a/kernel/sched/ext/sub.h b/kernel/sched/ext/sub.h
index 51400e174cd0..a1a2825dc04e 100644
--- a/kernel/sched/ext/sub.h
+++ b/kernel/sched/ext/sub.h
@@ -106,6 +106,9 @@ static inline u64 scx_missing_caps(struct scx_sched *sch, s32 cpu, u64 needed)
 /* map @enq_flags to the SCX_CAP_* bit required for the local-DSQ insert */
 static inline u64 scx_caps_for_enq(u64 enq_flags)
 {
+	/* a restored task must be put into the local DSQ regardless of caps */
+	if (enq_flags & SCX_ENQ_IGNORE_CAPS)
+		return 0;
 	return 0;
 }
 
-- 
2.55.0