[PATCH 1/2] sched_ext: scx_qmap: Restore unused idle claims from ops.dispatch()

[Tejun Heo <tj@kernel.org>]

scx_qmap tracks idle cids itself. pick_direct_dispatch_cid() claims a cid by
clearing its bit and the task is inserted into that cid's local DSQ, which
kicks the CPU. When the task does not arrive, for example because the insert
fell back to the global DSQ after an affinity change, the CPU wakes, finds
nothing and picks idle again. That is not an idle transition, so
ops.update_idle() is not called and the cid stays marked busy until an
unrelated task runs on it.

Restore the claim from ops.dispatch(). The kick guarantees a dispatch on the
kicked CPU, and when it finds nothing to run with a NULL @prev, the CPU is
going back to idle. Document the pattern in ops.update_idle(), which reports
only actual transitions.

Signed-off-by: Tejun Heo <tj@kernel.org>
Cc: Andrea Righi <arighi@nvidia.com>
---
 kernel/sched/ext/internal.h    |  6 ++++++
 tools/sched_ext/scx_qmap.bpf.c | 14 ++++++++++----
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 076a351bb3f2..0adaf649d5e0 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -572,6 +572,12 @@ struct sched_ext_ops {
 	 *
 	 * Specify the %SCX_OPS_KEEP_BUILTIN_IDLE flag to keep the built-in idle
 	 * tracking.
+	 *
+	 * Only actual transitions are reported. A CPU that is claimed with an
+	 * idle pick and kicked but dispatches no task returns to idle without a
+	 * transition. A scheduler tracking idle CPUs itself must restore the
+	 * idle state from ops.dispatch() when it returns without the next task
+	 * to run.
 	 */
 	void (*update_idle)(s32 cpu, bool idle);
 
diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index bda56c37acb5..67b7c01cae55 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -818,10 +818,10 @@ void BPF_STRUCT_OPS(qmap_dispatch, s32 cid, struct task_struct *prev)
 			batch--;
 			cpuc->dsp_cnt--;
 			if (!batch || !scx_bpf_dispatch_nr_slots()) {
-				if (scan_shared_dsq(false))
+				if (scan_shared_dsq(false) ||
+				    scx_bpf_dsq_move_to_local(SHARED_DSQ, needs_immed(cid)))
 					return;
-				scx_bpf_dsq_move_to_local(SHARED_DSQ, needs_immed(cid));
-				return;
+				goto prev;
 			}
 			if (!cpuc->dsp_cnt)
 				break;
@@ -832,10 +832,14 @@ void BPF_STRUCT_OPS(qmap_dispatch, s32 cid, struct task_struct *prev)
 
 	if (scan_shared_dsq(false))
 		return;
-
+prev:
 	/*
 	 * No other tasks. @prev will keep running. Update its core_sched_seq as
 	 * if the task were enqueued and dispatched immediately.
+	 *
+	 * No @prev to keep running means the CPU goes idle. If its claim was
+	 * never used, that is not a transition and ops.update_idle() stays
+	 * silent. Restore the claim here.
 	 */
 	if (prev) {
 		taskc = lookup_task_ctx(prev);
@@ -844,6 +848,8 @@ void BPF_STRUCT_OPS(qmap_dispatch, s32 cid, struct task_struct *prev)
 
 		taskc->core_sched_seq =
 			qa.core_sched_tail_seqs[weight_to_idx(prev->scx.weight)]++;
+	} else {
+		cmask_set(cid, &qa.idle_cids.mask);
 	}
 }
 
-- 
2.55.0