Merge tag 'sched_ext-for-7.3-rc3-fixes' of git://git.kernel.org/pub/scm/linux/kernel/git/tj/sched_ext

Pull sched_ext fixes from Tejun Heo:

 - An error raised by a BPF program before the scheduler finished
   enabling was consumed by the disable path's pre-enable shortcut,
   leaving a running scheduler that couldn't be disabled and was later
   freed while in use.

 - Two compat kfuncs dereferenced a NULL scheduler when handed an exited
   or idle task, oopsing the kernel.

 - Keep-running decisions in the dispatch path used the root scheduler's
   flags for tasks belonging to a sub-scheduler, causing warnings and
   stalls.

 - Schedulers with their own CPU ID mapping had no way to learn which
   IDs are online. Add a kernel-maintained online mask to plug the hole.

 - Cgroup idle state: the initial cpu.idle state wasn't passed on cgroup
   init and same-value rewrites delivered spurious callbacks.

 - Example scheduler fixes for a reenqueue loop on attach, placements on
   CPUs without effective grants, stalled partition work and stale idle
   tracking.

* tag 'sched_ext-for-7.3-rc3-fixes' of git://git.kernel.org/pub/scm/linux/kernel/git/tj/sched_ext:
  sched_ext: Maintain an online cid mask in the scheduler arena
  sched_ext: scx_qmap: Restore unused idle claims from ops.dispatch()
  sched_ext: Close the pre-enable ops error claim window
  sched_ext: scx_qmap: Fix pending partition work handoff
  sched_ext: scx_qmap: Place only on cids whose caps are in effect
  sched_ext: scx_qmap: Do not add IMMED to rescue inserts
  sched_ext: Use @prev's scheduler for the keep decisions in dispatch_one()
  sched_ext: Rename sch to root_sch in dispatch_one()
  sched_ext: Fix NULL sched deref in kfunc sub-sched error paths
  sched_ext: Don't deliver duplicate ops.cgroup_set_idle() for same value
  sched_ext: Pass the initial cpu.idle state in scx_cgroup_init_args

3 files changed, 115 insertions, 51 deletions

diff --git a/tools/sched_ext/include/scx/common.bpf.h b/tools/sched_ext/include/scx/common.bpf.h
index 76f5e025e107..2ddb01a059fd 100644
--- a/tools/sched_ext/include/scx/common.bpf.h
+++ b/tools/sched_ext/include/scx/common.bpf.h
@@ -113,6 +113,7 @@ s32 scx_bpf_this_cid(void) __ksym __weak;
 struct task_struct *scx_bpf_cid_curr(s32 cid) __ksym __weak;
 u32 scx_bpf_nr_cids(void) __ksym __weak;
 u32 scx_bpf_nr_online_cids(void) __ksym __weak;
+const void __arena *scx_bpf_online_cmask(void) __ksym __weak;
 u32 scx_bpf_cidperf_cap(s32 cid) __ksym __weak;
 u32 scx_bpf_cidperf_cur(s32 cid) __ksym __weak;
 s32 scx_bpf_cidperf_set(s32 cid, u32 perf) __ksym __weak;
diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index 9f6e61d7ca07..67b7c01cae55 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -24,6 +24,9 @@
  *            time-share that stays self-local.
  *   self   - The excl cpus the node kept for itself, plus all of held_shared.
  *   owner  - Who holds a cid - a child slot, CID_SELF, or CID_NONE.
+ *   avail  - Cpus whose caps are in effect, per ops.sub_ecaps_updated().
+ *   usable - self AND avail. Placement decisions use this: self is the
+ *            delegation split and can run ahead of what the cpus honor.
  *
  * The scheduler splits its held-excl cpus among self and the children in
  * proportion to each node's cpu.weight, handing each the floor of its share as
@@ -208,8 +211,8 @@ static int qmap_spin_lock(struct bpf_res_spin_lock *lock)
 }
 
 /*
- * Try prev_cid, then scan cpus_allowed AND idle_cids AND self_cids round-robin
- * from prev_cid + 1. Atomic claim retries on race; bounded by
+ * Try prev_cid, then scan cpus_allowed AND idle_cids AND usable_cids
+ * round-robin from prev_cid + 1. Atomic claim retries on race; bounded by
  * IDLE_PICK_RETRIES to keep the verifier's insn budget in check.
  */
 #define IDLE_PICK_RETRIES	16
@@ -221,7 +224,7 @@ static s32 pick_direct_dispatch_cid(struct task_struct *p, s32 prev_cid,
 	s32 cid;
 	u32 i;
 
-	if (cmask_test(prev_cid, &qa.self_cids.mask) &&
+	if (cmask_test(prev_cid, &qa.usable_cids.mask) &&
 	    cmask_test_and_clear(prev_cid, &qa.idle_cids.mask))
 		return prev_cid;
 
@@ -229,7 +232,7 @@ static s32 pick_direct_dispatch_cid(struct task_struct *p, s32 prev_cid,
 	bpf_for(i, 0, IDLE_PICK_RETRIES) {
 		cid = cmask_next_and2_set_wrap(&taskc->cpus_allowed,
 					       &qa.idle_cids.mask,
-					       &qa.self_cids.mask, cid + 1);
+					       &qa.usable_cids.mask, cid + 1);
 		barrier_var(cid);
 		if (cid >= nr_cids)
 			return -1;
@@ -358,8 +361,8 @@ s32 BPF_STRUCT_OPS(qmap_select_cid, struct task_struct *p,
 }
 
 /*
- * A received time-shared cid is held ENQ_IMMED-only, so inserts must set
- * SCX_ENQ_IMMED.
+ * A received time-shared cid is held ENQ_IMMED-only, so inserts meant to run
+ * there must set SCX_ENQ_IMMED.
  */
 static u64 needs_immed(s32 cid)
 {
@@ -444,9 +447,11 @@ void BPF_STRUCT_OPS(qmap_enqueue, struct task_struct *p, u64 enq_flags)
 	 * didn't grant them or we delegated them to children - would starve in
 	 * SHARED/FIFO since we only pull from those on self cids.
 	 *
-	 * Force it onto its first allowed cid's local DSQ. If we hold that cid
-	 * it runs. Otherwise the insert carries SCX_ENQ_RESCUE and the kernel
-	 * diverts the task to its rescue path.
+	 * Force it onto its first allowed cid's local DSQ with SCX_ENQ_RESCUE.
+	 * If we hold ENQ on that cid it runs. Otherwise the kernel diverts the
+	 * task to its rescue path. IMMED would turn the insert into a legal
+	 * placement on a time-shared cid and the kernel would bounce it back
+	 * here instead of rescuing it.
 	 */
 	if (!cmask_intersects(&taskc->cpus_allowed, &qa.self_cids.mask)) {
 		s32 c = cmask_next_set_wrap(&taskc->cpus_allowed, 0);
@@ -455,7 +460,7 @@ void BPF_STRUCT_OPS(qmap_enqueue, struct task_struct *p, u64 enq_flags)
 			taskc->force_local = false;
 			__sync_fetch_and_add(&qa.nr_rescue_dsp, 1);
 			scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL_ON | c, slice_ns,
-					   enq_flags | needs_immed(c) | SCX_ENQ_RESCUE);
+					   enq_flags | SCX_ENQ_RESCUE);
 			return;
 		}
 	}
@@ -540,7 +545,7 @@ void BPF_STRUCT_OPS(qmap_enqueue, struct task_struct *p, u64 enq_flags)
 		scx_bpf_dsq_insert(p, SHARED_DSQ, 0, enq_flags);
 		cid = cmask_next_and2_set_wrap(&taskc->cpus_allowed,
 					       &qa.idle_cids.mask,
-					       &qa.self_cids.mask, 0);
+					       &qa.usable_cids.mask, 0);
 		if (cid < scx_bpf_nr_cids())
 			scx_bpf_kick_cid(cid, SCX_KICK_IDLE);
 		return;
@@ -618,7 +623,7 @@ static bool scan_shared_dsq(bool from_timer)
 			if (c >= 0 && c < scx_bpf_nr_cids()) {
 				__sync_fetch_and_add(&qa.nr_rescue_dsp, 1);
 				scx_bpf_dsq_move(BPF_FOR_EACH_ITER, p, SCX_DSQ_LOCAL_ON | c,
-						 needs_immed(c) | SCX_ENQ_RESCUE);
+						 SCX_ENQ_RESCUE);
 			}
 			continue;
 		}
@@ -644,22 +649,27 @@ static bool scan_shared_dsq(bool from_timer)
 		if (!(taskc = lookup_task_ctx(p)))
 			return false;
 
-		/* only run highpri tasks on cids this node holds, not delegated ones */
+		/* only run highpri tasks on cids this node can use right now */
 		if (cmask_test(this_cid, &taskc->cpus_allowed) &&
-		    cmask_test(this_cid, &qa.self_cids.mask))
+		    cmask_test(this_cid, &qa.usable_cids.mask))
 			cid = this_cid;
 		else
 			cid = cmask_next_and_set_wrap(&taskc->cpus_allowed,
-						      &qa.self_cids.mask,
+						      &qa.usable_cids.mask,
 						      this_cid + 1);
 		if (cid >= nr_cids) {
-			/* stranded after the cull - rescue it from here */
-			s32 c = cmask_next_set_wrap(&taskc->cpus_allowed, 0);
+			s32 c;
+
+			/* self cids lack caps in effect yet, leave it queued */
+			if (cmask_intersects(&taskc->cpus_allowed, &qa.self_cids.mask))
+				continue;
 
+			/* stranded after the cull - rescue it from here */
+			c = cmask_next_set_wrap(&taskc->cpus_allowed, 0);
 			if (c >= 0 && c < nr_cids) {
 				__sync_fetch_and_add(&qa.nr_rescue_dsp, 1);
 				scx_bpf_dsq_move(BPF_FOR_EACH_ITER, p, SCX_DSQ_LOCAL_ON | c,
-						 needs_immed(c) | SCX_ENQ_RESCUE);
+						 SCX_ENQ_RESCUE);
 			}
 			continue;
 		}
@@ -808,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;
@@ -822,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);
@@ -834,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);
 	}
 }
 
@@ -1113,7 +1129,7 @@ void BPF_STRUCT_OPS(qmap_update_idle, s32 cid, bool idle)
 	/*
 	 * The kernel delivers update_idle() for every cid this node holds
 	 * SCX_CAP_BASE on. Track every cid's idle state regardless of
-	 * delegation: the direct-dispatch pick masks idle_cids with self_cids
+	 * delegation: the direct-dispatch pick masks idle_cids with usable_cids
 	 * at selection, so a cid already idle when it returns to self needs no
 	 * reseed here.
 	 */
@@ -1285,11 +1301,16 @@ struct {
 	__type(value, struct round_robin_timer);
 } round_robin_timer SEC(".maps");
 
+enum part_pending_flags {
+	PART_REFRESH = BIT_U64(0),
+	PART_REDISTRIBUTE = BIT_U64(1),
+};
+
 /*
  * Partition update synchronization. qa.part can be written from concurrent
  * contexts. This single-runner guard admits one writer at a time without
  * holding a lock across the grant/revoke kfuncs. part_pending coalesces
- * repartition requests that arrive while it is held.
+ * refresh and repartition requests that arrive while it is held.
  *
  * They live in .bss, not the arena: rr_advance() runs from a bpf_timer
  * callback, where the verifier rejects atomic ops on arena memory.
@@ -1538,6 +1559,19 @@ static __noinline void account_alloc(void)
 }
 
 /*
+ * usable_cids = self_cids & avail_cids. The inputs have separate writers,
+ * apply_partition() and qmap_sub_ecaps_updated(), so the result is rebuilt in
+ * full under the partition guard, in scratch first so that readers never see
+ * self_cids alone.
+ */
+static void refresh_usable(void)
+{
+	cmask_copy(&qa.usable_scratch.mask, &qa.self_cids.mask);
+	cmask_and(&qa.usable_scratch.mask, &qa.avail_cids.mask);
+	cmask_copy(&qa.usable_cids.mask, &qa.usable_scratch.mask);
+}
+
+/*
  * apply_partition - execute the plan compute_partition() built
  *
  * Turn the owner map into the per-child, shared and self cmasks and issue the
@@ -1559,6 +1593,7 @@ __noinline void apply_partition(void)
 	/* no excl cpu: run own tasks on the held shares, evict children */
 	if (!qa.part.nr_excl) {
 		cmask_copy(&qa.self_cids.mask, &qa.held_shared.mask);
+		refresh_usable();
 		bpf_for(i, 0, MAX_SUB_SCHEDS)
 			if (qa.sub_sched_ctxs[i].cgroup_id)
 				scx_bpf_sub_kill(qa.sub_sched_ctxs[i].cgroup_id,
@@ -1596,6 +1631,7 @@ __noinline void apply_partition(void)
 		else if (o == CID_SELF)
 			cmask_set(cid, &qa.self_cids.mask);
 	}
+	refresh_usable();
 
 	/*
 	 * Apply each child's exclusive cids as a delta against its previous
@@ -1643,33 +1679,46 @@ __noinline void apply_partition(void)
 	}
 }
 
-/*
- * Recompute the split off the node's held caps and apply it. The contexts this
- * runs from (the sub-sched and cgroup callbacks, the rr timer) are not
- * serialized by the kernel, so a single runner does the work. A caller that
- * finds the guard held leaves part_pending set; the holder drains it before
- * releasing, with the rr timer as a backstop.
+/**
+ * execute_partition - Run pending partition updates
+ *
+ * The rr timer is the backstop if the loop reaches its iteration limit.
  */
-static void redistribute(void)
+static void execute_partition(void)
 {
+	u64 pending;
 	s32 i;
 
-	__sync_fetch_and_or(&part_pending, 1);
+	bpf_for(i, 0, 1024) {
+		if (!part_try_start())
+			break;
 
-	if (!part_try_start())
-		return;
+		pending = __sync_fetch_and_and(&part_pending, 0);
+		if (pending & PART_REDISTRIBUTE) {
+			/* charge elapsed time before repartitioning */
+			account_alloc();
+			compute_partition();
+			apply_partition();
+		} else if (pending & PART_REFRESH) {
+			refresh_usable();
+		}
 
-	bpf_for(i, 0, 1024) {
-		__sync_fetch_and_and(&part_pending, 0);
-		/* charge elapsed time to the current partition before rebuilding it */
-		account_alloc();
-		compute_partition();
-		apply_partition();
+		/*
+		 * Requests are published before trying the guard. Releasing it
+		 * before checking pending work ensures a racing request is
+		 * either observed here or handled by a caller that acquires the
+		 * guard.
+		 */
+		part_end();
 		if (!__sync_fetch_and_or(&part_pending, 0))
 			break;
 	}
+}
 
-	part_end();
+static void redistribute(void)
+{
+	__sync_fetch_and_or(&part_pending, PART_REDISTRIBUTE);
+	execute_partition();
 }
 
 /*
@@ -1683,6 +1732,7 @@ int flush_alloc(void *ctx)
 	if (part_try_start()) {
 		account_alloc();
 		part_end();
+		execute_partition();
 	}
 	return 0;
 }
@@ -1740,9 +1790,7 @@ static void rr_advance(void)
 
 	part_end();
 
-	/* a resplit queued while we held the guard supersedes this rotation */
-	if (__sync_fetch_and_or(&part_pending, 0))
-		redistribute();
+	execute_partition();
 }
 
 /* advance the time-shared cid pool every round_robin_ns */
@@ -1837,8 +1885,11 @@ s32 BPF_STRUCT_OPS_SLEEPABLE(qmap_init)
 	cmask_init(&qa.rr_cids.mask, 0, nr_cids);
 	cmask_init(&qa.prev_rr_cids.mask, 0, nr_cids);
 	cmask_init(&qa.self_cids.mask, 0, nr_cids);
+	cmask_init(&qa.avail_cids.mask, 0, nr_cids);
+	cmask_init(&qa.usable_cids.mask, 0, nr_cids);
 	cmask_init(&qa.to_revoke_cids.mask, 0, nr_cids);
 	cmask_init(&qa.to_grant_cids.mask, 0, nr_cids);
+	cmask_init(&qa.usable_scratch.mask, 0, nr_cids);
 	cmask_init(&qa.held_excl.mask, 0, nr_cids);
 	cmask_init(&qa.held_shared.mask, 0, nr_cids);
 
@@ -1852,14 +1903,16 @@ s32 BPF_STRUCT_OPS_SLEEPABLE(qmap_init)
 	}
 
 	/*
-	 * The root starts holding every cid. qmap_sub_ecaps_updated() maintains
-	 * per-cid shared state as effective caps settle, and redistribute()
-	 * rebuilds owner and self from held caps. A non-root node starts with
-	 * nothing.
+	 * The root starts holding every cid and gets no ecaps notifications, so
+	 * its avail set is fixed here. qmap_sub_ecaps_updated() maintains the
+	 * per-cid state as effective caps settle, and redistribute() rebuilds
+	 * owner and self from held caps. A non-root node starts with nothing.
 	 */
 	bpf_for(i, 0, nr_cids) {
 		if (!sub_cgroup_id) {
 			cmask_set(i, &qa.self_cids.mask);
+			cmask_set(i, &qa.avail_cids.mask);
+			cmask_set(i, &qa.usable_cids.mask);
 			qa.part.cid_owner[i] = CID_SELF;
 		} else {
 			qa.part.cid_owner[i] = CID_NONE;
@@ -2000,12 +2053,19 @@ void BPF_STRUCT_OPS(qmap_sub_ecaps_updated, s32 cid, u64 before, u64 after)
 {
 	/*
 	 * Effective caps updated. Track which cids hold shared caps so a self
-	 * task placed there enqueues IMMED.
+	 * task placed there enqueues IMMED, and which cids have ENQ_IMMED in
+	 * effect at all (avail, see the header comment).
 	 */
-	if (after & SCX_CAP_ENQ_IMMED)
+	if (after & SCX_CAP_ENQ_IMMED) {
 		qa.cid_shared[cid] = (after & SCX_CAP_ENQ) ? 0 : 1;
-	else
+		cmask_set(cid, &qa.avail_cids.mask);
+	} else {
 		qa.cid_shared[cid] = 0;
+		cmask_clear(cid, &qa.avail_cids.mask);
+	}
+
+	__sync_fetch_and_or(&part_pending, PART_REFRESH);
+	execute_partition();
 }
 
 SCX_OPS_CID_DEFINE(qmap_ops,
diff --git a/tools/sched_ext/scx_qmap.h b/tools/sched_ext/scx_qmap.h
index c78d61806b39..e95fffcf7b23 100644
--- a/tools/sched_ext/scx_qmap.h
+++ b/tools/sched_ext/scx_qmap.h
@@ -165,12 +165,15 @@ struct qmap_arena {
 
 	/* bpf-internal cmasks (embedded, see struct qmap_cmask) */
 	struct qmap_cmask self_cids;	/* cids this node runs its own tasks on */
+	struct qmap_cmask avail_cids;	/* cids with caps in effect on the cpu */
+	struct qmap_cmask usable_cids;	/* self_cids & avail_cids, placeable right now */
 	struct qmap_cmask idle_cids;	/* idle state of all cids regardless of delegation */
 	struct qmap_cmask rr_cids;	/* the shared pool, as a mask for grant/revoke */
 
 	/* scratch cmasks */
 	struct qmap_cmask to_revoke_cids; /* delta cids to revoke */
 	struct qmap_cmask to_grant_cids; /* delta cids to grant */
+	struct qmap_cmask usable_scratch; /* refresh_usable() build area */
 	struct qmap_cmask prev_rr_cids; /* previous shared pool, to clear stale grants */
 	struct qmap_cmask held_excl;	/* cids held excl (ENQ): delegatable */
 	struct qmap_cmask held_shared;	/* cids held shared (ENQ_IMMED only): self-local */