xref: /linux/kernel/sched/ext/sub.c (revision bbda59d85341d6d4d957596233646c84d9d9a451)
1daf8e166STejun Heo // SPDX-License-Identifier: GPL-2.0
2daf8e166STejun Heo /*
3daf8e166STejun Heo  * BPF extensible scheduler class: Documentation/scheduler/sched-ext.rst
4daf8e166STejun Heo  *
5daf8e166STejun Heo  * Sub-scheduler hierarchy support.
6daf8e166STejun Heo  *
7daf8e166STejun Heo  * A sub-scheduler is an scx_sched attached to a cgroup subtree under another
8daf8e166STejun Heo  * scx_sched. This file holds the sub-scheduler implementation: the scheduler
9daf8e166STejun Heo  * tree walk, capability delegation, per-shard cap state and its sync, and the
10daf8e166STejun Heo  * sub-scheduler enable/disable paths. The core dispatch/enqueue machinery it
11daf8e166STejun Heo  * builds on lives in ext.c.
12daf8e166STejun Heo  *
13daf8e166STejun Heo  * Copyright (c) 2026 Meta Platforms, Inc. and affiliates.
14daf8e166STejun Heo  * Copyright (c) 2026 Tejun Heo <tj@kernel.org>
15daf8e166STejun Heo  */
16daf8e166STejun Heo #include <linux/rhashtable.h>
17daf8e166STejun Heo #include "internal.h"
18daf8e166STejun Heo #include "cid.h"
19daf8e166STejun Heo #include "arena.h"
20daf8e166STejun Heo #include "sub.h"
21daf8e166STejun Heo 
22daf8e166STejun Heo #ifdef CONFIG_EXT_SUB_SCHED
23daf8e166STejun Heo 
24daf8e166STejun Heo /**
25*bbda59d8STejun Heo  * scx_skip_subtree_pre - Skip @pos's subtree in a pre-order walk
26*bbda59d8STejun Heo  * @pos: current position
27*bbda59d8STejun Heo  * @root: walk root
28*bbda59d8STejun Heo  *
29*bbda59d8STejun Heo  * In a walk started by scx_next_descendant_pre(), continue past @pos's subtree:
30*bbda59d8STejun Heo  * return @pos's next sibling, or the closest ancestor's next sibling, or NULL
31*bbda59d8STejun Heo  * if @pos's subtree is the last under @root. Same locking rules.
32*bbda59d8STejun Heo  */
33*bbda59d8STejun Heo struct scx_sched *scx_skip_subtree_pre(struct scx_sched *pos, struct scx_sched *root)
34*bbda59d8STejun Heo {
35*bbda59d8STejun Heo 	struct scx_sched *next;
36*bbda59d8STejun Heo 
37*bbda59d8STejun Heo 	lockdep_assert(lockdep_is_held(&scx_enable_mutex) ||
38*bbda59d8STejun Heo 		       lockdep_is_held(&scx_sched_lock) ||
39*bbda59d8STejun Heo 		       rcu_read_lock_any_held());
40*bbda59d8STejun Heo 
41*bbda59d8STejun Heo 	while (pos != root) {
42*bbda59d8STejun Heo 		next = list_next_or_null_rcu(&scx_parent(pos)->children, &pos->sibling,
43*bbda59d8STejun Heo 					     struct scx_sched, sibling);
44*bbda59d8STejun Heo 		if (next)
45*bbda59d8STejun Heo 			return next;
46*bbda59d8STejun Heo 		pos = scx_parent(pos);
47*bbda59d8STejun Heo 	}
48*bbda59d8STejun Heo 	return NULL;
49*bbda59d8STejun Heo }
50*bbda59d8STejun Heo 
51*bbda59d8STejun Heo /**
52daf8e166STejun Heo  * scx_next_descendant_pre - find the next descendant for pre-order walk
53daf8e166STejun Heo  * @pos: the current position (%NULL to initiate traversal)
54daf8e166STejun Heo  * @root: sched whose descendants to walk
55daf8e166STejun Heo  *
56daf8e166STejun Heo  * To be used by scx_for_each_descendant_pre(). Find the next descendant to
57daf8e166STejun Heo  * visit for pre-order traversal of @root's descendants. @root is included in
58daf8e166STejun Heo  * the iteration and the first node to be visited.
59daf8e166STejun Heo  */
60daf8e166STejun Heo struct scx_sched *scx_next_descendant_pre(struct scx_sched *pos, struct scx_sched *root)
61daf8e166STejun Heo {
62daf8e166STejun Heo 	struct scx_sched *next;
63daf8e166STejun Heo 
64daf8e166STejun Heo 	lockdep_assert(lockdep_is_held(&scx_enable_mutex) ||
6570f8b178STejun Heo 		       lockdep_is_held(&scx_sched_lock) ||
6670f8b178STejun Heo 		       rcu_read_lock_any_held());
67daf8e166STejun Heo 
68daf8e166STejun Heo 	/* if first iteration, visit @root */
69daf8e166STejun Heo 	if (!pos)
70daf8e166STejun Heo 		return root;
71daf8e166STejun Heo 
72daf8e166STejun Heo 	/* visit the first child if exists */
7370f8b178STejun Heo 	next = list_first_or_null_rcu(&pos->children, struct scx_sched, sibling);
74daf8e166STejun Heo 	if (next)
75daf8e166STejun Heo 		return next;
76daf8e166STejun Heo 
77daf8e166STejun Heo 	/* no child, visit my or the closest ancestor's next sibling */
78*bbda59d8STejun Heo 	return scx_skip_subtree_pre(pos, root);
79daf8e166STejun Heo }
80daf8e166STejun Heo 
81daf8e166STejun Heo static struct scx_sched *scx_find_sub_sched(u64 cgroup_id)
82daf8e166STejun Heo {
83daf8e166STejun Heo 	return rhashtable_lookup(&scx_sched_hash, &cgroup_id,
84daf8e166STejun Heo 				 scx_sched_hash_params);
85daf8e166STejun Heo }
86daf8e166STejun Heo 
87daf8e166STejun Heo void scx_set_task_sched(struct task_struct *p, struct scx_sched *sch)
88daf8e166STejun Heo {
89daf8e166STejun Heo 	rcu_assign_pointer(p->scx.sched, sch);
90daf8e166STejun Heo }
91daf8e166STejun Heo 
92daf8e166STejun Heo struct cgroup *sch_cgroup(struct scx_sched *sch)
93daf8e166STejun Heo {
94daf8e166STejun Heo 	return sch->cgrp;
95daf8e166STejun Heo }
96daf8e166STejun Heo 
97daf8e166STejun Heo /* for each descendant of @cgrp including self, set ->scx_sched to @sch */
98daf8e166STejun Heo void set_cgroup_sched(struct cgroup *cgrp, struct scx_sched *sch)
99daf8e166STejun Heo {
100daf8e166STejun Heo 	struct cgroup *pos;
101daf8e166STejun Heo 	struct cgroup_subsys_state *css;
102daf8e166STejun Heo 
103daf8e166STejun Heo 	cgroup_for_each_live_descendant_pre(pos, css, cgrp)
104daf8e166STejun Heo 		rcu_assign_pointer(pos->scx_sched, sch);
105daf8e166STejun Heo }
106daf8e166STejun Heo 
1078dba3bbdSTejun Heo static void free_pshard(struct scx_pshard *pshard)
1088dba3bbdSTejun Heo {
1098dba3bbdSTejun Heo 	kfree(pshard);
1108dba3bbdSTejun Heo }
1118dba3bbdSTejun Heo 
1128dba3bbdSTejun Heo void scx_free_pshards(struct scx_sched *sch)
1138dba3bbdSTejun Heo {
1148dba3bbdSTejun Heo 	s32 si;
1158dba3bbdSTejun Heo 
1168dba3bbdSTejun Heo 	if (!sch->pshard)
1178dba3bbdSTejun Heo 		return;
1188dba3bbdSTejun Heo 	for (si = 0; si < sch->nr_pshards; si++)
1198dba3bbdSTejun Heo 		free_pshard(sch->pshard[si]);
1208dba3bbdSTejun Heo 	kfree(sch->pshard);
1218dba3bbdSTejun Heo }
1228dba3bbdSTejun Heo 
1238dba3bbdSTejun Heo static struct scx_pshard *alloc_pshard(struct scx_sched *sch, s32 shard_idx, s32 node)
1248dba3bbdSTejun Heo {
1258dba3bbdSTejun Heo 	return kzalloc_node(sizeof(struct scx_pshard), GFP_KERNEL, node);
1268dba3bbdSTejun Heo }
1278dba3bbdSTejun Heo 
1288dba3bbdSTejun Heo s32 scx_alloc_pshards(struct scx_sched *sch)
1298dba3bbdSTejun Heo {
1308dba3bbdSTejun Heo 	struct scx_pshard **pshard;
1318dba3bbdSTejun Heo 	s32 si;
1328dba3bbdSTejun Heo 
1338dba3bbdSTejun Heo 	if (!sch->is_cid_type || !sch->arena_pool)
1348dba3bbdSTejun Heo 		return 0;
1358dba3bbdSTejun Heo 
1368dba3bbdSTejun Heo 	pshard = kzalloc_objs(pshard[0], scx_nr_cid_shards, GFP_KERNEL);
1378dba3bbdSTejun Heo 	if (!pshard)
1388dba3bbdSTejun Heo 		return -ENOMEM;
1398dba3bbdSTejun Heo 
1408dba3bbdSTejun Heo 	for (si = 0; si < scx_nr_cid_shards; si++) {
1418dba3bbdSTejun Heo 		pshard[si] = alloc_pshard(sch, si, scx_shard_node[si]);
1428dba3bbdSTejun Heo 		if (!pshard[si]) {
1438dba3bbdSTejun Heo 			while (--si >= 0)
1448dba3bbdSTejun Heo 				free_pshard(pshard[si]);
1458dba3bbdSTejun Heo 			kfree(pshard);
1468dba3bbdSTejun Heo 			return -ENOMEM;
1478dba3bbdSTejun Heo 		}
1488dba3bbdSTejun Heo 	}
1498dba3bbdSTejun Heo 
1508dba3bbdSTejun Heo 	sch->nr_pshards = scx_nr_cid_shards;
1518dba3bbdSTejun Heo 	/*
1528dba3bbdSTejun Heo 	 * Publish only after every entry is built so a reader observing
1538dba3bbdSTejun Heo 	 * @sch->pshard never sees a partially-filled array. Pair the store
1548dba3bbdSTejun Heo 	 * with a barrier and READ_ONCE() on the read side.
1558dba3bbdSTejun Heo 	 */
1568dba3bbdSTejun Heo 	smp_wmb();
1578dba3bbdSTejun Heo 	WRITE_ONCE(sch->pshard, pshard);
1588dba3bbdSTejun Heo 	return 0;
1598dba3bbdSTejun Heo }
1608dba3bbdSTejun Heo 
161daf8e166STejun Heo static DECLARE_WAIT_QUEUE_HEAD(scx_unlink_waitq);
162daf8e166STejun Heo 
163daf8e166STejun Heo void drain_descendants(struct scx_sched *sch)
164daf8e166STejun Heo {
165daf8e166STejun Heo 	/*
166daf8e166STejun Heo 	 * Child scheds that finished the critical part of disabling will take
167daf8e166STejun Heo 	 * themselves off @sch->children. Wait for it to drain. As propagation
168daf8e166STejun Heo 	 * is recursive, empty @sch->children means that all proper descendant
169daf8e166STejun Heo 	 * scheds reached unlinking stage.
170daf8e166STejun Heo 	 */
171daf8e166STejun Heo 	wait_event(scx_unlink_waitq, list_empty(&sch->children));
172daf8e166STejun Heo }
173daf8e166STejun Heo 
174daf8e166STejun Heo static void scx_fail_parent(struct scx_sched *sch,
175daf8e166STejun Heo 			    struct task_struct *failed, s32 fail_code)
176daf8e166STejun Heo {
177daf8e166STejun Heo 	struct scx_sched *parent = scx_parent(sch);
178daf8e166STejun Heo 	struct scx_task_iter sti;
179daf8e166STejun Heo 	struct task_struct *p;
180daf8e166STejun Heo 
181daf8e166STejun Heo 	scx_error(parent, "ops.init_task() failed (%d) for %s[%d] while disabling a sub-scheduler",
182daf8e166STejun Heo 		  fail_code, failed->comm, failed->pid);
183daf8e166STejun Heo 
184daf8e166STejun Heo 	/*
185daf8e166STejun Heo 	 * Once $parent is bypassed, it's safe to put SCX_TASK_NONE tasks into
186daf8e166STejun Heo 	 * it. This may cause downstream failures on the BPF side but $parent is
187daf8e166STejun Heo 	 * dying anyway.
188daf8e166STejun Heo 	 */
189daf8e166STejun Heo 	scx_bypass(parent, true);
190daf8e166STejun Heo 
191daf8e166STejun Heo 	scx_task_iter_start(&sti, sch->cgrp);
192daf8e166STejun Heo 	while ((p = scx_task_iter_next_locked(&sti))) {
193daf8e166STejun Heo 		if (scx_task_on_sched(parent, p))
194daf8e166STejun Heo 			continue;
195daf8e166STejun Heo 
196daf8e166STejun Heo 		scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
197daf8e166STejun Heo 			scx_disable_and_exit_task(sch, p);
198daf8e166STejun Heo 			scx_set_task_sched(p, parent);
199daf8e166STejun Heo 		}
200daf8e166STejun Heo 	}
201daf8e166STejun Heo 	scx_task_iter_stop(&sti);
202daf8e166STejun Heo }
203daf8e166STejun Heo 
204daf8e166STejun Heo void scx_sub_disable(struct scx_sched *sch)
205daf8e166STejun Heo {
206daf8e166STejun Heo 	struct scx_sched *parent = scx_parent(sch);
207daf8e166STejun Heo 	struct scx_task_iter sti;
208daf8e166STejun Heo 	struct task_struct *p;
209daf8e166STejun Heo 	int ret;
210daf8e166STejun Heo 
211daf8e166STejun Heo 	/*
212daf8e166STejun Heo 	 * Guarantee forward progress and wait for descendants to be disabled.
213daf8e166STejun Heo 	 * To limit disruptions, $parent is not bypassed. Tasks are fully
214daf8e166STejun Heo 	 * prepped and then inserted back into $parent.
215daf8e166STejun Heo 	 */
216daf8e166STejun Heo 	scx_bypass(sch, true);
217daf8e166STejun Heo 	drain_descendants(sch);
218daf8e166STejun Heo 
219daf8e166STejun Heo 	/*
220daf8e166STejun Heo 	 * Here, every runnable task is guaranteed to make forward progress and
221daf8e166STejun Heo 	 * we can safely use blocking synchronization constructs. Actually
222daf8e166STejun Heo 	 * disable ops.
223daf8e166STejun Heo 	 */
224daf8e166STejun Heo 	mutex_lock(&scx_enable_mutex);
225daf8e166STejun Heo 	percpu_down_write(&scx_fork_rwsem);
226daf8e166STejun Heo 	scx_cgroup_lock();
227daf8e166STejun Heo 
228daf8e166STejun Heo 	set_cgroup_sched(sch_cgroup(sch), parent);
229daf8e166STejun Heo 
230daf8e166STejun Heo 	scx_task_iter_start(&sti, sch->cgrp);
231daf8e166STejun Heo 	while ((p = scx_task_iter_next_locked(&sti))) {
232daf8e166STejun Heo 		struct rq *rq;
233daf8e166STejun Heo 		struct rq_flags rf;
234daf8e166STejun Heo 
235daf8e166STejun Heo 		/* filter out duplicate visits */
236daf8e166STejun Heo 		if (scx_task_on_sched(parent, p))
237daf8e166STejun Heo 			continue;
238daf8e166STejun Heo 
239daf8e166STejun Heo 		/*
240daf8e166STejun Heo 		 * By the time control reaches here, all descendant schedulers
241daf8e166STejun Heo 		 * should already have been disabled.
242daf8e166STejun Heo 		 */
243daf8e166STejun Heo 		WARN_ON_ONCE(!scx_task_on_sched(sch, p));
244daf8e166STejun Heo 
245daf8e166STejun Heo 		/*
246daf8e166STejun Heo 		 * @p is pinned by the iter: css_task_iter_next() takes a
247daf8e166STejun Heo 		 * reference and holds it until the next iter_next() call, so
248daf8e166STejun Heo 		 * @p->usage is guaranteed > 0.
249daf8e166STejun Heo 		 */
250daf8e166STejun Heo 		get_task_struct(p);
251daf8e166STejun Heo 
252daf8e166STejun Heo 		scx_task_iter_unlock(&sti);
253daf8e166STejun Heo 
254daf8e166STejun Heo 		/*
255daf8e166STejun Heo 		 * $p is READY or ENABLED on @sch. Initialize for $parent,
256daf8e166STejun Heo 		 * disable and exit from @sch, and then switch over to $parent.
257daf8e166STejun Heo 		 *
258daf8e166STejun Heo 		 * If a task fails to initialize for $parent, the only available
259daf8e166STejun Heo 		 * action is disabling $parent too. While this allows disabling
260daf8e166STejun Heo 		 * of a child sched to cause the parent scheduler to fail, the
261daf8e166STejun Heo 		 * failure can only originate from ops.init_task() of the
262daf8e166STejun Heo 		 * parent. A child can't directly affect the parent through its
263daf8e166STejun Heo 		 * own failures.
264daf8e166STejun Heo 		 */
265daf8e166STejun Heo 		ret = __scx_init_task(parent, p, false);
266daf8e166STejun Heo 		if (ret) {
267daf8e166STejun Heo 			scx_fail_parent(sch, p, ret);
268daf8e166STejun Heo 			put_task_struct(p);
269daf8e166STejun Heo 			break;
270daf8e166STejun Heo 		}
271daf8e166STejun Heo 
272daf8e166STejun Heo 		rq = task_rq_lock(p, &rf);
273daf8e166STejun Heo 
274daf8e166STejun Heo 		if (scx_get_task_state(p) == SCX_TASK_DEAD) {
275daf8e166STejun Heo 			/*
276daf8e166STejun Heo 			 * sched_ext_dead() raced us between __scx_init_task()
277daf8e166STejun Heo 			 * and this rq lock and ran exit_task() on @sch (the
278daf8e166STejun Heo 			 * sched @p was on at that point), not on $parent.
279daf8e166STejun Heo 			 * $parent's just-completed init is owed an exit_task()
280daf8e166STejun Heo 			 * and we issue it here.
281daf8e166STejun Heo 			 */
282daf8e166STejun Heo 			scx_sub_init_cancel_task(parent, p);
283daf8e166STejun Heo 			task_rq_unlock(rq, p, &rf);
284daf8e166STejun Heo 			put_task_struct(p);
285daf8e166STejun Heo 			continue;
286daf8e166STejun Heo 		}
287daf8e166STejun Heo 
288daf8e166STejun Heo 		scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
289daf8e166STejun Heo 			/*
290daf8e166STejun Heo 			 * $p is initialized for $parent and still attached to
291daf8e166STejun Heo 			 * @sch. Disable and exit for @sch, switch over to
292daf8e166STejun Heo 			 * $parent, override the state to READY to account for
293daf8e166STejun Heo 			 * $p having already been initialized, and then enable.
294daf8e166STejun Heo 			 */
295daf8e166STejun Heo 			scx_disable_and_exit_task(sch, p);
296daf8e166STejun Heo 			scx_set_task_state(p, SCX_TASK_INIT_BEGIN);
297daf8e166STejun Heo 			scx_set_task_state(p, SCX_TASK_INIT);
298daf8e166STejun Heo 			scx_set_task_sched(p, parent);
299daf8e166STejun Heo 			scx_set_task_state(p, SCX_TASK_READY);
300daf8e166STejun Heo 			scx_enable_task(parent, p);
301daf8e166STejun Heo 		}
302daf8e166STejun Heo 
303daf8e166STejun Heo 		task_rq_unlock(rq, p, &rf);
304daf8e166STejun Heo 		put_task_struct(p);
305daf8e166STejun Heo 	}
306daf8e166STejun Heo 	scx_task_iter_stop(&sti);
307daf8e166STejun Heo 
308daf8e166STejun Heo 	scx_disable_dump(sch);
309daf8e166STejun Heo 
310daf8e166STejun Heo 	scx_cgroup_unlock();
311daf8e166STejun Heo 	percpu_up_write(&scx_fork_rwsem);
312daf8e166STejun Heo 
313daf8e166STejun Heo 	/*
314daf8e166STejun Heo 	 * All tasks are moved off of @sch but there may still be on-going
315daf8e166STejun Heo 	 * operations (e.g. ops.select_cpu()). Drain them by flushing RCU. Use
316daf8e166STejun Heo 	 * the expedited version as ancestors may be waiting in bypass mode.
317daf8e166STejun Heo 	 * Also, tell the parent that there is no need to keep running bypass
318daf8e166STejun Heo 	 * DSQs for us.
319daf8e166STejun Heo 	 */
320daf8e166STejun Heo 	synchronize_rcu_expedited();
321daf8e166STejun Heo 	scx_disable_bypass_dsp(sch);
322daf8e166STejun Heo 
323daf8e166STejun Heo 	scx_unlink_sched(sch);
324daf8e166STejun Heo 
325daf8e166STejun Heo 	mutex_unlock(&scx_enable_mutex);
326daf8e166STejun Heo 
327daf8e166STejun Heo 	/*
328daf8e166STejun Heo 	 * @sch is now unlinked from the parent's children list. Notify and call
329daf8e166STejun Heo 	 * ops.sub_detach/exit(). Note that ops.sub_detach/exit() must be called
330daf8e166STejun Heo 	 * after unlinking and releasing all locks. See scx_claim_exit().
331daf8e166STejun Heo 	 */
332daf8e166STejun Heo 	wake_up_all(&scx_unlink_waitq);
333daf8e166STejun Heo 
334daf8e166STejun Heo 	if (parent->ops.sub_detach && sch->sub_attached) {
335daf8e166STejun Heo 		struct scx_sub_detach_args sub_detach_args = {
336daf8e166STejun Heo 			.ops = &sch->ops,
337daf8e166STejun Heo 			.cgroup_path = sch->cgrp_path,
338daf8e166STejun Heo 		};
339daf8e166STejun Heo 		SCX_CALL_OP(parent, sub_detach, NULL,
340daf8e166STejun Heo 			    &sub_detach_args);
341daf8e166STejun Heo 	}
342daf8e166STejun Heo 
343daf8e166STejun Heo 	scx_log_sched_disable(sch);
344daf8e166STejun Heo 
345daf8e166STejun Heo 	if (sch->ops.exit)
346daf8e166STejun Heo 		SCX_CALL_OP(sch, exit, NULL, sch->exit_info);
347daf8e166STejun Heo 	if (sch->sub_kset)
348daf8e166STejun Heo 		kobject_del(&sch->sub_kset->kobj);
34980e6adaaSTejun Heo 	/* not added if enable failed before scx_sched_sysfs_add() */
35080e6adaaSTejun Heo 	if (sch->kobj.state_in_sysfs)
351daf8e166STejun Heo 		kobject_del(&sch->kobj);
352daf8e166STejun Heo }
353daf8e166STejun Heo 
354daf8e166STejun Heo /* verify that a scheduler can be attached to @cgrp and return the parent */
355daf8e166STejun Heo static struct scx_sched *find_parent_sched(struct cgroup *cgrp)
356daf8e166STejun Heo {
357daf8e166STejun Heo 	struct scx_sched *parent = cgrp->scx_sched;
358daf8e166STejun Heo 	struct scx_sched *pos;
359daf8e166STejun Heo 
360daf8e166STejun Heo 	lockdep_assert_held(&scx_sched_lock);
361daf8e166STejun Heo 
362daf8e166STejun Heo 	/* can't attach twice to the same cgroup */
363daf8e166STejun Heo 	if (parent->cgrp == cgrp)
364daf8e166STejun Heo 		return ERR_PTR(-EBUSY);
365daf8e166STejun Heo 
366daf8e166STejun Heo 	/* does $parent allow sub-scheds? */
367daf8e166STejun Heo 	if (!parent->ops.sub_attach)
368daf8e166STejun Heo 		return ERR_PTR(-EOPNOTSUPP);
369daf8e166STejun Heo 
370daf8e166STejun Heo 	/* can't insert between $parent and its exiting children */
371daf8e166STejun Heo 	list_for_each_entry(pos, &parent->children, sibling)
372daf8e166STejun Heo 		if (cgroup_is_descendant(pos->cgrp, cgrp))
373daf8e166STejun Heo 			return ERR_PTR(-EBUSY);
374daf8e166STejun Heo 
375daf8e166STejun Heo 	return parent;
376daf8e166STejun Heo }
377daf8e166STejun Heo 
378daf8e166STejun Heo static bool assert_task_ready_or_enabled(struct task_struct *p)
379daf8e166STejun Heo {
380daf8e166STejun Heo 	u32 state = scx_get_task_state(p);
381daf8e166STejun Heo 
382daf8e166STejun Heo 	switch (state) {
383daf8e166STejun Heo 	case SCX_TASK_READY:
384daf8e166STejun Heo 	case SCX_TASK_ENABLED:
385daf8e166STejun Heo 		return true;
386daf8e166STejun Heo 	default:
387daf8e166STejun Heo 		WARN_ONCE(true, "sched_ext: Invalid task state %d for %s[%d] during enabling sub sched",
388daf8e166STejun Heo 			  state, p->comm, p->pid);
389daf8e166STejun Heo 		return false;
390daf8e166STejun Heo 	}
391daf8e166STejun Heo }
392daf8e166STejun Heo 
393daf8e166STejun Heo void scx_sub_enable_workfn(struct kthread_work *work)
394daf8e166STejun Heo {
395daf8e166STejun Heo 	struct scx_enable_cmd *cmd = container_of(work, struct scx_enable_cmd, work);
396daf8e166STejun Heo 	struct sched_ext_ops *ops = cmd->ops;
397daf8e166STejun Heo 	struct cgroup *cgrp;
398daf8e166STejun Heo 	struct scx_sched *parent, *sch;
399daf8e166STejun Heo 	struct scx_task_iter sti;
400daf8e166STejun Heo 	struct task_struct *p;
401daf8e166STejun Heo 	s32 i, ret;
402daf8e166STejun Heo 
403daf8e166STejun Heo 	mutex_lock(&scx_enable_mutex);
404daf8e166STejun Heo 
405daf8e166STejun Heo 	if (!scx_enabled()) {
406daf8e166STejun Heo 		ret = -ENODEV;
407daf8e166STejun Heo 		goto out_unlock;
408daf8e166STejun Heo 	}
409daf8e166STejun Heo 
410daf8e166STejun Heo 	/* See scx_root_enable_workfn() for the @ops->priv check. */
411daf8e166STejun Heo 	if (rcu_access_pointer(ops->priv)) {
412daf8e166STejun Heo 		ret = -EBUSY;
413daf8e166STejun Heo 		goto out_unlock;
414daf8e166STejun Heo 	}
415daf8e166STejun Heo 
416daf8e166STejun Heo 	cgrp = cgroup_get_from_id(ops->sub_cgroup_id);
417daf8e166STejun Heo 	if (IS_ERR(cgrp)) {
418daf8e166STejun Heo 		ret = PTR_ERR(cgrp);
419daf8e166STejun Heo 		goto out_unlock;
420daf8e166STejun Heo 	}
421daf8e166STejun Heo 
422daf8e166STejun Heo 	raw_spin_lock_irq(&scx_sched_lock);
423daf8e166STejun Heo 	parent = find_parent_sched(cgrp);
424daf8e166STejun Heo 	if (IS_ERR(parent)) {
425daf8e166STejun Heo 		raw_spin_unlock_irq(&scx_sched_lock);
426daf8e166STejun Heo 		ret = PTR_ERR(parent);
427daf8e166STejun Heo 		goto out_put_cgrp;
428daf8e166STejun Heo 	}
429daf8e166STejun Heo 	kobject_get(&parent->kobj);
430daf8e166STejun Heo 	raw_spin_unlock_irq(&scx_sched_lock);
431daf8e166STejun Heo 
432daf8e166STejun Heo 	/* scx_alloc_and_add_sched() consumes @cgrp whether it succeeds or not */
433daf8e166STejun Heo 	sch = scx_alloc_and_add_sched(cmd, cgrp, parent);
434daf8e166STejun Heo 	kobject_put(&parent->kobj);
435daf8e166STejun Heo 	if (IS_ERR(sch)) {
436daf8e166STejun Heo 		ret = PTR_ERR(sch);
437daf8e166STejun Heo 		goto out_unlock;
438daf8e166STejun Heo 	}
439daf8e166STejun Heo 
440daf8e166STejun Heo 	ret = scx_link_sched(sch);
441daf8e166STejun Heo 	if (ret)
442daf8e166STejun Heo 		goto err_disable;
443daf8e166STejun Heo 
44480e6adaaSTejun Heo 	ret = scx_sched_sysfs_add(sch);
44580e6adaaSTejun Heo 	if (ret)
44680e6adaaSTejun Heo 		goto err_disable;
44780e6adaaSTejun Heo 
448daf8e166STejun Heo 	if (sch->level >= SCX_SUB_MAX_DEPTH) {
449daf8e166STejun Heo 		scx_error(sch, "max nesting depth %d violated",
450daf8e166STejun Heo 			  SCX_SUB_MAX_DEPTH);
451daf8e166STejun Heo 		goto err_disable;
452daf8e166STejun Heo 	}
453daf8e166STejun Heo 
454daf8e166STejun Heo 	if (sch->ops.init) {
455daf8e166STejun Heo 		ret = SCX_CALL_OP_RET(sch, init, NULL);
456daf8e166STejun Heo 		if (ret) {
457daf8e166STejun Heo 			ret = scx_ops_sanitize_err(sch, "init", ret);
458daf8e166STejun Heo 			scx_error(sch, "ops.init() failed (%d)", ret);
459daf8e166STejun Heo 			goto err_disable;
460daf8e166STejun Heo 		}
461daf8e166STejun Heo 		sch->exit_info->flags |= SCX_EFLAG_INITIALIZED;
462daf8e166STejun Heo 	}
463daf8e166STejun Heo 
464daf8e166STejun Heo 	ret = scx_arena_pool_init(sch);
465daf8e166STejun Heo 	if (ret)
466daf8e166STejun Heo 		goto err_disable;
467daf8e166STejun Heo 
468daf8e166STejun Heo 	ret = scx_set_cmask_scratch_alloc(sch);
469daf8e166STejun Heo 	if (ret)
470daf8e166STejun Heo 		goto err_disable;
471daf8e166STejun Heo 
472daf8e166STejun Heo 	if (scx_validate_ops(sch, ops))
473daf8e166STejun Heo 		goto err_disable;
474daf8e166STejun Heo 
475daf8e166STejun Heo 	struct scx_sub_attach_args sub_attach_args = {
476daf8e166STejun Heo 		.ops = &sch->ops,
477daf8e166STejun Heo 		.cgroup_path = sch->cgrp_path,
478daf8e166STejun Heo 	};
479daf8e166STejun Heo 
480daf8e166STejun Heo 	ret = SCX_CALL_OP_RET(parent, sub_attach, NULL,
481daf8e166STejun Heo 			      &sub_attach_args);
482daf8e166STejun Heo 	if (ret) {
483daf8e166STejun Heo 		ret = scx_ops_sanitize_err(sch, "sub_attach", ret);
484daf8e166STejun Heo 		scx_error(sch, "parent rejected (%d)", ret);
485daf8e166STejun Heo 		goto err_disable;
486daf8e166STejun Heo 	}
487daf8e166STejun Heo 	sch->sub_attached = true;
488daf8e166STejun Heo 
489daf8e166STejun Heo 	scx_bypass(sch, true);
490daf8e166STejun Heo 
491daf8e166STejun Heo 	for (i = SCX_OPI_BEGIN; i < SCX_OPI_END; i++)
492daf8e166STejun Heo 		if (((void (**)(void))ops)[i])
493daf8e166STejun Heo 			set_bit(i, sch->has_op);
494daf8e166STejun Heo 
495daf8e166STejun Heo 	percpu_down_write(&scx_fork_rwsem);
496daf8e166STejun Heo 	scx_cgroup_lock();
497daf8e166STejun Heo 
498daf8e166STejun Heo 	/*
499daf8e166STejun Heo 	 * Set cgroup->scx_sched's and check CSS_ONLINE. Either we see
500daf8e166STejun Heo 	 * !CSS_ONLINE or scx_cgroup_lifetime_notify() sees and shoots us down.
501daf8e166STejun Heo 	 */
502daf8e166STejun Heo 	set_cgroup_sched(sch_cgroup(sch), sch);
503daf8e166STejun Heo 	if (!(cgrp->self.flags & CSS_ONLINE)) {
504daf8e166STejun Heo 		scx_error(sch, "cgroup is not online");
505daf8e166STejun Heo 		goto err_unlock_and_disable;
506daf8e166STejun Heo 	}
507daf8e166STejun Heo 
508daf8e166STejun Heo 	/*
509daf8e166STejun Heo 	 * Initialize tasks for the new child $sch without exiting them for
510daf8e166STejun Heo 	 * $parent so that the tasks can always be reverted back to $parent
511daf8e166STejun Heo 	 * sched on child init failure.
512daf8e166STejun Heo 	 */
513daf8e166STejun Heo 	WARN_ON_ONCE(scx_enabling_sub_sched);
514daf8e166STejun Heo 	scx_enabling_sub_sched = sch;
515daf8e166STejun Heo 
516daf8e166STejun Heo 	scx_task_iter_start(&sti, sch->cgrp);
517daf8e166STejun Heo 	while ((p = scx_task_iter_next_locked(&sti))) {
518daf8e166STejun Heo 		struct rq *rq;
519daf8e166STejun Heo 		struct rq_flags rf;
520daf8e166STejun Heo 
521daf8e166STejun Heo 		/*
522daf8e166STejun Heo 		 * Task iteration may visit the same task twice when racing
523daf8e166STejun Heo 		 * against exiting. Use %SCX_TASK_SUB_INIT to mark tasks which
524daf8e166STejun Heo 		 * finished __scx_init_task() and skip if set.
525daf8e166STejun Heo 		 *
526daf8e166STejun Heo 		 * A task may exit and get freed between __scx_init_task()
527daf8e166STejun Heo 		 * completion and scx_enable_task(). In such cases,
528daf8e166STejun Heo 		 * scx_disable_and_exit_task() must exit the task for both the
529daf8e166STejun Heo 		 * parent and child scheds.
530daf8e166STejun Heo 		 */
531daf8e166STejun Heo 		if (p->scx.flags & SCX_TASK_SUB_INIT)
532daf8e166STejun Heo 			continue;
533daf8e166STejun Heo 
534daf8e166STejun Heo 		/* @p is pinned by the iter; see scx_sub_disable() */
535daf8e166STejun Heo 		get_task_struct(p);
536daf8e166STejun Heo 
537daf8e166STejun Heo 		if (!assert_task_ready_or_enabled(p)) {
538daf8e166STejun Heo 			ret = -EINVAL;
539daf8e166STejun Heo 			goto abort;
540daf8e166STejun Heo 		}
541daf8e166STejun Heo 
542daf8e166STejun Heo 		scx_task_iter_unlock(&sti);
543daf8e166STejun Heo 
544daf8e166STejun Heo 		/*
545daf8e166STejun Heo 		 * As $p is still on $parent, it can't be transitioned to INIT.
546daf8e166STejun Heo 		 * Let's worry about task state later. Use __scx_init_task().
547daf8e166STejun Heo 		 */
548daf8e166STejun Heo 		ret = __scx_init_task(sch, p, false);
549daf8e166STejun Heo 		if (ret)
550daf8e166STejun Heo 			goto abort;
551daf8e166STejun Heo 
552daf8e166STejun Heo 		rq = task_rq_lock(p, &rf);
553daf8e166STejun Heo 
554daf8e166STejun Heo 		if (scx_get_task_state(p) == SCX_TASK_DEAD) {
555daf8e166STejun Heo 			/*
556daf8e166STejun Heo 			 * sched_ext_dead() raced us between __scx_init_task()
557daf8e166STejun Heo 			 * and this rq lock and ran exit_task() on $parent (the
558daf8e166STejun Heo 			 * sched @p was on at that point), not on @sch. @sch's
559daf8e166STejun Heo 			 * just-completed init is owed an exit_task() and we
560daf8e166STejun Heo 			 * issue it here.
561daf8e166STejun Heo 			 */
562daf8e166STejun Heo 			scx_sub_init_cancel_task(sch, p);
563daf8e166STejun Heo 			task_rq_unlock(rq, p, &rf);
564daf8e166STejun Heo 			put_task_struct(p);
565daf8e166STejun Heo 			continue;
566daf8e166STejun Heo 		}
567daf8e166STejun Heo 
568daf8e166STejun Heo 		p->scx.flags |= SCX_TASK_SUB_INIT;
569daf8e166STejun Heo 		task_rq_unlock(rq, p, &rf);
570daf8e166STejun Heo 
571daf8e166STejun Heo 		put_task_struct(p);
572daf8e166STejun Heo 	}
573daf8e166STejun Heo 	scx_task_iter_stop(&sti);
574daf8e166STejun Heo 
575daf8e166STejun Heo 	/*
576daf8e166STejun Heo 	 * All tasks are prepped. Disable/exit tasks for $parent and enable for
577daf8e166STejun Heo 	 * the new @sch.
578daf8e166STejun Heo 	 */
579daf8e166STejun Heo 	scx_task_iter_start(&sti, sch->cgrp);
580daf8e166STejun Heo 	while ((p = scx_task_iter_next_locked(&sti))) {
581daf8e166STejun Heo 		/*
582daf8e166STejun Heo 		 * Use clearing of %SCX_TASK_SUB_INIT to detect and skip
583daf8e166STejun Heo 		 * duplicate iterations.
584daf8e166STejun Heo 		 */
585daf8e166STejun Heo 		if (!(p->scx.flags & SCX_TASK_SUB_INIT))
586daf8e166STejun Heo 			continue;
587daf8e166STejun Heo 
588daf8e166STejun Heo 		scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
589daf8e166STejun Heo 			/*
590daf8e166STejun Heo 			 * $p must be either READY or ENABLED. If ENABLED,
591daf8e166STejun Heo 			 * __scx_disabled_and_exit_task() first disables and
592daf8e166STejun Heo 			 * makes it READY. However, after exiting $p, it will
593daf8e166STejun Heo 			 * leave $p as READY.
594daf8e166STejun Heo 			 */
595daf8e166STejun Heo 			assert_task_ready_or_enabled(p);
596daf8e166STejun Heo 			__scx_disable_and_exit_task(parent, p);
597daf8e166STejun Heo 
598daf8e166STejun Heo 			/*
599daf8e166STejun Heo 			 * $p is now only initialized for @sch and READY, which
600daf8e166STejun Heo 			 * is what we want. Assign it to @sch and enable.
601daf8e166STejun Heo 			 */
602daf8e166STejun Heo 			scx_set_task_sched(p, sch);
603daf8e166STejun Heo 			scx_enable_task(sch, p);
604daf8e166STejun Heo 
605daf8e166STejun Heo 			p->scx.flags &= ~SCX_TASK_SUB_INIT;
606daf8e166STejun Heo 		}
607daf8e166STejun Heo 	}
608daf8e166STejun Heo 	scx_task_iter_stop(&sti);
609daf8e166STejun Heo 
610daf8e166STejun Heo 	scx_enabling_sub_sched = NULL;
611daf8e166STejun Heo 
612daf8e166STejun Heo 	scx_cgroup_unlock();
613daf8e166STejun Heo 	percpu_up_write(&scx_fork_rwsem);
614daf8e166STejun Heo 
615daf8e166STejun Heo 	scx_bypass(sch, false);
616daf8e166STejun Heo 
617daf8e166STejun Heo 	pr_info("sched_ext: BPF sub-scheduler \"%s\" enabled\n", sch->ops.name);
618daf8e166STejun Heo 	kobject_uevent(&sch->kobj, KOBJ_ADD);
619daf8e166STejun Heo 	ret = 0;
620daf8e166STejun Heo 	goto out_unlock;
621daf8e166STejun Heo 
622daf8e166STejun Heo out_put_cgrp:
623daf8e166STejun Heo 	cgroup_put(cgrp);
624daf8e166STejun Heo out_unlock:
625daf8e166STejun Heo 	mutex_unlock(&scx_enable_mutex);
626daf8e166STejun Heo 	cmd->ret = ret;
627daf8e166STejun Heo 	return;
628daf8e166STejun Heo 
629daf8e166STejun Heo abort:
630daf8e166STejun Heo 	put_task_struct(p);
631daf8e166STejun Heo 	scx_task_iter_stop(&sti);
632daf8e166STejun Heo 
633daf8e166STejun Heo 	/*
634daf8e166STejun Heo 	 * Undo __scx_init_task() for tasks we marked. scx_enable_task() never
635daf8e166STejun Heo 	 * ran for @sch on them, so calling scx_disable_task() here would invoke
636daf8e166STejun Heo 	 * ops.disable() without a matching ops.enable(). scx_enabling_sub_sched
637daf8e166STejun Heo 	 * must stay set until SUB_INIT is cleared from every marked task -
638daf8e166STejun Heo 	 * scx_disable_and_exit_task() reads it when a task exits concurrently.
639daf8e166STejun Heo 	 */
640daf8e166STejun Heo 	scx_task_iter_start(&sti, sch->cgrp);
641daf8e166STejun Heo 	while ((p = scx_task_iter_next_locked(&sti))) {
642daf8e166STejun Heo 		if (p->scx.flags & SCX_TASK_SUB_INIT) {
643daf8e166STejun Heo 			scx_sub_init_cancel_task(sch, p);
644daf8e166STejun Heo 			p->scx.flags &= ~SCX_TASK_SUB_INIT;
645daf8e166STejun Heo 		}
646daf8e166STejun Heo 	}
647daf8e166STejun Heo 	scx_task_iter_stop(&sti);
648daf8e166STejun Heo 	scx_enabling_sub_sched = NULL;
649daf8e166STejun Heo err_unlock_and_disable:
650daf8e166STejun Heo 	/* we'll soon enter disable path, keep bypass on */
651daf8e166STejun Heo 	scx_cgroup_unlock();
652daf8e166STejun Heo 	percpu_up_write(&scx_fork_rwsem);
653daf8e166STejun Heo err_disable:
654daf8e166STejun Heo 	mutex_unlock(&scx_enable_mutex);
655ad45691dSTejun Heo 	/*
656ad45691dSTejun Heo 	 * Some enable failures only return an errno (e.g. -ENOMEM from an
657ad45691dSTejun Heo 	 * allocation) without calling scx_error(). Record it so
658ad45691dSTejun Heo 	 * scx_flush_disable_work() runs the disable and ops.exit() fires.
659ad45691dSTejun Heo 	 */
660ad45691dSTejun Heo 	scx_error(sch, "scx_sub_enable() failed (%d)", ret);
661daf8e166STejun Heo 	scx_flush_disable_work(sch);
662daf8e166STejun Heo 	cmd->ret = 0;
663daf8e166STejun Heo }
664daf8e166STejun Heo 
665daf8e166STejun Heo static s32 scx_cgroup_lifetime_notify(struct notifier_block *nb,
666daf8e166STejun Heo 				      unsigned long action, void *data)
667daf8e166STejun Heo {
668daf8e166STejun Heo 	struct cgroup *cgrp = data;
669daf8e166STejun Heo 	struct cgroup *parent = cgroup_parent(cgrp);
670daf8e166STejun Heo 
671daf8e166STejun Heo 	if (!cgroup_on_dfl(cgrp))
672daf8e166STejun Heo 		return NOTIFY_OK;
673daf8e166STejun Heo 
674daf8e166STejun Heo 	switch (action) {
675daf8e166STejun Heo 	case CGROUP_LIFETIME_ONLINE:
676daf8e166STejun Heo 		/* inherit ->scx_sched from $parent */
677daf8e166STejun Heo 		if (parent)
678daf8e166STejun Heo 			rcu_assign_pointer(cgrp->scx_sched, parent->scx_sched);
679daf8e166STejun Heo 		break;
680daf8e166STejun Heo 	case CGROUP_LIFETIME_OFFLINE:
681daf8e166STejun Heo 		/* if there is a sched attached, shoot it down */
682daf8e166STejun Heo 		if (cgrp->scx_sched && cgrp->scx_sched->cgrp == cgrp)
683daf8e166STejun Heo 			scx_exit(cgrp->scx_sched, SCX_EXIT_UNREG_KERN,
684daf8e166STejun Heo 				 SCX_ECODE_RSN_CGROUP_OFFLINE,
685daf8e166STejun Heo 				 "cgroup %llu going offline", cgroup_id(cgrp));
686daf8e166STejun Heo 		break;
687daf8e166STejun Heo 	}
688daf8e166STejun Heo 
689daf8e166STejun Heo 	return NOTIFY_OK;
690daf8e166STejun Heo }
691daf8e166STejun Heo 
692daf8e166STejun Heo static struct notifier_block scx_cgroup_lifetime_nb = {
693daf8e166STejun Heo 	.notifier_call = scx_cgroup_lifetime_notify,
694daf8e166STejun Heo };
695daf8e166STejun Heo 
696daf8e166STejun Heo static s32 __init scx_cgroup_lifetime_notifier_init(void)
697daf8e166STejun Heo {
698daf8e166STejun Heo 	return blocking_notifier_chain_register(&cgroup_lifetime_notifier,
699daf8e166STejun Heo 						&scx_cgroup_lifetime_nb);
700daf8e166STejun Heo }
701daf8e166STejun Heo core_initcall(scx_cgroup_lifetime_notifier_init);
702daf8e166STejun Heo 
703daf8e166STejun Heo void scx_pstack_recursion_on_dispatch(struct bpf_prog *prog)
704daf8e166STejun Heo {
705daf8e166STejun Heo 	struct scx_sched *sch;
706daf8e166STejun Heo 
707daf8e166STejun Heo 	guard(rcu)();
708daf8e166STejun Heo 	sch = scx_prog_sched(prog->aux);
709daf8e166STejun Heo 	if (unlikely(!sch))
710daf8e166STejun Heo 		return;
711daf8e166STejun Heo 
712daf8e166STejun Heo 	scx_error(sch, "dispatch recursion detected");
713daf8e166STejun Heo }
714daf8e166STejun Heo 
715daf8e166STejun Heo __bpf_kfunc_start_defs();
716daf8e166STejun Heo 
717daf8e166STejun Heo /**
718daf8e166STejun Heo  * scx_bpf_sub_dispatch - Trigger dispatching on a child scheduler
719daf8e166STejun Heo  * @cgroup_id: cgroup ID of the child scheduler to dispatch
720daf8e166STejun Heo  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
721daf8e166STejun Heo  *
722daf8e166STejun Heo  * Allows a parent scheduler to trigger dispatching on one of its direct
723daf8e166STejun Heo  * child schedulers. The child scheduler runs its dispatch operation to
724daf8e166STejun Heo  * move tasks from dispatch queues to the local runqueue.
725daf8e166STejun Heo  *
726daf8e166STejun Heo  * Returns: true on success, false if cgroup_id is invalid, not a direct
727daf8e166STejun Heo  * child, or caller lacks dispatch permission.
728daf8e166STejun Heo  */
729daf8e166STejun Heo __bpf_kfunc bool scx_bpf_sub_dispatch(u64 cgroup_id, const struct bpf_prog_aux *aux)
730daf8e166STejun Heo {
731daf8e166STejun Heo 	struct rq *this_rq = this_rq();
732daf8e166STejun Heo 	struct scx_sched *parent, *child;
733daf8e166STejun Heo 
734daf8e166STejun Heo 	guard(rcu)();
735daf8e166STejun Heo 	parent = scx_prog_sched(aux);
736daf8e166STejun Heo 	if (unlikely(!parent))
737daf8e166STejun Heo 		return false;
738daf8e166STejun Heo 
739daf8e166STejun Heo 	child = scx_find_sub_sched(cgroup_id);
740daf8e166STejun Heo 
741daf8e166STejun Heo 	if (unlikely(!child))
742daf8e166STejun Heo 		return false;
743daf8e166STejun Heo 
744daf8e166STejun Heo 	if (unlikely(scx_parent(child) != parent)) {
745daf8e166STejun Heo 		scx_error(parent, "trying to dispatch a distant sub-sched on cgroup %llu",
746daf8e166STejun Heo 			  cgroup_id);
747daf8e166STejun Heo 		return false;
748daf8e166STejun Heo 	}
749daf8e166STejun Heo 
750daf8e166STejun Heo 	return scx_dispatch_sched(child, this_rq, this_rq->scx.sub_dispatch_prev,
751daf8e166STejun Heo 				  true);
752daf8e166STejun Heo }
753daf8e166STejun Heo 
754daf8e166STejun Heo __bpf_kfunc_end_defs();
755daf8e166STejun Heo 
756daf8e166STejun Heo #endif	/* CONFIG_EXT_SUB_SCHED */
757