linux - kernel/sched/ext/sub.c

2682 lines · 78.7 KiB

1	// SPDX-License-Identifier: GPL-2.0
2	/*
3	* BPF extensible scheduler class: Documentation/scheduler/sched-ext.rst
4	*
5	* Sub-scheduler hierarchy support.
6	*
7	* A sub-scheduler is an scx_sched attached to a cgroup subtree under another
8	* scx_sched. This file holds the sub-scheduler implementation: the scheduler
9	* tree walk, capability delegation, per-shard cap state and its sync, and the
10	* sub-scheduler enable/disable paths. The core dispatch/enqueue machinery it
11	* builds on lives in ext.c.
12	*
13	* Copyright (c) 2026 Meta Platforms, Inc. and affiliates.
14	* Copyright (c) 2026 Tejun Heo <[email protected]>
15	*/
16	#include <linux/rhashtable.h>
17	#include "internal.h"
18	#include "cid.h"
19	#include "arena.h"
20	#include "sub.h"
21	#include "inlines.h"
22
23	#ifdef CONFIG_EXT_SUB_SCHED
24
25	/*
26	* On while any sub-scheduler exists so that a root-only system doesn't pay for
27	* the sub-sched portions of hot paths. See scx_has_subs().
28	*/
29	DEFINE_STATIC_KEY_FALSE(__scx_has_subs);
30
31	/* latched at root enable before any rescue runs */
32	static s32 scx_rescue_bw_1024;
33	static s64 scx_rescue_quantum_ns;
34	static s64 scx_rescue_sat_delta_ns;
35	static unsigned long scx_rescue_decay_halflife;
36	static unsigned long scx_rescue_overload_after;
37
38	/**
39	* scx_skip_subtree_pre - Skip @pos's subtree in a pre-order walk
40	* @pos: current position
41	* @root: walk root
42	*
43	* In a walk started by scx_next_descendant_pre(), continue past @pos's subtree:
44	* return @pos's next sibling, or the closest ancestor's next sibling, or NULL
45	* if @pos's subtree is the last under @root. Same locking rules.
46	*/
47	struct scx_sched scx_skip_subtree_pre(struct scx_sched pos, struct scx_sched *root)
48	{
49	struct scx_sched *next;
50
51	lockdep_assert(lockdep_is_held(&scx_enable_mutex) \|\|
52	lockdep_is_held(&scx_sched_lock) \|\|
53	rcu_read_lock_any_held());
54
55	while (pos != root) {
56	next = list_next_or_null_rcu(&scx_parent(pos)->children, &pos->sibling,
57	struct scx_sched, sibling);
58	if (next)
59	return next;
60	pos = scx_parent(pos);
61	}
62	return NULL;
63	}
64
65	/**
66	* scx_next_descendant_pre - find the next descendant for pre-order walk
67	* @pos: the current position (%NULL to initiate traversal)
68	* @root: sched whose descendants to walk
69	*
70	* To be used by scx_for_each_descendant_pre(). Find the next descendant to
71	* visit for pre-order traversal of @root's descendants. @root is included in
72	* the iteration and the first node to be visited.
73	*/
74	struct scx_sched scx_next_descendant_pre(struct scx_sched pos, struct scx_sched *root)
75	{
76	struct scx_sched *next;
77
78	lockdep_assert(lockdep_is_held(&scx_enable_mutex) \|\|
79	lockdep_is_held(&scx_sched_lock) \|\|
80	rcu_read_lock_any_held());
81
82	/* if first iteration, visit @root */
83	if (!pos)
84	return root;
85
86	/* visit the first child if exists */
87	next = list_first_or_null_rcu(&pos->children, struct scx_sched, sibling);
88	if (next)
89	return next;
90
91	/* no child, visit my or the closest ancestor's next sibling */
92	return scx_skip_subtree_pre(pos, root);
93	}
94
95	static struct scx_sched *scx_find_sub_sched(u64 cgroup_id)
96	{
97	return rhashtable_lookup(&scx_sched_hash, &cgroup_id,
98	scx_sched_hash_params);
99	}
100
101	void scx_set_task_sched(struct task_struct p, struct scx_sched sch)
102	{
103	rcu_assign_pointer(p->scx.sched, sch);
104	}
105
106	struct cgroup sch_cgroup(struct scx_sched sch)
107	{
108	return sch->cgrp;
109	}
110
111	/* for each descendant of @cgrp including self, set ->scx_sched to @sch */
112	void set_cgroup_sched(struct cgroup cgrp, struct scx_sched sch)
113	{
114	struct cgroup *pos;
115	struct cgroup_subsys_state *css;
116
117	cgroup_for_each_live_descendant_pre(pos, css, cgrp)
118	rcu_assign_pointer(pos->scx_sched, sch);
119	}
120
121	static void free_pshard(struct scx_pshard *pshard)
122	{
123	struct scx_caps_updated *cu;
124
125	if (!pshard)
126	return;
127	cu = &pshard->caps_updated;
128	if (cu->cmask_arena_out)
129	scx_arena_free(pshard->sch, cu->cmask_arena_out,
130	struct_size_t(struct scx_cmask, bits,
131	SCX_CMASK_NR_WORDS(pshard->nr_cids)));
132	kfree(pshard);
133	}
134
135	void scx_free_pshards(struct scx_sched *sch)
136	{
137	s32 si;
138
139	if (!sch->pshard)
140	return;
141	for (si = 0; si < sch->nr_pshards; si++)
142	free_pshard(sch->pshard[si]);
143	kfree(sch->pshard);
144	}
145
146	static struct scx_pshard alloc_pshard(struct scx_sched sch, s32 shard_idx, s32 node)
147	{
148	const struct scx_cid_shard *shard =
149	&rcu_dereference_protected(scx_cid_shard_ranges,
150	lockdep_is_held(&scx_enable_mutex))[shard_idx];
151	size_t cmask_size = struct_size_t(struct scx_cmask, bits,
152	SCX_CMASK_NR_WORDS(shard->nr_cids));
153	struct scx_pshard *pshard;
154	struct scx_caps_updated *cu;
155	s32 i;
156
157	pshard = kzalloc_node(sizeof(*pshard), GFP_KERNEL, node);
158	if (!pshard)
159	return NULL;
160
161	raw_spin_lock_init(&pshard->lock);
162	pshard->sch = sch;
163	pshard->base = shard->base_cid;
164	pshard->nr_cids = shard->nr_cids;
165
166	for (i = 0; i < __SCX_NR_CAPS; i++)
167	scx_cmask_init(&pshard->caps[i].cmask, shard->base_cid, shard->nr_cids);
168
169	cu = &pshard->caps_updated;
170	raw_spin_lock_init(&cu->lock);
171	INIT_LIST_HEAD(&cu->node_in_flight);
172	__scx_cmask_init(&cu->cmask, shard->base_cid, shard->nr_cids, SCX_CID_SHARD_MAX_CPUS);
173
174	cu->cmask_arena_out = scx_arena_alloc(sch, cmask_size);
175	if (!cu->cmask_arena_out) {
176	free_pshard(pshard);
177	return NULL;
178	}
179
180	scx_cmask_init(cu->cmask_arena_out, shard->base_cid, shard->nr_cids);
181
182	return pshard;
183	}
184
185	s32 scx_alloc_pshards(struct scx_sched *sch)
186	{
187	struct scx_pshard **pshard;
188	s32 *shard_node;
189	s32 si;
190
191	if (!sch->is_cid_type \|\| !sch->arena_pool)
192	return 0;
193
194	shard_node = rcu_dereference_protected(scx_shard_node,
195	lockdep_is_held(&scx_enable_mutex));
196
197	pshard = kzalloc_objs(pshard[0], scx_nr_cid_shards);
198	if (!pshard)
199	return -ENOMEM;
200
201	for (si = 0; si < scx_nr_cid_shards; si++) {
202	pshard[si] = alloc_pshard(sch, si, shard_node[si]);
203	if (!pshard[si]) {
204	while (--si >= 0)
205	free_pshard(pshard[si]);
206	kfree(pshard);
207	return -ENOMEM;
208	}
209	}
210
211	sch->nr_pshards = scx_nr_cid_shards;
212	/*
213	* Publish only after every entry is built so a reader observing
214	* @sch->pshard never sees a partially-filled array or unpublished cid
215	* tables. Pair the store with a barrier and an acquire load on the
216	* read side.
217	*/
218	smp_wmb();
219	WRITE_ONCE(sch->pshard, pshard);
220	return 0;
221	}
222
223	/*
224	* Seed the root's caps fully. Root owns all cids on all caps at enable time.
225	* Children acquire caps via scx_bpf_sub_grant().
226	*/
227	void scx_init_root_caps(struct scx_sched *sch)
228	{
229	s32 si, i;
230
231	for (si = 0; si < sch->nr_pshards; si++) {
232	struct scx_pshard *ps = sch->pshard[si];
233
234	for (i = 0; i < __SCX_NR_CAPS; i++)
235	scx_cmask_fill(&ps->caps[i].cmask);
236	}
237	}
238
239	/* unserved remainder of @rq's rescuee's admitted slice, 0 once fully served */
240	static s64 scx_rescue_slice_remaining(struct rq *rq)
241	{
242	s64 served = rq->scx.rescue.curr->se.sum_exec_runtime - rq->scx.rescue.exec_snap;
243
244	return max(rq->scx.rescue.slice - served, 0);
245	}
246
247	/*
248	* Decay @pcpu's rescue usage average in place, halving per the knob-derived
249	* halflife, see scx_rescue_set_knobs(). The timestamp advances only by whole
250	* halflives.
251	*/
252	static u64 scx_rescue_decay_avg(struct scx_sched_pcpu *pcpu)
253	{
254	unsigned long halflife = scx_rescue_decay_halflife;
255	u64 n = div_u64(get_jiffies_64() - pcpu->rescue_avg_at, halflife);
256
257	if (n) {
258	pcpu->rescue_avg = n < 64 ? pcpu->rescue_avg >> n : 0;
259	pcpu->rescue_avg_at += n * halflife;
260	}
261	return pcpu->rescue_avg;
262	}
263
264	/**
265	* scx_rescue_charge - Charge the rescuee's runtime
266	* @rq: rq the rescuee is running on
267	* @delta_exec: runtime being charged
268	*
269	* Also ends the rescue once the admitted slice has been served in full. Ending
270	* on served time rather than slice exhaustion bounds both the rescue and the
271	* charging when a scheduler extends the rescuee's slice.
272	*/
273	void scx_rescue_charge(struct rq *rq, s64 delta_exec)
274	{
275	struct scx_sched_pcpu *pcpu;
276
277	lockdep_assert_rq_held(rq);
278
279	/*
280	* A rescue slice is bounded by one quantum and tick-driven expiry can
281	* overshoot by up to a tick. Clamp to avoid wild over-charges on VMs.
282	*/
283	delta_exec = min_t(s64, delta_exec, scx_rescue_quantum_ns + TICK_NSEC);
284
285	rq->scx.rescue.budget -= delta_exec;
286
287	/* per-cpu usage average feeds the overload victim pick */
288	pcpu = per_cpu_ptr(scx_task_sched(rq->curr)->pcpu, cpu_of(rq));
289	pcpu->rescue_avg = scx_rescue_decay_avg(pcpu) + delta_exec;
290
291	if (!scx_rescue_slice_remaining(rq))
292	scx_task_slice_ended(rq, rq->scx.rescue.curr);
293	}
294
295	/**
296	* scx_rescue_end - End the rescue execution on @rq
297	* @rq: rq of interest
298	*
299	* When no rescuee is left pending, the session is over and the balance above
300	* one quantum dies with it - it would otherwise become a banked license to
301	* preempt the cid owner long after the starvation ended. While waiters remain,
302	* the accrued deficit belongs to the queue and carries into the next rescue.
303	*/
304	void scx_rescue_end(struct rq *rq)
305	{
306	lockdep_assert_rq_held(rq);
307
308	rq->scx.rescue.curr = NULL;
309	if (list_empty(&rq->scx.rescue.dsq.list))
310	rq->scx.rescue.budget = min(rq->scx.rescue.budget, scx_rescue_quantum_ns);
311	}
312
313	/**
314	* scx_rescue_keep - Keep the rescue going for a preempted-out rescuee
315	* @rq: rq @p is running on
316	* @p: task under rescue whose slice is exhausted
317	*
318	* Called from put_prev_task_scx() to decide what an exhausted slice means for
319	* the rescuee. scx_rescue_charge() ends the rescue the moment the admitted
320	* slice is fully served, so arriving here with the rescue still open means @p
321	* was preempted. Restore the unserved remainder and return %true - @p stays the
322	* rescuee and the caller reinserts it at the tail of the local DSQ, behind
323	* whatever preempted the rescuee.
324	*
325	* Return %false to end the rescue instead - the slice is already fully served,
326	* @p is leaving the rq or bypass is dismantling rescues.
327	*/
328	bool scx_rescue_keep(struct rq rq, struct task_struct p)
329	{
330	s64 remaining = scx_rescue_slice_remaining(rq);
331
332	lockdep_assert_rq_held(rq);
333
334	if (!remaining \|\| !(p->scx.flags & SCX_TASK_QUEUED) \|\|
335	scx_bypassing(scx_task_sched(p), cpu_of(rq)))
336	return false;
337
338	scx_set_task_slice(p, remaining);
339	return true;
340	}
341
342	/**
343	* scx_rescue_accrue - Accrue budget at the configured fraction of elapsed time
344	* @rq: rq of interest
345	*
346	* A session spans from the first arrival until no rescuee is left, pending or
347	* admitted. While one is active the cap is three quanta and the balance drives
348	* escalation, see scx_rescue_timerfn(). Outside a session the cap is one
349	* quantum, so an idle gap funds the next arrival's admission but never an
350	* escalation.
351	*/
352	static void scx_rescue_accrue(struct rq *rq)
353	{
354	bool in_session = rq->scx.rescue.curr \|\| !list_empty(&rq->scx.rescue.dsq.list);
355	s64 cap = in_session ? 3 * scx_rescue_quantum_ns : scx_rescue_quantum_ns;
356	s64 delta;
357	u64 now;
358
359	lockdep_assert_rq_held(rq);
360
361	/* not every path here holds an updated rq clock, use __scx_bpf_now() */
362	now = __scx_bpf_now(rq);
363	delta = now - rq->scx.rescue.clock;
364	rq->scx.rescue.clock = now;
365
366	/*
367	* Avoid multiplication overflows by taking a shortcut when the gap is
368	* large enough to fill the budget.
369	*/
370	if (delta >= scx_rescue_sat_delta_ns)
371	rq->scx.rescue.budget = cap;
372	else
373	rq->scx.rescue.budget =
374	min(cap, rq->scx.rescue.budget +
375	((delta * scx_rescue_bw_1024) >> SCHED_CAPACITY_SHIFT));
376	}
377
378	/*
379	* The slice for the next admission - the quantum divided across the stranded
380	* tasks so that a crowded queue round-robins on shorter slices.
381	*/
382	static s64 scx_rescue_next_slice(struct rq *rq)
383	{
384	s64 min_slice = max_t(s64, SCX_RESCUE_MIN_SLICE_US * NSEC_PER_USEC, TICK_NSEC);
385	u32 depth = rq->scx.rescue.dsq.nr ?: 1;
386
387	return clamp(div_s64(scx_rescue_quantum_ns, depth), min_slice, scx_rescue_quantum_ns);
388	}
389
390	static void scx_rescue_timer_arm(struct rq *rq)
391	{
392	struct timer_list *timer = &rq->scx.rescue.timer;
393	s64 delay = scx_rescue_quantum_ns / 4; /* should be granular enough */
394
395	if (timer_pending(timer))
396	return;
397
398	/*
399	* While the head waiter can't be admitted because the bucket is short
400	* of a full quantum, stretch to the full funding delay.
401	*/
402	if (!rq->scx.rescue.curr && rq->scx.rescue.budget < scx_rescue_quantum_ns) {
403	s64 deficit = scx_rescue_quantum_ns - rq->scx.rescue.budget;
404
405	delay = max(delay,
406	div_s64(deficit << SCHED_CAPACITY_SHIFT, scx_rescue_bw_1024));
407	}
408
409	/* +1 rounds up so the beat is due by the time the timer fires */
410	timer->expires = jiffies + nsecs_to_jiffies(delay) + 1;
411	add_timer_on(timer, cpu_of(rq));
412	}
413
414	/**
415	* scx_rescue_admit - Start rescuing @p on @rq
416	* @rq: rq @p is being admitted on
417	* @p: task being admitted, off any DSQ
418	* @slice: CPU time to grant
419	*
420	* The schedulers keep their normal control over @p and may preempt or reslice
421	* it. @slice is measured on served CPU time against the snapshot taken here, so
422	* neither shortens the rescue, see scx_rescue_charge() and scx_rescue_keep().
423	* Prolonged denial escalates into protected execution, see
424	* scx_rescue_timerfn().
425	*/
426	static void scx_rescue_admit(struct rq rq, struct task_struct p, s64 slice)
427	{
428	lockdep_assert_rq_held(rq);
429	WARN_ON_ONCE(rq->scx.rescue.curr);
430
431	rq->scx.rescue.curr = p;
432	rq->scx.rescue.slice = slice;
433	rq->scx.rescue.exec_snap = p->se.sum_exec_runtime;
434	scx_set_task_slice(p, slice);
435	scx_rescue_timer_arm(rq);
436	}
437
438	/**
439	* scx_rescue_try_admit - Try to admit a freshly stranded task
440	* @rq: rq @p is being inserted on
441	* @p: stranded task being diverted to rescue
442	*
443	* One rescue at a time and earlier arrivals go first. Admission needs a full
444	* quantum of budget, spent as the rescue runs. Return %true if @p was admitted
445	* and should be inserted at the tail of @rq's local DSQ, %false if it has to
446	* park on the rescue DSQ, with the timer armed to admit it later.
447	*/
448	static bool scx_rescue_try_admit(struct rq rq, struct task_struct p)
449	{
450	scx_rescue_accrue(rq);
451
452	if (!rq->scx.rescue.curr && list_empty(&rq->scx.rescue.dsq.list) &&
453	rq->scx.rescue.budget >= scx_rescue_quantum_ns) {
454	scx_rescue_admit(rq, p, scx_rescue_quantum_ns);
455	return true;
456	}
457
458	scx_rescue_timer_arm(rq);
459	return false;
460	}
461
462	/**
463	* scx_rescue_check_overload - Eject the top rescue consumer on a stuck rescue
464	* @rq: rq whose rescue timer fired
465	*
466	* If the oldest waiter on @rq's rescue DSQ has been queued for too long, rescue
467	* demand on this cpu persistently exceeds the configured bandwidth. Eject the
468	* sub with the highest recent rescue consumption instead of letting the
469	* scheduler stall path blame the waiter's owner, who may just be crowded out.
470	*/
471	static void scx_rescue_check_overload(struct rq *rq)
472	{
473	struct scx_sched victim = NULL, pos;
474	struct task_struct *p;
475	int cpu = cpu_of(rq);
476	u64 max_avg = 0;
477	u32 dur_ms;
478
479	lockdep_assert_rq_held(rq);
480
481	p = list_first_entry_or_null(&rq->scx.rescue.dsq.list, struct task_struct,
482	scx.dsq_list.node);
483	if (!p)
484	return;
485
486	/* has the head waiter been queued for longer than the threshold? */
487	if (time_before(jiffies, p->scx.rescue_at + scx_rescue_overload_after))
488	return;
489
490	/*
491	* Grace period after the last ejection on this cpu - the freed
492	* bandwidth gets one threshold's worth of time to drain the backlog
493	* before another sub is judged.
494	*/
495	if (time_before64(get_jiffies_64(), rq->scx.rescue.kill_at +
496	scx_rescue_overload_after))
497	return;
498
499	list_for_each_entry_rcu(pos, &scx_sched_all, all) {
500	u64 avg = scx_rescue_decay_avg(per_cpu_ptr(pos->pcpu, cpu));
501
502	/* skip an already-exiting sub, else the ejection is wasted */
503	if (pos->level && avg > max_avg &&
504	atomic_read(&pos->exit_kind) == SCX_EXIT_NONE) {
505	max_avg = avg;
506	victim = pos;
507	}
508	}
509	if (!victim)
510	return;
511
512	rq->scx.rescue.kill_at = get_jiffies_64();
513	dur_ms = jiffies_to_msecs(jiffies - p->scx.rescue_at);
514	__scx_exit(victim, SCX_EXIT_ERROR_RESCUE, 0, cpu,
515	"used too much rescue CPU time (%llums) while %s[%d] waited %u.%03us to be rescued",
516	div_u64(max_avg, NSEC_PER_MSEC), p->comm, p->pid, dur_ms / 1000,
517	dur_ms % 1000);
518	}
519
520	/**
521	* scx_rescue_timerfn - Drive and pace rescue execution
522	* @timer: rq->scx.rescue.timer
523	*
524	* Runs every quarter quantum while a rescuee exists, pending or admitted, see
525	* scx_rescue_timer_arm(). The head waiter is admitted once the bucket holds a
526	* full quantum and granted its slice, see scx_rescue_next_slice(). A session
527	* whose budget accumulates over two quanta with the admitted rescuee still
528	* waiting escalates - the rescuee's remaining slice turns into protected
529	* execution and it preempts the current task. An overloaded rescue queue ejects
530	* the top consumer, see scx_rescue_check_overload().
531	*/
532	static void scx_rescue_timerfn(struct timer_list *timer)
533	{
534	struct rq *rq = timer_container_of(rq, timer, scx.rescue.timer);
535	struct task_struct *p;
536
537	guard(rq_lock_irqsave)(rq);
538
539	p = rq->scx.rescue.curr;
540	if (!p && list_empty(&rq->scx.rescue.dsq.list))
541	return;
542
543	scx_rescue_accrue(rq);
544	scx_rescue_check_overload(rq);
545
546	if (!p) {
547	s64 slice = scx_rescue_next_slice(rq);
548
549	/* no rescue in progress */
550	if (rq->scx.rescue.budget < scx_rescue_quantum_ns)
551	goto out_arm;
552
553	/* there's enough budget to start rescuing the next one */
554	p = list_first_entry(&rq->scx.rescue.dsq.list, struct task_struct,
555	scx.dsq_list.node);
556	scx_task_unlink_from_dsq(p, &rq->scx.rescue.dsq);
557	scx_rescue_admit(rq, p, slice);
558	scx_move_local_task_to_local_dsq(scx_task_sched(p), p,
559	SCX_ENQ_IGNORE_CAPS, rq);
560	if (sched_class_above(&ext_sched_class, rq->curr->sched_class))
561	resched_curr(rq);
562	} else if (p->scx.dsq && rq->scx.rescue.budget > 2 * scx_rescue_quantum_ns) {
563	/*
564	* The rescuee waited for the CPU for too long. Escalate - grant
565	* the unserved remainder, protect it from the schedulers and
566	* preempt the current task. The slice is set before the
567	* protection. Repeat beats only repeat the head move - the
568	* slice write is refused on a protected task.
569	*/
570	scx_set_task_slice(p, scx_rescue_slice_remaining(rq));
571	p->scx.flags \|= SCX_TASK_PROTECTED;
572	scx_task_unlink_from_dsq(p, &rq->scx.local_dsq);
573	scx_move_local_task_to_local_dsq(scx_task_sched(p), p,
574	SCX_ENQ_HEAD \| SCX_ENQ_PREEMPT \| SCX_ENQ_IGNORE_CAPS,
575	rq);
576	}
577	out_arm:
578	scx_rescue_timer_arm(rq);
579	}
580
581	/* flush out tasks waiting for rescue before a CPU goes down */
582	void scx_rescue_flush(struct rq *rq)
583	{
584	struct task_struct p, n;
585
586	lockdep_assert_rq_held(rq);
587
588	/* end the current rescue */
589	if (rq->scx.rescue.curr)
590	scx_task_slice_ended(rq, rq->scx.rescue.curr);
591
592	/* and flush out all pending ones */
593	list_for_each_entry_safe(p, n, &rq->scx.rescue.dsq.list, scx.dsq_list.node) {
594	scx_task_unlink_from_dsq(p, &rq->scx.rescue.dsq);
595	scx_move_local_task_to_local_dsq(scx_task_sched(p), p,
596	SCX_ENQ_IGNORE_CAPS, rq);
597	}
598
599	timer_delete(&rq->scx.rescue.timer);
600	}
601
602	void scx_rescue_dump(struct seq_buf s, struct rq rq)
603	{
604	struct task_struct *p = rq->scx.rescue.curr;
605
606	scx_dump_line(s, " rescue=%u budget=%lldus rescuing=%s[%d]",
607	rq->scx.rescue.dsq.nr,
608	div_s64(rq->scx.rescue.budget, NSEC_PER_USEC),
609	p ? p->comm : "none", p ? p->pid : -1);
610	}
611
612	/*
613	* A scheduler whose stall watchdog is shorter than the overload threshold gets
614	* stall-killed over its parked waiters before the overload check can eject the
615	* actual top consumer. The root's knobs set the threshold, warn on any
616	* scheduler that doesn't fit it.
617	*/
618	static void scx_rescue_check_timeout(struct scx_sched *sch)
619	{
620	if (!scx_rescue_bw_1024 \|\| sch->watchdog_timeout > scx_rescue_overload_after)
621	return;
622
623	pr_warn("sched_ext: %s: watchdog timeout %ums <= rescue overload threshold %ums\n",
624	sch->ops.name, jiffies_to_msecs(sch->watchdog_timeout),
625	jiffies_to_msecs(scx_rescue_overload_after));
626	}
627
628	/* latch the rescue parameters on root scheduler enable */
629	void scx_rescue_set_knobs(struct scx_sched *sch)
630	{
631	s32 bw_ppt = sch->ops.rescue_bandwidth_ppt ?: SCX_RESCUE_DFL_BW_PPT;
632	s64 quantum_us = sch->ops.rescue_quantum_us ?: SCX_RESCUE_DFL_QUANTUM_US;
633	s64 period_ns;
634
635	if (sch->ops.rescue_bandwidth_ppt == SCX_RESCUE_DISABLE) {
636	scx_rescue_bw_1024 = 0;
637	return;
638	}
639
640	scx_rescue_bw_1024 = bw_ppt * SCHED_CAPACITY_SCALE / 1000;
641	scx_rescue_quantum_ns = max(quantum_us * NSEC_PER_USEC, TICK_NSEC);
642	scx_rescue_sat_delta_ns =
643	div_s64((4 * scx_rescue_quantum_ns + TICK_NSEC) << SCHED_CAPACITY_SHIFT,
644	scx_rescue_bw_1024);
645
646	/*
647	* The overload threshold and the decay halflife scale with the funding
648	* period - the time the bucket takes to fund one full quantum.
649	*/
650	period_ns = div_s64(scx_rescue_quantum_ns << SCHED_CAPACITY_SHIFT, scx_rescue_bw_1024);
651	scx_rescue_overload_after =
652	clamp(nsecs_to_jiffies(SCX_RESCUE_OVERLOAD_MULT * period_ns),
653	msecs_to_jiffies(SCX_RESCUE_MIN_OVERLOAD_MS),
654	msecs_to_jiffies(SCX_RESCUE_MAX_OVERLOAD_MS));
655	scx_rescue_decay_halflife = scx_rescue_overload_after / 4;
656
657	/* a single in-budget wait must not cross the overload trigger */
658	if (nsecs_to_jiffies(period_ns) > scx_rescue_overload_after / 2)
659	pr_warn("sched_ext: %s: rescue funding period %lldms > overload threshold %ums / 2\n",
660	sch->ops.name, div_s64(period_ns, NSEC_PER_MSEC),
661	jiffies_to_msecs(scx_rescue_overload_after));
662
663	scx_rescue_check_timeout(sch);
664	}
665
666	void scx_rescue_init(struct rq *rq)
667	{
668	BUG_ON(scx_init_dsq(&rq->scx.rescue.dsq, SCX_DSQ_RESCUE, NULL));
669	timer_setup(&rq->scx.rescue.timer, scx_rescue_timerfn, TIMER_PINNED);
670	rq->scx.rescue.kill_at = get_jiffies_64();
671	}
672
673	/**
674	* scx_resolve_local_dsq - Pick the local, rescue or reject DSQ for an insert
675	* @sch: enqueuing sub-sched
676	* @rq: rq whose local DSQ @p targets
677	* @p: task being inserted
678	* @enq_flags: in/out, unhonored flags are cleared
679	*
680	* Return @rq's local DSQ if @sch holds the required caps on @rq's cid.
681	* Otherwise, return @rq's rescue DSQ if the insert carries %SCX_ENQ_RESCUE and
682	* rescue is enabled, or @rq's reject DSQ after recording the reenq reason on
683	* @p.
684	*
685	* %SCX_ENQ_IMMED, %SCX_ENQ_PREEMPT and %SCX_ENQ_HEAD are cleared when diverting
686	* to rescue or reject. %SCX_ENQ_PREEMPT is also cleared on a fallback
687	* migration-disabled admission.
688	*
689	* Bypass doesn't need special-casing as a bypassing sched's tasks are enqueued
690	* to and run by its nearest non-bypassing ancestor. If root is bypassing, it
691	* always holds all caps.
692	*/
693	struct scx_dispatch_q scx_resolve_local_dsq(struct scx_sched sch, struct rq *rq,
694	struct task_struct p, u64 enq_flags)
695	{
696	if (!scx_has_subs())
697	return &rq->scx.local_dsq;
698
699	s32 cid = __scx_cpu_to_cid(cpu_of(rq));
700	struct scx_sched *asch = rq->scx.remote_activate_sch ?: sch;
701	u64 needed = scx_caps_for_enq(*enq_flags);
702	u64 missing;
703
704	/*
705	* On a remote activation the scheduling sched (@asch) differs from
706	* @p's owner (@sch). Check caps against the scheduling sched.
707	*/
708	if (*enq_flags & SCX_ENQ_PREEMPT)
709	needed \|= scx_caps_for_preempt(asch, rq, *enq_flags);
710	missing = scx_missing_caps(asch, cpu_of(rq), needed);
711
712	/* requirements met */
713	if (likely(!missing))
714	return &rq->scx.local_dsq;
715
716	/*
717	* The task must run on this CPU regardless of caps: the rq is draining
718	* offline (BPF scheduler bypassed), the task is migration-disabled, or a
719	* migration is pending. Admit despite the missing caps and count it.
720	* Refuse preemptions.
721	*/
722	if (unlikely(!scx_rq_online(rq) \|\| is_migration_disabled(p) \|\|
723	p->migration_pending)) {
724	__scx_add_event(sch, SCX_EV_SUB_FORCED_ADMIT, 1);
725	*enq_flags &= ~SCX_ENQ_PREEMPT;
726	return &rq->scx.local_dsq;
727	}
728
729	/*
730	* Diverting to rescue or reject, neither of which honors IMMED, PREEMPT
731	* or HEAD - a diversion has no priority and IMMED is not allowed on
732	* non-local DSQs. Strip the enq and task flags along with the slice.
733	*/
734	*enq_flags &= ~(SCX_ENQ_IMMED \| SCX_ENQ_PREEMPT \| SCX_ENQ_HEAD \|
735	SCX_ENQ_APPLY_SLICE \| SCX_ENQ_SLICE_DFL);
736	p->scx.flags &= ~SCX_TASK_IMMED;
737
738	/* the enqueuer opted for rescue instead of rejection and reenqueue */
739	if ((*enq_flags & SCX_ENQ_RESCUE) && likely(scx_rescue_bw_1024)) {
740	__scx_add_event(sch, SCX_EV_SUB_RESCUE, 1);
741	if (scx_rescue_try_admit(rq, p))
742	return &rq->scx.local_dsq;
743
744	/* queueing, the overload trigger measures the wait from here */
745	p->scx.rescue_at = jiffies;
746	return &rq->scx.rescue.dsq;
747	}
748
749	p->scx.reenq_reason_caps = missing;
750	p->scx.reenq_reason_cid = cid;
751
752	return &rq->scx.reject_dsq;
753	}
754
755	/* @p lost the caps needed to stay on @rq's local DSQ? Record reason if so. */
756	bool scx_task_reenq_on_cap_revoke(struct rq rq, struct task_struct p)
757	{
758	u64 missing;
759
760	/* migration-disabled tasks and the rescuee are admitted capless */
761	if (is_migration_disabled(p) \|\| p == scx_rescuee(rq))
762	return false;
763
764	missing = scx_missing_caps(scx_task_sched(p), cpu_of(rq), scx_caps_for_task(p));
765	if (likely(!missing))
766	return false;
767
768	p->scx.reenq_reason_caps = missing;
769	p->scx.reenq_reason_cid = __scx_cpu_to_cid(cpu_of(rq));
770	return true;
771	}
772
773	/*
774	* Drain @rq->scx.reject_dsq, reenqueueing each task so the BPF re-decides
775	* from p->scx.reenq_reason_*.
776	*
777	* A task can be re-rejected repeatedly. The reenqueue is bounded per task in
778	* scx_do_enqueue_task(), which ejects the owning sub past SCX_REENQ_MAX_REPEAT.
779	* Rejection can't happen for root.
780	*/
781	void scx_reenq_reject(struct rq *rq)
782	{
783	LIST_HEAD(tasks);
784	struct task_struct p, n;
785
786	lockdep_assert_rq_held(rq);
787
788	if (!scx_has_subs() \|\| list_empty(&rq->scx.reject_dsq.list))
789	return;
790
791	/*
792	* Move to a private list so a task re-rejected by the
793	* scx_do_enqueue_task() below isn't revisited this round.
794	*/
795	list_for_each_entry_safe(p, n, &rq->scx.reject_dsq.list, scx.dsq_list.node) {
796	/* migration_pending tasks should have bypassed to local DSQ */
797	if (WARN_ON_ONCE(p->migration_pending))
798	continue;
799
800	scx_wait_dispatching(p);
801	scx_dispatch_dequeue(rq, p);
802
803	if (WARN_ON_ONCE(p->scx.flags & SCX_TASK_REENQ_REASON_MASK))
804	p->scx.flags &= ~SCX_TASK_REENQ_REASON_MASK;
805	p->scx.flags \|= SCX_TASK_REENQ_CAP;
806
807	list_add_tail(&p->scx.dsq_list.node, &tasks);
808	}
809
810	list_for_each_entry_safe(p, n, &tasks, scx.dsq_list.node) {
811	list_del_init(&p->scx.dsq_list.node);
812
813	scx_do_enqueue_task(rq, p, SCX_ENQ_REENQ, -1);
814
815	p->scx.flags &= ~SCX_TASK_REENQ_REASON_MASK;
816	}
817	}
818
819	/* record a caps change, see struct scx_caps_updated */
820	static void caps_updated_record(struct scx_pshard ps, const struct scx_cmask cids, u64 caps,
821	struct list_head *to_deliver)
822	{
823	struct scx_caps_updated *cu = &ps->caps_updated;
824
825	guard(raw_spinlock)(&cu->lock);
826	scx_cmask_or(&cu->cmask, cids);
827	cu->caps \|= caps;
828	if (list_empty(&cu->node_in_flight))
829	list_add_tail(&cu->node_in_flight, to_deliver);
830	}
831
832	/* deliver queued caps_updated callbacks, see struct scx_caps_updated */
833	static void caps_updated_deliver(struct list_head *to_deliver)
834	{
835	struct scx_caps_updated cu, tmp;
836
837	list_for_each_entry_safe(cu, tmp, to_deliver, node_in_flight) {
838	struct scx_pshard *ps = container_of(cu, struct scx_pshard, caps_updated);
839	struct scx_sched *sch = ps->sch;
840
841	while (true) {
842	u64 caps = 0;
843
844	/*
845	* During enable, has_op is set after ops.sub_attach(),
846	* so !has_op means the op is absent or the sched isn't
847	* live yet - e.g. caps grant from ops.sub_attach().
848	* Either way don't consume - leave for
849	* scx_sub_seed_caps() to deliver once live.
850	*/
851	scoped_guard (raw_spinlock, &cu->lock) {
852	if (cu->caps && SCX_HAS_OP(sch, sub_caps_updated) &&
853	likely(!READ_ONCE(sch->aborting))) {
854	struct scx_cmask_ref ref;
855
856	caps = cu->caps;
857	scx_cmask_ref_init_kern(sch, cu->cmask_arena_out,
858	ps->base, ps->nr_cids, &ref);
859	scx_cmask_ref_copy(&ref, &cu->cmask);
860	scx_cmask_clear(&cu->cmask);
861	cu->caps = 0;
862	} else {
863	list_del_init(&cu->node_in_flight);
864	}
865	}
866	if (!caps)
867	break;
868
869	/* caps != 0 only when deliverable (has_op, above) */
870	SCX_CALL_OP(sch, sub_caps_updated, NULL, cu->cmask_arena_out, caps);
871	}
872	}
873	}
874
875	/*
876	* Deliver caps owed to @sch that couldn't be delivered earlier (e.g. a grant
877	* taken during its sub_attach(), before has_op was set). Called once @sch is
878	* enabled.
879	*/
880	static void scx_sub_seed_caps(struct scx_sched *sch)
881	{
882	LIST_HEAD(to_deliver);
883	s32 si;
884
885	guard(irqsave)();
886
887	for (si = 0; si < sch->nr_pshards; si++) {
888	struct scx_pshard *ps = sch->pshard[si];
889	struct scx_caps_updated *cu = &ps->caps_updated;
890
891	scoped_guard (raw_spinlock, &cu->lock) {
892	if (cu->caps && list_empty(&cu->node_in_flight))
893	list_add_tail(&cu->node_in_flight, &to_deliver);
894	}
895	}
896	caps_updated_deliver(&to_deliver);
897	}
898
899	static u64 calc_effective_caps(struct scx_pshard *ps, s32 cid)
900	{
901	u64 ecaps = 0;
902	u32 cap_bit;
903
904	for (cap_bit = 0; cap_bit < __SCX_NR_CAPS; cap_bit++)
905	if (scx_cmask_test(cid, &ps->caps[cap_bit].cmask))
906	ecaps \|= BIT_U64(cap_bit) \| scx_caps_implied(BIT_U64(cap_bit));
907	return ecaps;
908	}
909
910	/**
911	* queue_sync_ecaps - Queue ecaps update for a (sch, cid) pair
912	* @sch: sched to update
913	* @cid: cid to update
914	*
915	* Queue an ecaps update for @sch's @cid and kick the cpu so that it syncs in
916	* dispatch_one().
917	*/
918	static void queue_sync_ecaps(struct scx_sched *sch, s32 cid)
919	{
920	s32 cpu = __scx_cid_to_cpu(cid);
921	struct scx_sched_pcpu *pcpu = per_cpu_ptr(sch->pcpu, cpu);
922
923	/*
924	* Pairs with smp_mb() in scx_process_sync_ecaps(). Either the check
925	* below sees the node off the list and queues it, or the in-flight sync
926	* sees the caps[] update made before this call.
927	*/
928	smp_mb();
929
930	/* @cid's pshard->lock excludes concurrent queueing attempts */
931	if (llist_on_list(&pcpu->ecaps_to_sync_node))
932	return;
933	if (llist_add(&pcpu->ecaps_to_sync_node, &cpu_rq(cpu)->scx.ecaps_to_sync))
934	scx_kick_cpu(sch->ancestors[0], cpu, 0);
935	}
936
937	/* discard @rq's queued ecaps syncs */
938	static void discard_queued_syncs(struct rq *rq)
939	{
940	struct llist_node pos, tmp;
941
942	lockdep_assert_rq_held(rq);
943
944	llist_for_each_safe(pos, tmp, llist_del_all(&rq->scx.ecaps_to_sync))
945	init_llist_node(pos);
946	}
947
948	/**
949	* scx_process_sync_ecaps - Sync this cpu's ecaps to pshard->caps[]
950	* @rq: the cid's cpu rq
951	* @prev: @rq's previous task from the in-progress dispatch
952	*
953	* pshard->caps[] is the target configuration. pcpu->ecaps is the effective
954	* transposed copy owned by the cid's cpu and written only here under @rq's
955	* lock.
956	*
957	* A sched that newly gains baseline access here is owed an update_idle() so it
958	* learns the cid's idle state. Such a gain arms the per-rq
959	* %SCX_RQ_SUB_IDLE_RENOTIFY gate so the next idle pick delivers it.
960	*/
961	void scx_process_sync_ecaps(struct rq rq, struct task_struct prev)
962	{
963	s32 cpu = cpu_of(rq);
964	s32 cid, shard;
965	struct llist_node batch, pos, *tmp;
966	u64 lost_all = 0;
967
968	lockdep_assert_rq_held(rq);
969
970	if (!scx_has_subs() \|\| likely(llist_empty(&rq->scx.ecaps_to_sync)))
971	return;
972
973	/*
974	* ecaps are zeroed while the cpu is inactive and must stay zero.
975	* Discard queued syncs instead of processing them - the
976	* scx_online_ecaps() reseed re-syncs every sched on activation.
977	* cpu_active() clears before the offline zeroing and sets before the
978	* reseed is queued, so this test can neither miss a racing sync nor
979	* eat the reseed.
980	*/
981	if (unlikely(!cpu_active(cpu))) {
982	discard_queued_syncs(rq);
983	return;
984	}
985
986	/* @cid is valid here: the cpu is active with queued syncs */
987	cid = __scx_cpu_to_cid(cpu);
988	shard = rcu_dereference_all(scx_cid_to_shard)[cid];
989
990	batch = llist_del_all(&rq->scx.ecaps_to_sync);
991	llist_for_each_safe(pos, tmp, batch) {
992	struct scx_sched_pcpu *pcpu =
993	container_of(pos, struct scx_sched_pcpu, ecaps_to_sync_node);
994	struct scx_pshard *ps = pcpu->sch->pshard[shard];
995	u64 old, ecaps, lost, gained;
996
997	init_llist_node(pos);
998
999	/* pairs with smp_mb() in queue_sync_ecaps(), see there */
1000	smp_mb();
1001
1002	old = READ_ONCE(pcpu->ecaps);
1003	ecaps = calc_effective_caps(ps, cid);
1004	WRITE_ONCE(pcpu->ecaps, ecaps);
1005
1006	lost = old & ~ecaps;
1007	gained = ecaps & ~old;
1008	lost_all \|= lost;
1009
1010	/*
1011	* Tell the sched its effective caps on this cid changed. The
1012	* invocation is equivalent to the dispatch path and may drop
1013	* and re-acquire the rq lock temporarily while the rest of
1014	* @batch is held privately, see scx_discard_ecaps_to_sync().
1015	* The dispatch kfuncs resolve their context on the executing
1016	* cpu, which under core scheduling can differ from @rq's cpu,
1017	* so the context is set up there. The rq recorded in it keeps
1018	* the dispatches targeting @rq.
1019	*/
1020	if (ecaps != pcpu->reported_ecaps &&
1021	SCX_HAS_OP(pcpu->sch, sub_ecaps_updated) &&
1022	!scx_bypassing(pcpu->sch, cpu)) {
1023	struct scx_dsp_ctx *dspc = &this_cpu_ptr(pcpu->sch->pcpu)->dsp_ctx;
1024
1025	dspc->rq = rq;
1026	/* stash @prev so nested dispatches can access it */
1027	rq->scx.sub_dispatch_prev = prev;
1028	SCX_CALL_OP(pcpu->sch, sub_ecaps_updated, rq, scx_cpu_arg(cpu),
1029	pcpu->reported_ecaps, ecaps);
1030	rq->scx.sub_dispatch_prev = NULL;
1031	scx_flush_dispatch_buf(pcpu->sch, rq);
1032	pcpu->reported_ecaps = ecaps;
1033	}
1034
1035	/*
1036	* Gaining baseline access owes an update_idle() so the sched
1037	* learns the cpu's idle state. Arm the per-rq gate so the next
1038	* idle pick flushes it. Losing access drops any pending notify.
1039	*/
1040	if (gained & SCX_CAP_BASE) {
1041	pcpu->idle_renotify = true;
1042	rq->scx.flags \|= SCX_RQ_SUB_IDLE_RENOTIFY;
1043	} else if (lost & SCX_CAP_BASE) {
1044	pcpu->idle_renotify = false;
1045	}
1046	}
1047
1048	/*
1049	* Losing a cap can strand already-queued tasks. Schedule a reenq scan
1050	* to move the now-capless ones off the local DSQ. The scan tests
1051	* against the effective caps and thus must come after the ecaps sync.
1052	*/
1053	if (lost_all & SCX_CAPS_REENQ_ON_LOSS)
1054	scx_schedule_reenq_local(rq, SCX_REENQ_CAP_REVOKE);
1055	}
1056
1057	/**
1058	* scx_unbypass_replay_ecaps - Replay a bypass-suppressed ecaps notification
1059	* @rq: rq of the cpu leaving bypass
1060	* @sch: scheduler that just left bypass on @rq's cpu
1061	*
1062	* scx_process_sync_ecaps() consumes syncs while bypassing without delivering
1063	* ops.sub_ecaps_updated(), leaving reported_ecaps stale. Nothing re-queues a
1064	* sync when bypass lifts, so without a replay a cid that never changes again
1065	* would never be notified. The attach-time initial grants are the acute case
1066	* as they are consumed during the enable bypass window. Re-queue a sync for
1067	* any undelivered delta so the next dispatch delivers it.
1068	*/
1069	void scx_unbypass_replay_ecaps(struct rq rq, struct scx_sched sch)
1070	{
1071	s32 cpu = cpu_of(rq);
1072	struct scx_sched_pcpu *pcpu = per_cpu_ptr(sch->pcpu, cpu);
1073	struct scx_pshard *ps;
1074	s32 cid;
1075
1076	lockdep_assert_rq_held(rq);
1077
1078	/* root holds every cap and never uses ecaps */
1079	if (!sch->level)
1080	return;
1081
1082	if (READ_ONCE(pcpu->ecaps) == pcpu->reported_ecaps)
1083	return;
1084
1085	cid = __scx_cpu_to_cid(cpu);
1086	ps = sch->pshard[rcu_dereference_all(scx_cid_to_shard)[cid]];
1087
1088	guard(raw_spinlock)(&ps->lock);
1089	queue_sync_ecaps(sch, cid);
1090	}
1091
1092	/*
1093	* A cpu came back. Re-seed each sub-sched's ecaps on the cpu's cid. The sync
1094	* recomputes effective caps from the pshard and fires ops.sub_ecaps_updated()
1095	* only on a real change since offline.
1096	*/
1097	void scx_online_ecaps(struct rq *rq)
1098	{
1099	struct scx_sched root, pos;
1100	s32 cid, shard;
1101
1102	/*
1103	* Only a live hierarchy can have ecaps to reseed. This also keeps the
1104	* table reads below away from an enable that failed before publishing
1105	* the tables. A concurrent disable can't retire them, see
1106	* handle_hotplug().
1107	*/
1108	if (!scx_enabled())
1109	return;
1110
1111	guard(rq_lock_irqsave)(rq);
1112
1113	root = scx_root_protected();
1114	cid = __scx_cpu_to_cid(cpu_of(rq));
1115	shard = rcu_dereference_all(scx_cid_to_shard)[cid];
1116
1117	scx_for_each_descendant_pre(pos, root) {
1118	struct scx_pshard *ps;
1119
1120	/* root holds every cap and never uses ecaps */
1121	if (!pos->level)
1122	continue;
1123
1124	ps = pos->pshard[shard];
1125	guard(raw_spinlock)(&ps->lock);
1126	queue_sync_ecaps(pos, cid);
1127	}
1128	}
1129
1130	/*
1131	* A cpu is going down. Zero each sub-sched's in-effect ecaps so cap checks
1132	* treat the cpu as capless while offline. Pending and late-queued syncs are
1133	* discarded at consumption by scx_process_sync_ecaps() while the cpu is
1134	* inactive. Leave reported_ecaps. Ownership is unchanged, so the
1135	* scx_online_ecaps() reseed reports only a genuine delta. No callback fires
1136	* here.
1137	*/
1138	void scx_offline_ecaps(struct rq *rq)
1139	{
1140	s32 cpu = cpu_of(rq);
1141	struct scx_sched root, pos;
1142
1143	guard(rq_lock_irqsave)(rq);
1144
1145	root = scx_root_protected();
1146
1147	scx_for_each_descendant_pre(pos, root) {
1148	/* root holds every cap and never uses ecaps */
1149	if (!pos->level)
1150	continue;
1151
1152	WRITE_ONCE(per_cpu_ptr(pos->pcpu, cpu)->ecaps, 0);
1153	}
1154	}
1155
1156	/*
1157	* @pcpu's sched was unhashed before the grace period, so nothing re-queues its
1158	* sync node. Remove the node from @rq's pending list so the pcpu can be freed.
1159	*/
1160	void scx_discard_ecaps_to_sync(s32 cpu, struct scx_sched_pcpu *pcpu)
1161	{
1162	struct rq *rq = cpu_rq(cpu);
1163	struct llist_node head = NULL, tail = NULL;
1164	struct llist_node pos, tmp;
1165
1166	/*
1167	* llist can't unlink a single node. Take all queued nodes, drop @pcpu's
1168	* and resplice the rest. Nodes in the taken batch read as on-list
1169	* throughout, so queue_sync_ecaps() stays correct.
1170	*/
1171	if (llist_on_list(&pcpu->ecaps_to_sync_node)) {
1172	scoped_guard (rq_lock_irqsave, rq) {
1173	llist_for_each_safe(pos, tmp, llist_del_all(&rq->scx.ecaps_to_sync)) {
1174	if (pos == &pcpu->ecaps_to_sync_node) {
1175	init_llist_node(pos);
1176	} else {
1177	pos->next = head;
1178	head = pos;
1179	if (!tail)
1180	tail = pos;
1181	}
1182	}
1183	if (head)
1184	llist_add_batch(head, tail, &rq->scx.ecaps_to_sync);
1185	}
1186	}
1187
1188	/*
1189	* An in-flight scx_process_sync_ecaps() batch may still hold the node
1190	* privately across dispatch-induced rq unlocks, reading as on-list.
1191	*
1192	* Because a bypassing sched gets no op call, init_llist_node() and all
1193	* @pcpu accesses share one contiguous lock hold, off-list under the rq
1194	* lock means @pcpu won't be accessed again.
1195	*/
1196	while (true) {
1197	scoped_guard (rq_lock_irqsave, rq) {
1198	if (!llist_on_list(&pcpu->ecaps_to_sync_node))
1199	return;
1200	}
1201	cpu_relax();
1202	}
1203	}
1204
1205	/**
1206	* scx_discard_stale_ecaps_syncs - Discard ecaps syncs from earlier schedulers
1207	*
1208	* To be called during root enable before the scheduler goes live. An earlier
1209	* root's sub-sched may not have gone through its RCU free path yet (e.g. a
1210	* still-open link fd defers it) and can leave queued ecaps syncs behind.
1211	* Processing them would decode the dead sched's pshards with the current cid
1212	* layout. Discard them instead. The backing scx_sched_pcpu's are still
1213	* allocated as the free path removes ecaps_to_sync_node before freeing.
1214	*/
1215	void scx_discard_stale_ecaps_syncs(void)
1216	{
1217	s32 cpu;
1218
1219	for_each_possible_cpu(cpu) {
1220	struct rq *rq = cpu_rq(cpu);
1221
1222	guard(rq_lock_irqsave)(rq);
1223	discard_queued_syncs(rq);
1224	}
1225	}
1226
1227	static DECLARE_WAIT_QUEUE_HEAD(scx_unlink_waitq);
1228
1229	void drain_descendants(struct scx_sched *sch)
1230	{
1231	/*
1232	* Child scheds that finished the critical part of disabling will take
1233	* themselves off @sch->children. Wait for it to drain. As propagation
1234	* is recursive, empty @sch->children means that all proper descendant
1235	* scheds reached unlinking stage.
1236	*/
1237	wait_event(scx_unlink_waitq, list_empty(&sch->children));
1238	}
1239
1240	/**
1241	* scx_rehome_task - Move a task to a sched it has been initialized for
1242	* @to: sched taking over @p, @p's init on it already complete
1243	* @p: task to re-home
1244	*
1245	* Exit @p from its current sched and switch it over to @to, overriding the
1246	* state to %SCX_TASK_READY to account for the already completed init. A task
1247	* on a non-ext class, possible under an %SCX_OPS_SWITCH_PARTIAL root, stays
1248	* %READY and is enabled by switching_to_scx() if it switches over.
1249	*/
1250	static void scx_rehome_task(struct scx_sched to, struct task_struct p)
1251	{
1252	lockdep_assert_held(&p->pi_lock);
1253	lockdep_assert_rq_held(task_rq(p));
1254
1255	scoped_guard (sched_change, p, DEQUEUE_SAVE \| DEQUEUE_MOVE) {
1256	scx_disable_and_exit_task(scx_task_sched(p), p);
1257	scx_set_task_state(p, SCX_TASK_INIT_BEGIN);
1258	scx_set_task_state(p, SCX_TASK_INIT);
1259	scx_set_task_sched(p, to);
1260	scx_set_task_state(p, SCX_TASK_READY);
1261	if (p->sched_class == &ext_sched_class)
1262	scx_enable_task(to, p);
1263	}
1264	}
1265
1266	/**
1267	* scx_punt_task - Hand a task to a failed sched without initialization
1268	* @to: failed and bypassed sched taking custody of @p
1269	* @p: task to punt
1270	*
1271	* Take @p off its current sched and put it on @to at %SCX_TASK_NONE. @to is
1272	* dying and its teardown will re-home @p properly.
1273	*
1274	* Used when @to must take over @p but failed to initialize it. Bypass keeps
1275	* scheduling decisions away from @to but @p can still trigger its task ops,
1276	* which may confuse the BPF side. @to is dying anyway. The exit paths skip
1277	* %NONE tasks (see __scx_disable_and_exit_task() and switched_from_scx()).
1278	*/
1279	static void scx_punt_task(struct scx_sched to, struct task_struct p)
1280	{
1281	lockdep_assert_held(&p->pi_lock);
1282	lockdep_assert_rq_held(task_rq(p));
1283	WARN_ON_ONCE(!READ_ONCE(to->bypass_depth));
1284
1285	scoped_guard (sched_change, p, DEQUEUE_SAVE \| DEQUEUE_MOVE) {
1286	scx_disable_and_exit_task(scx_task_sched(p), p);
1287	scx_set_task_sched(p, to);
1288	}
1289	}
1290
1291	static void scx_fail_parent(struct scx_sched *sch,
1292	struct task_struct *failed, s32 fail_code)
1293	{
1294	struct scx_sched *parent = scx_parent(sch);
1295	struct scx_task_iter sti;
1296	struct task_struct *p;
1297
1298	scx_error(parent, "ops.init_task() failed (%d) for %s[%d] while disabling a sub-scheduler",
1299	fail_code, failed->comm, failed->pid);
1300
1301	/*
1302	* Once $parent is bypassed, tasks can be punted into it. This may
1303	* cause downstream failures on the BPF side but $parent is dying
1304	* anyway.
1305	*/
1306	scx_bypass(parent, true);
1307
1308	scx_task_iter_start(&sti, sch->cgrp);
1309	while ((p = scx_task_iter_next_locked(&sti))) {
1310	if (scx_task_on_sched(parent, p))
1311	continue;
1312
1313	scx_punt_task(parent, p);
1314	}
1315	scx_task_iter_stop(&sti);
1316	}
1317
1318	#ifdef CONFIG_EXT_GROUP_SCHED
1319	/**
1320	* scx_cgroup_claim_subtree - Claim the subtree's cgroups for an enabling sub
1321	* @sch: sub-scheduler being enabled
1322	*
1323	* Called while enabling @sch, after the subtree's cgrp->scx_sched's are pointed
1324	* at @sch and before any task is claimed. This mirrors root enable's
1325	* cgroups-before-tasks order. The ops.init_task() args are task_group-granular
1326	* and can still reference a cgroup outside the handed-over set when the cpu
1327	* controller is coarser than the sub topology or mounted on cgroup1.
1328	*
1329	* First init each of the parent sched's subtree cgroups on @sch, and only then
1330	* exit them from the parent, so that a failed init can be unwound with the
1331	* parent untouched. The both-inited transient is invisible outside
1332	* scx_cgroup_lock(). %SCX_TG_SUB_INIT tracks the first pass's progress.
1333	* %SCX_TG_INITED stays set throughout, except for a task_group whose
1334	* ops.cgroup_init() failed on the parent (see scx_cgroup_return_subtree()):
1335	* there is nothing to exit from the parent and %SCX_TG_INITED is set back with
1336	* the transfer.
1337	*
1338	* Dying but not yet offlined task_groups are included: a removed cgroup keeps
1339	* hosting scheduling events until its dying tasks finish their final context
1340	* switches, so it still needs to be inited on a sched, and its offline-time
1341	* ops.cgroup_exit() follows the last of those events.
1342	*
1343	* Return 0 on success, -errno on failure. On failure, @sch has been
1344	* scx_error()'d and is left with no cgroups.
1345	*/
1346	static s32 scx_cgroup_claim_subtree(struct scx_sched *sch)
1347	{
1348	struct cgroup *sub_cgrp = sch_cgroup(sch);
1349	struct cgroup_subsys_state *ecss = cgroup_e_css(sub_cgrp, &cpu_cgrp_subsys);
1350	struct scx_sched *parent = scx_parent(sch);
1351	struct cgroup_subsys_state *css;
1352	int ret;
1353
1354	css_for_each_descendant_pre(css, ecss) {
1355	struct task_group *tg = css_tg(css);
1356	struct scx_cgroup_init_args args = {
1357	.weight = tg->scx.weight,
1358	.bw_period_us = tg->scx.bw_period_us,
1359	.bw_quota_us = tg->scx.bw_quota_us,
1360	.bw_burst_us = tg->scx.bw_burst_us,
1361	.sched_idle = tg->scx.idle,
1362	};
1363
1364	if (tg->scx.sched != parent \|\|
1365	!cgroup_is_descendant(css->cgroup, sub_cgrp))
1366	continue;
1367
1368	if (SCX_HAS_OP(sch, cgroup_init)) {
1369	ret = SCX_CALL_OP_RET(sch, cgroup_init, NULL, css->cgroup, &args);
1370	if (ret) {
1371	scx_error(sch, "ops.cgroup_init() failed (%d)", ret);
1372	goto err;
1373	}
1374	}
1375	tg->scx.flags \|= SCX_TG_SUB_INIT;
1376	}
1377
1378	css_for_each_descendant_post(css, ecss) {
1379	struct task_group *tg = css_tg(css);
1380
1381	/*
1382	* SUB_INIT is pass 1's progress mark: pass 2 and the err path
1383	* must visit exactly the tgs pass 1 inited.
1384	*/
1385	if (!(tg->scx.flags & SCX_TG_SUB_INIT))
1386	continue;
1387
1388	/* skip the exit if the parent's ops.cgroup_init() failed */
1389	if ((tg->scx.flags & SCX_TG_INITED) && SCX_HAS_OP(parent, cgroup_exit))
1390	SCX_CALL_OP(parent, cgroup_exit, NULL, css->cgroup);
1391	tg->scx.sched = sch;
1392	tg->scx.flags \|= SCX_TG_INITED;
1393	tg->scx.flags &= ~SCX_TG_SUB_INIT;
1394	}
1395
1396	return 0;
1397
1398	err:
1399	css_for_each_descendant_post(css, ecss) {
1400	struct task_group *tg = css_tg(css);
1401
1402	if (!(tg->scx.flags & SCX_TG_SUB_INIT))
1403	continue;
1404
1405	if (SCX_HAS_OP(sch, cgroup_exit))
1406	SCX_CALL_OP(sch, cgroup_exit, NULL, css->cgroup);
1407	tg->scx.flags &= ~SCX_TG_SUB_INIT;
1408	}
1409	return ret;
1410	}
1411
1412	/**
1413	* scx_cgroup_return_subtree - Return the subtree's cgroups to the parent sched
1414	* @sch: sub-scheduler being disabled
1415	*
1416	* Called while disabling @sch, after the subtree's cgrp->scx_sched's are reset
1417	* to the parent sched and before tasks are re-homed, mirroring root disable's
1418	* cgroups-before-tasks teardown order. The reverse of
1419	* scx_cgroup_claim_subtree(): exit @sch's cgroups from @sch, then init them on
1420	* the parent with the current tg->scx.* values, resyncing settings that changed
1421	* while @sch had them.
1422	*
1423	* When an init on the parent fails, the parent is failed - the same policy as
1424	* task re-homing. The remaining task_groups are punted: they move to the parent
1425	* anyway with %SCX_TG_INITED cleared, as ops.cgroup_init() failed or never ran
1426	* for them. A punted task_group gets no cgroup ops. The dying parent's own
1427	* disable moves it one sched up, initing it there. Root ends the chain: root
1428	* teardown drops cgroup ops entirely and the next enable's bulk init re-inits
1429	* every online task_group.
1430	*
1431	* The task re-home that follows still delivers ops.init_task() to the dying
1432	* parent, including for tasks in punted cgroups it never inited - tolerated
1433	* like the downstream failures of task punting (see scx_punt_task()).
1434	*/
1435	static void scx_cgroup_return_subtree(struct scx_sched *sch)
1436	{
1437	struct cgroup *sub_cgrp = sch_cgroup(sch);
1438	struct cgroup_subsys_state *ecss = cgroup_e_css(sub_cgrp, &cpu_cgrp_subsys);
1439	struct scx_sched *parent = scx_parent(sch);
1440	struct cgroup_subsys_state *css;
1441	bool parent_failed = false;
1442	int ret;
1443
1444	css_for_each_descendant_post(css, ecss) {
1445	struct task_group *tg = css_tg(css);
1446
1447	if (tg->scx.sched != sch \|\|
1448	!cgroup_is_descendant(css->cgroup, sub_cgrp))
1449	continue;
1450
1451	/* skip the exit if @sch's ops.cgroup_init() failed for the tg */
1452	if ((tg->scx.flags & SCX_TG_INITED) && SCX_HAS_OP(sch, cgroup_exit))
1453	SCX_CALL_OP(sch, cgroup_exit, NULL, css->cgroup);
1454	tg->scx.sched = parent;
1455	tg->scx.flags \|= SCX_TG_SUB_INIT;
1456	}
1457
1458	css_for_each_descendant_pre(css, ecss) {
1459	struct task_group *tg = css_tg(css);
1460	struct scx_cgroup_init_args args = {
1461	.weight = tg->scx.weight,
1462	.bw_period_us = tg->scx.bw_period_us,
1463	.bw_quota_us = tg->scx.bw_quota_us,
1464	.bw_burst_us = tg->scx.bw_burst_us,
1465	.sched_idle = tg->scx.idle,
1466	};
1467
1468	/* the first pass must have transferred everything */
1469	WARN_ON_ONCE(tg->scx.sched == sch);
1470
1471	/*
1472	* SUB_INIT distinguishes the tgs pass 1 moved. The sched test
1473	* can't: a tg punted to the parent by an earlier failure would
1474	* also match.
1475	*/
1476	if (!(tg->scx.flags & SCX_TG_SUB_INIT))
1477	continue;
1478	tg->scx.flags &= ~(SCX_TG_SUB_INIT \| SCX_TG_INITED);
1479
1480	/*
1481	* A re-init on $parent failed. The task_groups from here on are
1482	* punted: they stay on the dying $parent with INITED clear and
1483	* move onward when it disables.
1484	*/
1485	if (parent_failed)
1486	continue;
1487
1488	if (SCX_HAS_OP(parent, cgroup_init)) {
1489	ret = SCX_CALL_OP_RET(parent, cgroup_init, NULL, css->cgroup, &args);
1490	if (ret) {
1491	scx_error(parent, "ops.cgroup_init() failed (%d) while disabling a sub-scheduler",
1492	ret);
1493	parent_failed = true;
1494	continue;
1495	}
1496	}
1497	tg->scx.flags \|= SCX_TG_INITED;
1498	}
1499	}
1500	#else
1501	static inline s32 scx_cgroup_claim_subtree(struct scx_sched *sch) { return 0; }
1502	static inline void scx_cgroup_return_subtree(struct scx_sched *sch) {}
1503	#endif
1504
1505	void scx_sub_disable(struct scx_sched *sch)
1506	{
1507	struct scx_sched *parent = scx_parent(sch);
1508	struct scx_task_iter sti;
1509	struct task_struct *p;
1510	int ret;
1511
1512	/*
1513	* Guarantee forward progress and wait for descendants to be disabled.
1514	* To limit disruptions, $parent is not bypassed. Tasks are fully
1515	* prepped and then inserted back into $parent.
1516	*/
1517	scx_bypass(sch, true);
1518	drain_descendants(sch);
1519
1520	/*
1521	* Here, every runnable task is guaranteed to make forward progress and
1522	* we can safely use blocking synchronization constructs. Actually
1523	* disable ops.
1524	*/
1525	mutex_lock(&scx_enable_mutex);
1526	percpu_down_write(&scx_fork_rwsem);
1527	scx_cgroup_lock();
1528
1529	/*
1530	* An enable that failed before scx_link_sched() succeeded never owned a
1531	* cgroup or task and won't be waited on by an ancestor's
1532	* drain_descendants(). Nothing to reparent and walking the tasks can
1533	* misbehave as the task ownership invariant (either owned by self or
1534	* parent) does not hold. ->sibling can't identify this case - an undone
1535	* link leaves it non-empty.
1536	*/
1537	if (!sch->linked)
1538	goto dump;
1539
1540	set_cgroup_sched(sch_cgroup(sch), parent);
1541
1542	/*
1543	* Return the subtree's cgroups before re-homing tasks so that any
1544	* ops.init_task() on $parent only sees cgroups it has initialized.
1545	*/
1546	scx_cgroup_return_subtree(sch);
1547
1548	scx_task_iter_start(&sti, sch->cgrp);
1549	while ((p = scx_task_iter_next_locked(&sti))) {
1550	struct rq *rq;
1551	struct rq_flags rf;
1552
1553	/* filter out duplicate visits */
1554	if (scx_task_on_sched(parent, p))
1555	continue;
1556
1557	/*
1558	* By the time control reaches here, all linked descendant
1559	* schedulers should have been disabled.
1560	*/
1561	WARN_ON_ONCE(!scx_task_on_sched(sch, p));
1562
1563	/*
1564	* @p is pinned by the iter: css_task_iter_next() takes a
1565	* reference and holds it until the next iter_next() call, so
1566	* @p->usage is guaranteed > 0.
1567	*/
1568	get_task_struct(p);
1569
1570	scx_task_iter_unlock(&sti);
1571
1572	/*
1573	* $p is READY or ENABLED on @sch. Initialize for $parent,
1574	* disable and exit from @sch, and then switch over to $parent.
1575	*
1576	* If a task fails to initialize for $parent, the only available
1577	* action is disabling $parent too. While this allows disabling
1578	* of a child sched to cause the parent scheduler to fail, the
1579	* failure can only originate from ops.init_task() of the
1580	* parent. A child can't directly affect the parent through its
1581	* own failures.
1582	*/
1583	ret = __scx_init_task(parent, p, NULL, false);
1584	if (ret) {
1585	scx_fail_parent(sch, p, ret);
1586	put_task_struct(p);
1587	break;
1588	}
1589
1590	rq = task_rq_lock(p, &rf);
1591
1592	if (scx_get_task_state(p) == SCX_TASK_DEAD) {
1593	/*
1594	* sched_ext_dead() raced us between __scx_init_task()
1595	* and this rq lock and ran exit_task() on @sch (the
1596	* sched @p was on at that point), not on $parent.
1597	* $parent's just-completed init is owed an exit_task()
1598	* and we issue it here.
1599	*/
1600	scx_sub_init_cancel_task(parent, p);
1601	task_rq_unlock(rq, p, &rf);
1602	put_task_struct(p);
1603	continue;
1604	}
1605
1606	scx_rehome_task(parent, p);
1607
1608	task_rq_unlock(rq, p, &rf);
1609	put_task_struct(p);
1610	}
1611	scx_task_iter_stop(&sti);
1612
1613	dump:
1614	scx_disable_dump(sch);
1615
1616	scx_cgroup_unlock();
1617	percpu_up_write(&scx_fork_rwsem);
1618
1619	/*
1620	* All tasks are moved off of @sch but there may still be on-going
1621	* operations (e.g. ops.select_cpu()). Drain them by flushing RCU. Use
1622	* the expedited version as ancestors may be waiting in bypass mode.
1623	* Also, tell the parent that there is no need to keep running bypass
1624	* DSQs for us.
1625	*/
1626	synchronize_rcu_expedited();
1627	scx_disable_bypass_dsp(sch);
1628
1629	scx_unlink_sched(sch);
1630
1631	mutex_unlock(&scx_enable_mutex);
1632
1633	/*
1634	* @sch is now unlinked from the parent's children list. Notify and call
1635	* ops.sub_detach/exit(). Note that ops.sub_detach/exit() must be called
1636	* after unlinking and releasing all locks. See scx_claim_exit().
1637	*/
1638	wake_up_all(&scx_unlink_waitq);
1639
1640	if (parent->ops.sub_detach && sch->sub_attached) {
1641	struct scx_sub_detach_args sub_detach_args = {
1642	.ops = &sch->ops,
1643	.cgroup_path = sch->cgrp_path,
1644	};
1645	SCX_CALL_OP(parent, sub_detach, NULL,
1646	&sub_detach_args);
1647	}
1648
1649	scx_log_sched_disable(sch);
1650
1651	if (sch->ops.exit)
1652	SCX_CALL_OP(sch, exit, NULL, sch->exit_info);
1653
1654	/*
1655	* @sch's non-ops programs such as timers and tracers can fire after
1656	* ops.exit(). Now that exit is complete, stop scx_prog_sched() from
1657	* resolving to @sch and drain in-flight resolvers.
1658	*/
1659	WRITE_ONCE(sch->dead, true);
1660	synchronize_rcu();
1661
1662	if (sch->sub_kset)
1663	kobject_del(&sch->sub_kset->kobj);
1664	/* not added if enable failed before scx_sched_sysfs_add() */
1665	if (sch->kobj.state_in_sysfs)
1666	kobject_del(&sch->kobj);
1667	}
1668
1669	/* verify that a scheduler can be attached to @cgrp and return the parent */
1670	static struct scx_sched find_parent_sched(struct cgroup cgrp)
1671	{
1672	struct scx_sched *parent = scx_cgroup_sched(cgrp);
1673	struct scx_sched *pos;
1674
1675	lockdep_assert_held(&scx_sched_lock);
1676
1677	/* can't attach twice to the same cgroup */
1678	if (parent->cgrp == cgrp)
1679	return ERR_PTR(-EBUSY);
1680
1681	/* does $parent allow sub-scheds? */
1682	if (!parent->ops.sub_attach)
1683	return ERR_PTR(-EOPNOTSUPP);
1684
1685	/* can't insert between $parent and its exiting children */
1686	list_for_each_entry(pos, &parent->children, sibling)
1687	if (cgroup_is_descendant(pos->cgrp, cgrp))
1688	return ERR_PTR(-EBUSY);
1689
1690	return parent;
1691	}
1692
1693	static bool assert_task_ready_or_enabled(struct task_struct *p)
1694	{
1695	u32 state = scx_get_task_state(p);
1696
1697	switch (state) {
1698	case SCX_TASK_READY:
1699	case SCX_TASK_ENABLED:
1700	return true;
1701	default:
1702	WARN_ONCE(true, "sched_ext: Invalid task state %d for %s[%d] during enabling sub sched",
1703	state, p->comm, p->pid);
1704	return false;
1705	}
1706	}
1707
1708	void scx_sub_enable_workfn(struct kthread_work *work)
1709	{
1710	struct scx_enable_cmd *cmd = container_of(work, struct scx_enable_cmd, work);
1711	struct sched_ext_ops *ops = cmd->ops;
1712	struct cgroup *cgrp;
1713	struct scx_sched parent, sch;
1714	struct scx_task_iter sti;
1715	struct task_struct *p;
1716	s32 i, ret;
1717
1718	mutex_lock(&scx_enable_mutex);
1719
1720	if (!scx_enabled()) {
1721	ret = -ENODEV;
1722	goto out_unlock;
1723	}
1724
1725	/* See scx_root_enable_workfn() for the @ops->priv check. */
1726	if (rcu_access_pointer(ops->priv)) {
1727	ret = -EBUSY;
1728	goto out_unlock;
1729	}
1730
1731	cgrp = cgroup_get_from_id(ops->sub_cgroup_id);
1732	if (IS_ERR(cgrp)) {
1733	ret = PTR_ERR(cgrp);
1734	goto out_unlock;
1735	}
1736
1737	raw_spin_lock_irq(&scx_sched_lock);
1738	parent = find_parent_sched(cgrp);
1739	if (IS_ERR(parent)) {
1740	raw_spin_unlock_irq(&scx_sched_lock);
1741	ret = PTR_ERR(parent);
1742	goto out_put_cgrp;
1743	}
1744	kobject_get(&parent->kobj);
1745	raw_spin_unlock_irq(&scx_sched_lock);
1746
1747	/*
1748	* Flip the hot-path gates before ops->priv is published - the sub's
1749	* programs can e.g. kick cpus from that point on. The matching dec is
1750	* at the end of scx_sched_free_rcu_work().
1751	*/
1752	static_branch_inc(&__scx_has_subs);
1753
1754	/* scx_alloc_and_add_sched() consumes @cgrp whether it succeeds or not */
1755	sch = scx_alloc_and_add_sched(cmd, cgrp, parent);
1756	kobject_put(&parent->kobj);
1757	if (IS_ERR(sch)) {
1758	static_branch_dec(&__scx_has_subs);
1759	ret = PTR_ERR(sch);
1760	goto out_unlock;
1761	}
1762
1763	/*
1764	* Validate before scx_link_sched() publishes @sch, so an invalid sub
1765	* never becomes visible with an unallocated pshard.
1766	*/
1767	ret = scx_validate_ops(sch, ops);
1768	if (ret)
1769	goto err_disable;
1770
1771	scx_rescue_check_timeout(sch);
1772
1773	/*
1774	* Allocate pshard[] before scx_link_sched() publishes @sch into the
1775	* parent's RCU children list. A concurrent revoke walking the tree
1776	* would otherwise dereference sch->pshard[si] while it's still NULL.
1777	* Unlike the root path, the cid shard layout is stable at this point.
1778	*
1779	* scx_alloc_pshards() skips allocation when @sch's arena pool isn't
1780	* initialized, so scx_arena_pool_init() must run first.
1781	*/
1782	ret = scx_arena_pool_init(sch);
1783	if (ret)
1784	goto err_disable;
1785
1786	ret = scx_alloc_pshards(sch);
1787	if (ret)
1788	goto err_disable;
1789
1790	ret = scx_link_sched(sch);
1791	if (ret)
1792	goto err_disable;
1793
1794	ret = scx_sched_sysfs_add(sch);
1795	if (ret)
1796	goto err_disable;
1797
1798	if (sch->level >= SCX_SUB_MAX_DEPTH) {
1799	scx_error(sch, "max nesting depth %d violated",
1800	SCX_SUB_MAX_DEPTH);
1801	ret = -EINVAL;
1802	goto err_disable;
1803	}
1804
1805	scoped_guard(cpus_read_lock) {
1806	ret = scx_alloc_kern_arena_objs(sch);
1807	if (ret)
1808	goto err_disable;
1809	}
1810
1811	if (sch->ops.init) {
1812	ret = SCX_CALL_OP_RET(sch, init, NULL);
1813	if (ret) {
1814	ret = scx_ops_sanitize_err(sch, "init", ret);
1815	scx_error(sch, "ops.init() failed (%d)", ret);
1816	goto err_disable;
1817	}
1818	sch->exit_info->flags \|= SCX_EFLAG_INITIALIZED;
1819	}
1820
1821	struct scx_sub_attach_args sub_attach_args = {
1822	.ops = &sch->ops,
1823	.cgroup_path = sch->cgrp_path,
1824	};
1825
1826	ret = SCX_CALL_OP_RET(parent, sub_attach, NULL,
1827	&sub_attach_args);
1828	if (ret) {
1829	ret = scx_ops_sanitize_err(sch, "sub_attach", ret);
1830	scx_error(sch, "parent rejected (%d)", ret);
1831	goto err_disable;
1832	}
1833	sch->sub_attached = true;
1834
1835	scx_bypass(sch, true);
1836
1837	for (i = SCX_OPI_BEGIN; i < SCX_OPI_END; i++)
1838	if (((void (**)(void))ops)[i])
1839	set_bit(i, sch->has_op);
1840
1841	percpu_down_write(&scx_fork_rwsem);
1842	scx_cgroup_lock();
1843
1844	/*
1845	* Set cgroup->scx_sched's and check CSS_ONLINE. Either we see
1846	* !CSS_ONLINE or scx_cgroup_lifetime_notify() sees and shoots us down.
1847	*/
1848	set_cgroup_sched(sch_cgroup(sch), sch);
1849	if (!(cgrp->self.flags & CSS_ONLINE)) {
1850	scx_error(sch, "cgroup is not online");
1851	ret = -ENODEV;
1852	goto err_unlock_and_disable;
1853	}
1854
1855	/*
1856	* Take over the subtree's cgroups before any task is claimed,
1857	* mirroring root enable's cgroups-before-tasks order.
1858	*/
1859	ret = scx_cgroup_claim_subtree(sch);
1860	if (ret)
1861	goto err_unlock_and_disable;
1862
1863	/*
1864	* Initialize tasks for the new child $sch without exiting them for
1865	* $parent so that the tasks can always be reverted back to $parent
1866	* sched on child init failure.
1867	*/
1868	WARN_ON_ONCE(scx_enabling_sub_sched);
1869	scx_enabling_sub_sched = sch;
1870
1871	scx_task_iter_start(&sti, sch->cgrp);
1872	while ((p = scx_task_iter_next_locked(&sti))) {
1873	struct rq *rq;
1874	struct rq_flags rf;
1875
1876	/*
1877	* Task iteration may visit the same task twice when racing
1878	* against exiting. Use %SCX_TASK_SUB_INIT to mark tasks which
1879	* finished __scx_init_task() and skip if set.
1880	*
1881	* A task may exit and get freed between __scx_init_task()
1882	* completion and scx_enable_task(). In such cases,
1883	* scx_disable_and_exit_task() must exit the task for both the
1884	* parent and child scheds.
1885	*/
1886	if (p->scx.flags & SCX_TASK_SUB_INIT)
1887	continue;
1888
1889	/* @p is pinned by the iter; see scx_sub_disable() */
1890	get_task_struct(p);
1891
1892	if (!assert_task_ready_or_enabled(p)) {
1893	ret = -EINVAL;
1894	goto abort;
1895	}
1896
1897	scx_task_iter_unlock(&sti);
1898
1899	/*
1900	* As $p is still on $parent, it can't be transitioned to INIT.
1901	* Let's worry about task state later. Use __scx_init_task().
1902	*/
1903	ret = __scx_init_task(sch, p, NULL, false);
1904	if (ret)
1905	goto abort;
1906
1907	rq = task_rq_lock(p, &rf);
1908
1909	if (scx_get_task_state(p) == SCX_TASK_DEAD) {
1910	/*
1911	* sched_ext_dead() raced us between __scx_init_task()
1912	* and this rq lock and ran exit_task() on $parent (the
1913	* sched @p was on at that point), not on @sch. @sch's
1914	* just-completed init is owed an exit_task() and we
1915	* issue it here.
1916	*/
1917	scx_sub_init_cancel_task(sch, p);
1918	task_rq_unlock(rq, p, &rf);
1919	put_task_struct(p);
1920	continue;
1921	}
1922
1923	p->scx.flags \|= SCX_TASK_SUB_INIT;
1924	task_rq_unlock(rq, p, &rf);
1925
1926	put_task_struct(p);
1927	}
1928	scx_task_iter_stop(&sti);
1929
1930	/*
1931	* All tasks are prepped. Disable/exit tasks for $parent and enable for
1932	* the new @sch.
1933	*/
1934	scx_task_iter_start(&sti, sch->cgrp);
1935	while ((p = scx_task_iter_next_locked(&sti))) {
1936	/*
1937	* Use clearing of %SCX_TASK_SUB_INIT to detect and skip
1938	* duplicate iterations.
1939	*/
1940	if (!(p->scx.flags & SCX_TASK_SUB_INIT))
1941	continue;
1942
1943	scoped_guard (sched_change, p, DEQUEUE_SAVE \| DEQUEUE_MOVE) {
1944	/*
1945	* $p must be either READY or ENABLED. If ENABLED,
1946	* __scx_disabled_and_exit_task() first disables and
1947	* makes it READY. However, after exiting $p, it will
1948	* leave $p as READY.
1949	*/
1950	assert_task_ready_or_enabled(p);
1951	__scx_disable_and_exit_task(parent, p);
1952
1953	/*
1954	* $p is now only initialized for @sch and READY, which
1955	* is what we want. Assign it to @sch and, if it's on
1956	* the ext class, enable. A non-ext task, possible under
1957	* an %SCX_OPS_SWITCH_PARTIAL root, stays READY and is
1958	* enabled by switching_to_scx() if it switches over.
1959	*/
1960	scx_set_task_sched(p, sch);
1961	if (p->sched_class == &ext_sched_class)
1962	scx_enable_task(sch, p);
1963
1964	p->scx.flags &= ~SCX_TASK_SUB_INIT;
1965	}
1966	}
1967	scx_task_iter_stop(&sti);
1968
1969	scx_enabling_sub_sched = NULL;
1970
1971	scx_cgroup_unlock();
1972	percpu_up_write(&scx_fork_rwsem);
1973
1974	scx_bypass(sch, false);
1975
1976	/* @sch is enabled; deliver any caps owed since its sub_attach() */
1977	scx_sub_seed_caps(sch);
1978
1979	pr_info("sched_ext: BPF sub-scheduler \"%s\" enabled\n", sch->ops.name);
1980	kobject_uevent(&sch->kobj, KOBJ_ADD);
1981	ret = 0;
1982	goto out_unlock;
1983
1984	out_put_cgrp:
1985	cgroup_put(cgrp);
1986	out_unlock:
1987	mutex_unlock(&scx_enable_mutex);
1988	cmd->ret = ret;
1989	return;
1990
1991	abort:
1992	put_task_struct(p);
1993	scx_task_iter_stop(&sti);
1994
1995	/*
1996	* Undo __scx_init_task() for tasks we marked. scx_enable_task() never
1997	* ran for @sch on them, so calling scx_disable_task() here would invoke
1998	* ops.disable() without a matching ops.enable(). scx_enabling_sub_sched
1999	* must stay set until SUB_INIT is cleared from every marked task -
2000	* scx_disable_and_exit_task() reads it when a task exits concurrently.
2001	*/
2002	scx_task_iter_start(&sti, sch->cgrp);
2003	while ((p = scx_task_iter_next_locked(&sti))) {
2004	if (p->scx.flags & SCX_TASK_SUB_INIT) {
2005	scx_sub_init_cancel_task(sch, p);
2006	p->scx.flags &= ~SCX_TASK_SUB_INIT;
2007	}
2008	}
2009	scx_task_iter_stop(&sti);
2010	scx_enabling_sub_sched = NULL;
2011	err_unlock_and_disable:
2012	/* we'll soon enter disable path, keep bypass on */
2013	scx_cgroup_unlock();
2014	percpu_up_write(&scx_fork_rwsem);
2015	err_disable:
2016	mutex_unlock(&scx_enable_mutex);
2017	/*
2018	* Some enable failures only return an errno (e.g. -ENOMEM from an
2019	* allocation) without calling scx_error(). Record it so
2020	* scx_flush_disable_work() runs the disable and ops.exit() fires.
2021	*/
2022	scx_error(sch, "scx_sub_enable() failed (%d)", ret);
2023	scx_flush_disable_work(sch);
2024	cmd->ret = 0;
2025	}
2026
2027	/**
2028	* scx_cgroup_task_migrating - Prepare a task for a cgroup migration
2029	* @ctx: migration being prepared
2030	*
2031	* A task's sched must match its cgroup's owner, so a migration that crosses a
2032	* sched boundary re-homes the task once committed. Run the fallible part here,
2033	* before the migration commits: initialize the task for the destination sched.
2034	* A rejection fails the cgroup.procs write.
2035	*/
2036	static s32 scx_cgroup_task_migrating(struct cgroup_task_migrate_ctx *ctx)
2037	{
2038	struct task_struct *p = ctx->task;
2039	struct scx_sched *to;
2040	int ret;
2041
2042	/*
2043	* Cleared under scx_cgroup_lock() before root disable starts tearing
2044	* down tasks. As cgroup_mutex is held, a set flag guarantees that the
2045	* teardown loop is not running concurrently.
2046	*/
2047	if (!scx_cgroup_enabled)
2048	return NOTIFY_OK;
2049
2050	to = scx_cgroup_sched(ctx->dst_dcgrp);
2051	if (scx_task_on_sched(to, p))
2052	return NOTIFY_OK;
2053
2054	ret = __scx_init_task(to, p, ctx->dst_dcgrp, false);
2055	if (ret)
2056	return notifier_from_errno(ret);
2057
2058	return NOTIFY_OK;
2059	}
2060
2061	/**
2062	* scx_cgroup_task_migrated - Re-home a task that changed cgroups
2063	* @ctx: committed migration
2064	*
2065	* Move the task to its new cgroup's sched, which scx_cgroup_task_migrating()
2066	* already initialized it for. Can't fail.
2067	*
2068	* This is safe against all phases of the destination sched's destruction. A
2069	* disable resets cgroup ownership to the parent and re-homes tasks in one
2070	* scx_cgroup_lock() section. If that section already ran, the destination would
2071	* be the parent. Otherwise, the re-home loop is still ahead and guaranteed to
2072	* visit the task, now in the destination cgroup.
2073	*/
2074	static void scx_cgroup_task_migrated(struct cgroup_task_migrate_ctx *ctx)
2075	{
2076	struct task_struct *p = ctx->task;
2077	struct scx_sched *to;
2078	struct rq *rq;
2079	struct rq_flags rf;
2080
2081	if (!scx_cgroup_enabled)
2082	return;
2083
2084	to = scx_cgroup_sched(ctx->dst_dcgrp);
2085	if (scx_task_on_sched(to, p))
2086	return;
2087
2088	rq = task_rq_lock(p, &rf);
2089	scx_rehome_task(to, p);
2090	task_rq_unlock(rq, p, &rf);
2091	}
2092
2093	/**
2094	* scx_cgroup_task_migrate_canceled - Undo migration preparation
2095	* @ctx: canceled migration
2096	*
2097	* The migration failed after scx_cgroup_task_migrating() initialized the task
2098	* for the destination sched. The task stays on its current sched in the source
2099	* cgroup. Undo the destination's init.
2100	*/
2101	static void scx_cgroup_task_migrate_canceled(struct cgroup_task_migrate_ctx *ctx)
2102	{
2103	struct task_struct *p = ctx->task;
2104	struct scx_sched *to;
2105	struct rq *rq;
2106	struct rq_flags rf;
2107
2108	if (!scx_cgroup_enabled)
2109	return;
2110
2111	to = scx_cgroup_sched(ctx->dst_dcgrp);
2112	if (scx_task_on_sched(to, p))
2113	return;
2114
2115	rq = task_rq_lock(p, &rf);
2116	scx_sub_init_cancel_task(to, p);
2117	task_rq_unlock(rq, p, &rf);
2118	}
2119
2120	static s32 scx_cgroup_lifetime_notify(struct notifier_block *nb,
2121	unsigned long action, void *data)
2122	{
2123	struct cgroup *cgrp = data;
2124	struct cgroup *parent = cgroup_parent(cgrp);
2125	struct scx_sched *sch;
2126
2127	if (!cgroup_on_dfl(cgrp))
2128	return NOTIFY_OK;
2129
2130	switch (action) {
2131	case CGROUP_LIFETIME_ONLINE:
2132	/* inherit ->scx_sched from $parent */
2133	if (parent)
2134	rcu_assign_pointer(cgrp->scx_sched, scx_cgroup_sched(parent));
2135	break;
2136	case CGROUP_LIFETIME_OFFLINE:
2137	/* if there is a sched attached, shoot it down */
2138	sch = scx_cgroup_sched(cgrp);
2139	if (sch && sch->cgrp == cgrp)
2140	scx_exit(sch, SCX_EXIT_UNREG_KERN,
2141	SCX_ECODE_RSN_CGROUP_OFFLINE,
2142	"cgroup %llu going offline", cgroup_id(cgrp));
2143	break;
2144	}
2145
2146	return NOTIFY_OK;
2147	}
2148
2149	static struct notifier_block scx_cgroup_lifetime_nb = {
2150	.notifier_call = scx_cgroup_lifetime_notify,
2151	};
2152
2153	static s32 scx_cgroup_task_notify(struct notifier_block *nb,
2154	unsigned long action, void *data)
2155	{
2156	struct cgroup_task_migrate_ctx *ctx = data;
2157
2158	switch (action) {
2159	case CGROUP_TASK_MIGRATING:
2160	return scx_cgroup_task_migrating(ctx);
2161	case CGROUP_TASK_MIGRATED:
2162	scx_cgroup_task_migrated(ctx);
2163	break;
2164	case CGROUP_TASK_MIGRATE_CANCELED:
2165	scx_cgroup_task_migrate_canceled(ctx);
2166	break;
2167	}
2168
2169	return NOTIFY_OK;
2170	}
2171
2172	static struct notifier_block scx_cgroup_task_nb = {
2173	.notifier_call = scx_cgroup_task_notify,
2174	};
2175
2176	static s32 __init scx_cgroup_notifier_init(void)
2177	{
2178	s32 ret;
2179
2180	ret = blocking_notifier_chain_register(&cgroup_lifetime_notifier,
2181	&scx_cgroup_lifetime_nb);
2182	if (ret)
2183	return ret;
2184
2185	return blocking_notifier_chain_register(&cgroup_task_notifier,
2186	&scx_cgroup_task_nb);
2187	}
2188	core_initcall(scx_cgroup_notifier_init);
2189
2190	static void scx_pstack_recursion(struct bpf_prog prog, const char op)
2191	{
2192	struct scx_sched *sch;
2193
2194	guard(rcu)();
2195	sch = scx_prog_sched(prog->aux);
2196	if (unlikely(!sch))
2197	return;
2198
2199	scx_error(sch, "%s recursion detected", op);
2200	}
2201
2202	void scx_pstack_recursion_on_dispatch(struct bpf_prog *prog)
2203	{
2204	scx_pstack_recursion(prog, "dispatch");
2205	}
2206
2207	void scx_pstack_recursion_on_caps_updated(struct bpf_prog *prog)
2208	{
2209	scx_pstack_recursion(prog, "sub_caps_updated");
2210	}
2211
2212	__bpf_kfunc_start_defs();
2213
2214	/**
2215	* scx_bpf_sub_dispatch - Trigger dispatching on a child scheduler
2216	* @cgroup_id: cgroup ID of the child scheduler to dispatch
2217	* @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
2218	*
2219	* Allows a parent scheduler to trigger dispatching on one of its direct
2220	* child schedulers. The child scheduler runs its dispatch operation to
2221	* move tasks from dispatch queues to the local runqueue.
2222	*
2223	* Returns: true on success, false if cgroup_id is invalid, not a direct
2224	* child, or caller lacks dispatch permission.
2225	*/
2226	__bpf_kfunc bool scx_bpf_sub_dispatch(u64 cgroup_id, const struct bpf_prog_aux *aux)
2227	{
2228	struct rq *rq = scx_locked_rq();
2229	struct scx_sched parent, child;
2230
2231	guard(rcu)();
2232	parent = scx_prog_sched(aux);
2233	if (unlikely(!parent))
2234	return false;
2235
2236	child = scx_find_sub_sched(cgroup_id);
2237
2238	if (unlikely(!child))
2239	return false;
2240
2241	if (unlikely(scx_parent(child) != parent)) {
2242	scx_error(parent, "trying to dispatch a distant sub-sched on cgroup %llu",
2243	cgroup_id);
2244	return false;
2245	}
2246
2247	/*
2248	* Skip a child that does not effectively hold the base cap on this cpu:
2249	* its inserts would only be rejected. ecaps are synced at the top of
2250	* dispatch_one() before dispatch, so this reflects the in-effect state.
2251	*/
2252	if (scx_missing_caps(child, cpu_of(rq), SCX_CAP_BASE))
2253	return false;
2254
2255	return scx_dispatch_sched(child, rq, rq->scx.sub_dispatch_prev, true) !=
2256	SCX_DSP_NONE;
2257	}
2258
2259	/* Validate common inputs. On success, parent_out and child_out are set. */
2260	static s32 sub_cap_preamble(u64 cgroup_id, u64 caps, const struct bpf_prog_aux *aux,
2261	struct scx_sched parent_out, struct scx_sched child_out)
2262	{
2263	struct scx_sched parent, child;
2264
2265	parent = scx_prog_sched(aux);
2266	if (unlikely(!parent))
2267	return -ENODEV;
2268
2269	if (!scx_is_cid_type()) {
2270	scx_error(parent, "sub-cap kfuncs require a cid-form scheduler");
2271	return -EOPNOTSUPP;
2272	}
2273
2274	child = scx_find_sub_sched(cgroup_id);
2275	if (unlikely(!child))
2276	return -ENODEV;
2277
2278	if (unlikely(scx_parent(child) != parent)) {
2279	scx_error(parent, "%s: sub-%llu is not a direct child",
2280	parent->cgrp_path, cgroup_id);
2281	return -EINVAL;
2282	}
2283
2284	if (unlikely(caps & ~__SCX_CAP_ALL)) {
2285	scx_error(parent, "invalid caps 0x%llx", caps);
2286	return -EINVAL;
2287	}
2288
2289	*parent_out = parent;
2290	*child_out = child;
2291	return 0;
2292	}
2293
2294	/**
2295	* scx_bpf_sub_grant - Grant @caps on a cmask's cids to a direct child
2296	* @cgroup_id: cgroup id of the direct child sub-sched
2297	* @caps: bitmask of SCX_CAP_* to grant
2298	* @cmask__arena: cid cmask to grant @caps on
2299	* @denied_out__arena__nullable: optional cmask accumulating refused cids
2300	* @aux: implicit BPF argument
2301	*
2302	* A cid in @cmask__arena is granted to the child only if the parent holds every
2303	* requested cap on it. Refused cids are OR'd into the denied mask when
2304	* provided. Refusals outside the denied mask's range are not recorded.
2305	*
2306	* All-or-nothing keeps the caller-visible result binary per cid, so the denied
2307	* mask is one mask to interpret rather than a per-cap matrix.
2308	*
2309	* Return 0 on full success, -EPERM if any cid was refused, or a negative
2310	* errno on other failures.
2311	*/
2312	__bpf_kfunc s32 scx_bpf_sub_grant(u64 cgroup_id, u64 caps,
2313	const struct scx_cmask *cmask__arena,
2314	struct scx_cmask *denied_out__arena__nullable,
2315	const struct bpf_prog_aux *aux)
2316	{
2317	struct scx_cmask_ref ref, denied_ref;
2318	struct scx_sched parent, child;
2319	bool any_denied = false;
2320	LIST_HEAD(to_deliver);
2321	s32 si, ret;
2322
2323	guard(irqsave)();
2324
2325	ret = sub_cap_preamble(cgroup_id, caps, aux, &parent, &child);
2326	if (ret)
2327	return ret;
2328
2329	ret = scx_cmask_ref_init(parent, cmask__arena, &ref);
2330	if (ret) {
2331	scx_error(parent, "invalid cmask (%d)", ret);
2332	return ret;
2333	}
2334
2335	if (denied_out__arena__nullable) {
2336	ret = scx_cmask_ref_init(parent, denied_out__arena__nullable, &denied_ref);
2337	if (ret) {
2338	scx_error(parent, "invalid denied_out (%d)", ret);
2339	return ret;
2340	}
2341	}
2342
2343	/* apply the grant one shard at a time */
2344	for (si = ref.shard_first; si < ref.shard_end; si++) {
2345	SCX_CMASK_DEFINE_SHARD(slice, 0, SCX_CID_SHARD_MAX_CPUS);
2346	struct scx_pshard *pps = parent->pshard[si];
2347	struct scx_pshard *cps = child->pshard[si];
2348	u64 granted_caps = 0;
2349	u32 cap_bit;
2350
2351	scx_cmask_ref_shard(&ref, si, slice);
2352	if (scx_cmask_empty(slice))
2353	continue;
2354
2355	SCX_CMASK_DEFINE_SHARD(granted_cids, slice->base, slice->nr_cids);
2356	SCX_CMASK_DEFINE_SHARD(changed_cids, slice->base, slice->nr_cids);
2357	SCX_CMASK_DEFINE_SHARD(delta, slice->base, slice->nr_cids);
2358
2359	scx_cmask_copy(granted_cids, slice);
2360
2361	scoped_guard (raw_spinlock, &pps->lock) {
2362	guard(raw_spinlock_nested)(&cps->lock);
2363
2364	/*
2365	* Narrow granted_cids to cids the parent holds every
2366	* requested cap on. All-or-nothing per cid.
2367	*/
2368	scx_for_each_cap_bit(cap_bit, caps)
2369	scx_cmask_and(granted_cids, &pps->caps[cap_bit].cmask);
2370
2371	/*
2372	* For each requested cap, fold the newly-set cids into
2373	* the child and accumulate the delta.
2374	*/
2375	scx_for_each_cap_bit(cap_bit, caps) {
2376	struct scx_cmask *ccm = &cps->caps[cap_bit].cmask;
2377
2378	scx_cmask_copy(delta, granted_cids);
2379	scx_cmask_andnot(delta, ccm);
2380	if (scx_cmask_empty(delta))
2381	continue;
2382
2383	scx_cmask_or(ccm, delta);
2384	scx_cmask_or(changed_cids, delta);
2385	granted_caps \|= BIT_U64(cap_bit);
2386	}
2387
2388	if (granted_caps) {
2389	s32 cid;
2390
2391	caps_updated_record(cps, changed_cids, granted_caps,
2392	&to_deliver);
2393	/*
2394	* The sync arms an update_idle() re-notify if
2395	* the cid gains baseline access, so the holder
2396	* learns of an already-idle cid.
2397	*/
2398	scx_cmask_for_each_cid(cid, changed_cids)
2399	queue_sync_ecaps(child, cid);
2400	}
2401	}
2402
2403	/* record cids that didn't make it into the denied mask */
2404	if (!scx_cmask_subset(slice, granted_cids)) {
2405	any_denied = true;
2406	if (denied_out__arena__nullable) {
2407	SCX_CMASK_DEFINE_SHARD(denied, slice->base, slice->nr_cids);
2408
2409	scx_cmask_copy(denied, slice);
2410	scx_cmask_andnot(denied, granted_cids);
2411	scx_cmask_ref_or(&denied_ref, denied);
2412	}
2413	}
2414	}
2415
2416	caps_updated_deliver(&to_deliver);
2417
2418	return any_denied ? -EPERM : 0;
2419	}
2420
2421	/**
2422	* scx_bpf_sub_revoke - Revoke @caps on a cmask's cids from a direct child
2423	* @cgroup_id: cgroup id of the direct child sub-sched
2424	* @caps: bitmask of SCX_CAP_* to revoke
2425	* @cmask__arena: cid cmask to revoke @caps on
2426	* @aux: implicit BPF argument
2427	*
2428	* Clear @caps bits on @cmask__arena from the child named by @cgroup_id and all
2429	* its descendants. The origin parent's pshard lock is held across the subtree
2430	* walk so a concurrent grant from the origin parent observes the revoked state.
2431	*/
2432	__bpf_kfunc void scx_bpf_sub_revoke(u64 cgroup_id, u64 caps,
2433	const struct scx_cmask *cmask__arena,
2434	const struct bpf_prog_aux *aux)
2435	{
2436	struct scx_cmask_ref ref;
2437	struct scx_sched parent, child, *pos;
2438	LIST_HEAD(to_deliver);
2439	s32 si, ret;
2440
2441	guard(irqsave)();
2442
2443	if (sub_cap_preamble(cgroup_id, caps, aux, &parent, &child))
2444	return;
2445
2446	ret = scx_cmask_ref_init(parent, cmask__arena, &ref);
2447	if (ret) {
2448	scx_error(parent, "invalid cmask (%d)", ret);
2449	return;
2450	}
2451
2452	/* per-shard, walk child's subtree and clear @caps */
2453	for (si = ref.shard_first; si < ref.shard_end; si++) {
2454	SCX_CMASK_DEFINE_SHARD(slice, 0, SCX_CID_SHARD_MAX_CPUS);
2455
2456	scx_cmask_ref_shard(&ref, si, slice);
2457	if (scx_cmask_empty(slice))
2458	continue;
2459
2460	/*
2461	* Pre-order with subtree skip: a descendant that cleared
2462	* nothing means no descendant of it can hold @caps on these
2463	* cids either.
2464	*/
2465	guard(raw_spinlock)(&parent->pshard[si]->lock);
2466	pos = scx_next_descendant_pre(NULL, child);
2467	while (pos) {
2468	struct scx_pshard *ps = pos->pshard[si];
2469	SCX_CMASK_DEFINE_SHARD(changed_cids, slice->base, slice->nr_cids);
2470	SCX_CMASK_DEFINE_SHARD(delta, slice->base, slice->nr_cids);
2471	u64 revoked_caps = 0;
2472	u32 cap_bit;
2473
2474	scoped_guard (raw_spinlock_nested, &ps->lock) {
2475	/*
2476	* For each cap, clear lost cids and accumulate
2477	* the per-cap diff for notification.
2478	*/
2479	scx_for_each_cap_bit(cap_bit, caps) {
2480	struct scx_cmask *cm = &ps->caps[cap_bit].cmask;
2481
2482	scx_cmask_copy(delta, cm);
2483	scx_cmask_and(delta, slice);
2484	if (scx_cmask_empty(delta))
2485	continue;
2486
2487	scx_cmask_andnot(cm, delta);
2488	scx_cmask_or(changed_cids, delta);
2489	revoked_caps \|= BIT_U64(cap_bit);
2490	}
2491
2492	if (revoked_caps) {
2493	s32 cid;
2494
2495	caps_updated_record(ps, changed_cids, revoked_caps,
2496	&to_deliver);
2497	scx_cmask_for_each_cid(cid, changed_cids)
2498	queue_sync_ecaps(pos, cid);
2499	}
2500	}
2501
2502	if (revoked_caps)
2503	pos = scx_next_descendant_pre(pos, child);
2504	else
2505	pos = scx_skip_subtree_pre(pos, child);
2506	}
2507	}
2508
2509	caps_updated_deliver(&to_deliver);
2510	}
2511
2512	/**
2513	* scx_bpf_sub_caps - Read self's or a direct child's cap cmasks
2514	* @cgroup_id: 0 for self, or a direct child's cgroup id
2515	* @caps: one or more SCX_CAP_* bits
2516	* @out__arena: cmask to receive the union of @caps within its range
2517	* @aux: implicit BPF argument
2518	*
2519	* Read the cap cmasks granted on each cid for self (@cgroup_id 0) or a direct
2520	* child - the literal granted set. A sched can read only itself or a direct
2521	* child.
2522	*
2523	* Return 0, -ENODEV if @cgroup_id names no direct child, or -EINVAL on bad
2524	* inputs.
2525	*/
2526	__bpf_kfunc s32 scx_bpf_sub_caps(u64 cgroup_id, u64 caps, struct scx_cmask *out__arena,
2527	const struct bpf_prog_aux *aux)
2528	{
2529	struct scx_cmask_ref ref;
2530	struct scx_sched sch, target;
2531	struct scx_pshard **pshard;
2532	s32 si, ret;
2533
2534	guard(irqsave)();
2535
2536	sch = scx_prog_sched(aux);
2537	if (unlikely(!sch))
2538	return -ENODEV;
2539
2540	if (!scx_is_cid_type()) {
2541	scx_error(sch, "sub-cap kfuncs require a cid-form scheduler");
2542	return -EOPNOTSUPP;
2543	}
2544
2545	if (unlikely(caps & ~__SCX_CAP_ALL)) {
2546	scx_error(sch, "invalid caps 0x%llx", caps);
2547	return -EINVAL;
2548	}
2549
2550	/* @cgroup_id 0 reads self, otherwise a direct child */
2551	if (cgroup_id) {
2552	target = scx_find_sub_sched(cgroup_id);
2553	if (unlikely(!target))
2554	return -ENODEV;
2555	if (unlikely(scx_parent(target) != sch)) {
2556	scx_error(sch, "%s: sub-%llu is not a direct child",
2557	sch->cgrp_path, cgroup_id);
2558	return -EINVAL;
2559	}
2560	} else {
2561	target = sch;
2562	}
2563
2564	/*
2565	* The target's caps storage may not be set up yet (e.g. a self-read
2566	* during ops.init_cids()). Pairs with the publish in
2567	* scx_alloc_pshards(): a non-NULL pshard has every element set and the
2568	* acquire also orders the cid table reads below against it.
2569	*/
2570	pshard = smp_load_acquire(&target->pshard);
2571	if (unlikely(!pshard)) {
2572	scx_error(sch, "scx_bpf_sub_caps() called before caps storage is initialized");
2573	return -ENODEV;
2574	}
2575
2576	ret = scx_cmask_ref_init(sch, out__arena, &ref);
2577	if (ret) {
2578	scx_error(sch, "invalid out (%d)", ret);
2579	return ret;
2580	}
2581
2582	for (si = ref.shard_first; si < ref.shard_end; si++) {
2583	const struct scx_cid_shard *shard =
2584	&rcu_dereference_all(scx_cid_shard_ranges)[si];
2585	SCX_CMASK_DEFINE_SHARD(local_out, shard->base_cid, shard->nr_cids);
2586	u32 cap_bit;
2587
2588	scx_for_each_cap_bit(cap_bit, caps)
2589	scx_cmask_or(local_out, &pshard[si]->caps[cap_bit].cmask);
2590	scx_cmask_ref_copy(&ref, local_out);
2591	}
2592	return 0;
2593	}
2594
2595	/**
2596	* scx_bpf_sub_kill_bstr - Kill a direct child sub-scheduler
2597	* @cgroup_id: cgroup id of the direct child to kill
2598	* @fmt: reason message format string
2599	* @data: format string parameters packaged using ___bpf_fill() macro
2600	* @data__sz: @data len, must end in '__sz' for the verifier
2601	* @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
2602	*
2603	* Evict a direct child sub-scheduler, disabling it with the supplied reason.
2604	* The child and its subtree are torn down asynchronously through the usual
2605	* disable path.
2606	*
2607	* Unlike scx_bpf_exit(), no exit code is taken: the child is a separate
2608	* scheduler with its own exit-code semantics, so a code chosen by the parent
2609	* would have no defined meaning. The reason string carries the intent.
2610	*
2611	* Return 0 on success or -ENODEV if @cgroup_id names no sub-scheduler, which
2612	* can race with the child detaching on its own and so is not a scheduler error.
2613	* Naming a sched that exists but is not a direct child aborts the parent.
2614	*/
2615	__printf(2, 0)
2616	__bpf_kfunc s32 scx_bpf_sub_kill_bstr(u64 cgroup_id, char *fmt,
2617	unsigned long long *data, u32 data__sz,
2618	const struct bpf_prog_aux *aux)
2619	{
2620	struct scx_sched parent, child;
2621
2622	guard(rcu)();
2623
2624	parent = scx_prog_sched(aux);
2625	if (unlikely(!parent))
2626	return -ENODEV;
2627
2628	if (!scx_is_cid_type()) {
2629	scx_error(parent, "sub-cap kfuncs require a cid-form scheduler");
2630	return -EOPNOTSUPP;
2631	}
2632
2633	child = scx_find_sub_sched(cgroup_id);
2634	if (unlikely(!child))
2635	return -ENODEV;
2636
2637	if (unlikely(scx_parent(child) != parent)) {
2638	scx_error(parent, "%s: sub-%llu is not a direct child",
2639	parent->cgrp_path, cgroup_id);
2640	return -EINVAL;
2641	}
2642
2643	scx_exit_bstr(child, SCX_EXIT_PARENT_KILL, 0, parent, fmt, data, data__sz);
2644	return 0;
2645	}
2646
2647	__bpf_kfunc_end_defs();
2648
2649	#else /* !CONFIG_EXT_SUB_SCHED */
2650
2651	__bpf_kfunc_start_defs();
2652
2653	__bpf_kfunc s32 scx_bpf_sub_grant(u64 cgroup_id, u64 caps,
2654	const struct scx_cmask *cmask__arena,
2655	struct scx_cmask *denied_out__arena__nullable,
2656	const struct bpf_prog_aux *aux)
2657	{
2658	return -EOPNOTSUPP;
2659	}
2660
2661	__bpf_kfunc void scx_bpf_sub_revoke(u64 cgroup_id, u64 caps,
2662	const struct scx_cmask *cmask__arena,
2663	const struct bpf_prog_aux *aux)
2664	{
2665	}
2666
2667	__bpf_kfunc s32 scx_bpf_sub_caps(u64 cgroup_id, u64 caps, struct scx_cmask *out__arena,
2668	const struct bpf_prog_aux *aux)
2669	{
2670	return -EOPNOTSUPP;
2671	}
2672
2673	__bpf_kfunc s32 scx_bpf_sub_kill_bstr(u64 cgroup_id, char *fmt,
2674	unsigned long long *data, u32 data__sz,
2675	const struct bpf_prog_aux *aux)
2676	{
2677	return -EOPNOTSUPP;
2678	}
2679
2680	__bpf_kfunc_end_defs();
2681
2682	#endif /* CONFIG_EXT_SUB_SCHED */