[PATCH 11/12] sched_ext: scx_qmap: Add proxy execution support

[Andrea Righi]

Blocked sched_ext tasks must remain available to the scheduler before
the core can use their scheduling contexts to execute mutex owners. BPF
schedulers opt into receiving these donor requests through
SCX_OPS_ENQ_BLOCKED and are responsible for deciding when to admit them.

Enable SCX_OPS_ENQ_BLOCKED in scx_qmap and immediately admit every
blocked donor. Query the mutex owner's cid and use it when allowed by
the donor's affinity, falling back to the donor's current cid otherwise.
Insert the donor at the head of the selected local DSQ.

Applying the same policy after a proxy migration expedites the donor on
the owner's CPU so proxy execution can begin without waiting behind
unrelated local work.

Signed-off-by: Andrea Righi <arighi@xxxxxxxxxx>
---
 tools/sched_ext/scx_qmap.bpf.c | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index fd9a82a676278..5ecd2ffe05270 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -396,6 +396,21 @@ void BPF_STRUCT_OPS(qmap_enqueue, struct task_struct *p, u64 enq_flags)
 	 */
 	taskc->core_sched_seq = qa.core_sched_tail_seqs[idx]++;
 
+	/*
+	 * Admit blocked proxy donors immediately. Dispatch directly on the mutex
+	 * owner's cid when available to facilitate the proxy execution switch.
+	 * SCX_ENQ_HEAD also puts the donor at the head of the local DSQ and
+	 * expedites the destination enqueue after a proxy migration.
+	 */
+	if (scx_bpf_task_is_blocked(p)) {
+		cid = scx_bpf_task_proxy_cid(p);
+		if (cid < 0 || !cmask_test(cid, &taskc->cpus_allowed))
+			cid = scx_bpf_task_cid(p);
+		scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | cid,
+				   slice_ns, enq_flags | SCX_ENQ_HEAD);
+		return;
+	}
+
 	/*
 	 * IMMED stress testing: Every immed_stress_nth'th enqueue, dispatch
 	 * directly to prev_cpu's local DSQ even when busy to force dsq->nr > 1
@@ -1216,7 +1231,8 @@ void BPF_STRUCT_OPS(qmap_sub_detach, struct scx_sub_detach_args *args)
 }
 
 SCX_OPS_CID_DEFINE(qmap_ops,
-	       .flags			= SCX_OPS_ENQ_EXITING | SCX_OPS_TID_TO_TASK,
+	       .flags			= SCX_OPS_ENQ_EXITING | SCX_OPS_TID_TO_TASK |
+					  SCX_OPS_ENQ_BLOCKED,
 	       .select_cid		= (void *)qmap_select_cid,
 	       .enqueue			= (void *)qmap_enqueue,
 	       .dequeue			= (void *)qmap_dequeue,
-- 
2.55.0