xref: /linux/tools/sched_ext/scx_pair.bpf.c (revision 11260c335ec6071af5543aef73000b28f041c124)
1 /* SPDX-License-Identifier: GPL-2.0 */
2 /*
3  * A demo sched_ext core-scheduler which always makes every sibling CPU pair
4  * execute from the same CPU cgroup.
5  *
6  * This scheduler is a minimal implementation and would need some form of
7  * priority handling both inside each cgroup and across the cgroups to be
8  * practically useful.
9  *
10  * Each CPU in the system is paired with exactly one other CPU, according to a
11  * "stride" value that can be specified when the BPF scheduler program is first
12  * loaded. Throughout the runtime of the scheduler, these CPU pairs guarantee
13  * that they will only ever schedule tasks that belong to the same CPU cgroup.
14  *
15  * Scheduler Initialization
16  * ------------------------
17  *
18  * The scheduler BPF program is first initialized from user space, before it is
19  * enabled. During this initialization process, each CPU on the system is
20  * assigned several values that are constant throughout its runtime:
21  *
22  * 1. *Pair CPU*: The CPU that it synchronizes with when making scheduling
23  *		  decisions. Paired CPUs always schedule tasks from the same
24  *		  CPU cgroup, and synchronize with each other to guarantee
25  *		  that this constraint is not violated.
26  * 2. *Pair ID*:  Each CPU pair is assigned a Pair ID, which is used to access
27  *		  a struct pair_ctx object that is shared between the pair.
28  * 3. *In-pair-index*: An index, 0 or 1, that is assigned to each core in the
29  *		       pair. Each struct pair_ctx has an active_mask field,
30  *		       which is a bitmap used to indicate whether each core
31  *		       in the pair currently has an actively running task.
32  *		       This index specifies which entry in the bitmap corresponds
33  *		       to each CPU in the pair.
34  *
35  * During this initialization, the CPUs are paired according to a "stride" that
36  * may be specified when invoking the user space program that initializes and
37  * loads the scheduler. By default, the stride is 1/2 the total number of CPUs.
38  *
39  * Tasks and cgroups
40  * -----------------
41  *
42  * Every cgroup in the system is registered with the scheduler using the
43  * pair_cgroup_init() callback, and every task in the system is associated with
44  * exactly one cgroup. At a high level, the idea with the pair scheduler is to
45  * always schedule tasks from the same cgroup within a given CPU pair. When a
46  * task is enqueued (i.e. passed to the pair_enqueue() callback function), its
47  * cgroup ID is read from its task struct, and then a corresponding queue map
48  * is used to FIFO-enqueue the task for that cgroup.
49  *
50  * If you look through the implementation of the scheduler, you'll notice that
51  * there is quite a bit of complexity involved with looking up the per-cgroup
52  * FIFO queue that we enqueue tasks in. For example, there is a cgrp_q_idx_hash
53  * BPF hash map that is used to map a cgroup ID to a globally unique ID that's
54  * allocated in the BPF program. This is done because we use separate maps to
55  * store the FIFO queue of tasks, and the length of that map, per cgroup. This
56  * complexity is only present because of current deficiencies in BPF that will
57  * soon be addressed. The main point to keep in mind is that newly enqueued
58  * tasks are added to their cgroup's FIFO queue.
59  *
60  * Dispatching tasks
61  * -----------------
62  *
63  * This section will describe how enqueued tasks are dispatched and scheduled.
64  * Tasks are dispatched in pair_dispatch(), and at a high level the workflow is
65  * as follows:
66  *
67  * 1. Fetch the struct pair_ctx for the current CPU. As mentioned above, this is
68  *    the structure that's used to synchronize amongst the two pair CPUs in their
69  *    scheduling decisions. After any of the following events have occurred:
70  *
71  * - The cgroup's slice run has expired, or
72  * - The cgroup becomes empty, or
73  * - Either CPU in the pair is preempted by a higher priority scheduling class
74  *
75  * The cgroup transitions to the draining state and stops executing new tasks
76  * from the cgroup.
77  *
78  * 2. If the pair is still executing a task, mark the pair_ctx as draining, and
79  *    wait for the pair CPU to be preempted.
80  *
81  * 3. Otherwise, if the pair CPU is not running a task, we can move onto
82  *    scheduling new tasks. Pop the next cgroup id from the top_q queue.
83  *
84  * 4. Pop a task from that cgroup's FIFO task queue, and begin executing it.
85  *
86  * Note again that this scheduling behavior is simple, but the implementation
87  * is complex mostly because this it hits several BPF shortcomings and has to
88  * work around in often awkward ways. Most of the shortcomings are expected to
89  * be resolved in the near future which should allow greatly simplifying this
90  * scheduler.
91  *
92  * Dealing with preemption
93  * -----------------------
94  *
95  * SCX is the lowest priority sched_class, and could be preempted by them at
96  * any time. To address this, the scheduler watches every sched_switch from
97  * a tracepoint and edge-detects when a CPU leaves and returns to SCX
98  * control.
99  *
100  * When a higher-priority class takes a CPU away from a running SCX task -
101  * a sched_switch from an SCX task to a higher-priority task - we mark the
102  * pair_ctx as having been preempted and then invoke:
103  *
104  * scx_bpf_kick_cpu(pair_cpu, SCX_KICK_PREEMPT | SCX_KICK_WAIT);
105  *
106  * This preempts the pair CPU, and waits until it has re-entered the scheduler
107  * before returning. This is necessary to ensure that the higher priority
108  * sched_class that preempted our scheduler does not schedule a task
109  * concurrently with our pair CPU.
110  *
111  * When the CPU returns to SCX or idle, we unmark the preemption in the
112  * pair_ctx and send another resched IPI to the pair CPU to re-enable pair
113  * scheduling.
114  *
115  * A switch from idle straight to a higher-priority task is not a release:
116  * the CPU was not running an SCX task, so there is nothing to drain and no
117  * reason to make the pair wait. Kicking SCX_KICK_WAIT on every such wakeup
118  * would stall the pair CPU behind rt bursts it was never coupled to.
119  *
120  * Note that sched_setscheduler() on a running task changes its class in
121  * place without a context switch, so such transitions are only observed at
122  * the task's next switch. Until then the stale active_mask bit makes the
123  * pair wait in try_dispatch(), which is bounded by that next switch.
124  *
125  * Copyright (c) 2022 Meta Platforms, Inc. and affiliates.
126  * Copyright (c) 2022 Tejun Heo <tj@kernel.org>
127  * Copyright (c) 2022 David Vernet <dvernet@meta.com>
128  */
129 #include <scx/common.bpf.h>
130 #include "scx_pair.h"
131 
132 #define MAX_RT_PRIO	100
133 
134 char _license[] SEC("license") = "GPL";
135 
136 /* !0 for veristat, set during init */
137 const volatile u32 nr_cpu_ids = 1;
138 
139 /* a pair of CPUs stay on a cgroup for this duration */
140 const volatile u32 pair_batch_dur_ns;
141 
142 /* cpu ID -> pair cpu ID */
143 const volatile s32 RESIZABLE_ARRAY(rodata, pair_cpu);
144 
145 /* cpu ID -> pair_id */
146 const volatile u32 RESIZABLE_ARRAY(rodata, pair_id);
147 
148 /* CPU ID -> CPU # in the pair (0 or 1) */
149 const volatile u32 RESIZABLE_ARRAY(rodata, in_pair_idx);
150 
151 struct pair_ctx {
152 	struct bpf_spin_lock	lock;
153 
154 	/* the cgroup the pair is currently executing */
155 	u64			cgid;
156 
157 	/* the pair started executing the current cgroup at */
158 	u64			started_at;
159 
160 	/* whether the current cgroup is draining */
161 	bool			draining;
162 
163 	/* the CPUs that are currently active on the cgroup */
164 	u32			active_mask;
165 
166 	/*
167 	 * the CPUs that are currently preempted and running tasks in a
168 	 * different scheduler.
169 	 */
170 	u32			preempted_mask;
171 };
172 
173 struct {
174 	__uint(type, BPF_MAP_TYPE_ARRAY);
175 	__type(key, u32);
176 	__type(value, struct pair_ctx);
177 } pair_ctx SEC(".maps");
178 
179 /* queue of cgrp_q's possibly with tasks on them */
180 struct {
181 	__uint(type, BPF_MAP_TYPE_QUEUE);
182 	/*
183 	 * Because it's difficult to build strong synchronization encompassing
184 	 * multiple non-trivial operations in BPF, this queue is managed in an
185 	 * opportunistic way so that we guarantee that a cgroup w/ active tasks
186 	 * is always on it but possibly multiple times. Once we have more robust
187 	 * synchronization constructs and e.g. linked list, we should be able to
188 	 * do this in a prettier way but for now just size it big enough.
189 	 */
190 	__uint(max_entries, 4 * MAX_CGRPS);
191 	__type(value, u64);
192 } top_q SEC(".maps");
193 
194 /* per-cgroup q which FIFOs the tasks from the cgroup */
195 struct cgrp_q {
196 	__uint(type, BPF_MAP_TYPE_QUEUE);
197 	__uint(max_entries, MAX_QUEUED);
198 	__type(value, u32);
199 };
200 
201 /*
202  * Ideally, we want to allocate cgrp_q and cgrq_q_len in the cgroup local
203  * storage; however, a cgroup local storage can only be accessed from the BPF
204  * progs attached to the cgroup. For now, work around by allocating array of
205  * cgrp_q's and then allocating per-cgroup indices.
206  *
207  * Another caveat: It's difficult to populate a large array of maps statically
208  * or from BPF. Initialize it from userland.
209  */
210 struct {
211 	__uint(type, BPF_MAP_TYPE_ARRAY_OF_MAPS);
212 	__uint(max_entries, MAX_CGRPS);
213 	__type(key, s32);
214 	__array(values, struct cgrp_q);
215 } cgrp_q_arr SEC(".maps");
216 
217 static u64 cgrp_q_len[MAX_CGRPS];
218 
219 /*
220  * This and cgrp_q_idx_hash combine into a poor man's IDR. This likely would be
221  * useful to have as a map type.
222  */
223 static u32 cgrp_q_idx_cursor;
224 static u64 cgrp_q_idx_busy[MAX_CGRPS];
225 
226 /*
227  * All added up, the following is what we do:
228  *
229  * 1. When a cgroup is enabled, RR cgroup_q_idx_busy array doing cmpxchg looking
230  *    for a free ID. If not found, fail cgroup creation with -EBUSY.
231  *
232  * 2. Hash the cgroup ID to the allocated cgrp_q_idx in the following
233  *    cgrp_q_idx_hash.
234  *
235  * 3. Whenever a cgrp_q needs to be accessed, first look up the cgrp_q_idx from
236  *    cgrp_q_idx_hash and then access the corresponding entry in cgrp_q_arr.
237  *
238  * This is sadly complicated for something pretty simple. Hopefully, we should
239  * be able to simplify in the future.
240  */
241 struct {
242 	__uint(type, BPF_MAP_TYPE_HASH);
243 	__uint(max_entries, MAX_CGRPS);
244 	__uint(key_size, sizeof(u64));		/* cgrp ID */
245 	__uint(value_size, sizeof(s32));	/* cgrp_q idx */
246 } cgrp_q_idx_hash SEC(".maps");
247 
248 /* statistics */
249 u64 nr_total, nr_dispatched, nr_missing, nr_kicks, nr_preemptions;
250 u64 nr_exps, nr_exp_waits, nr_exp_empty;
251 u64 nr_cgrp_next, nr_cgrp_coll, nr_cgrp_empty;
252 
253 UEI_DEFINE(uei);
254 
BPF_STRUCT_OPS(pair_enqueue,struct task_struct * p,u64 enq_flags)255 void BPF_STRUCT_OPS(pair_enqueue, struct task_struct *p, u64 enq_flags)
256 {
257 	struct cgroup *cgrp;
258 	struct cgrp_q *cgq;
259 	s32 pid = p->pid;
260 	u64 cgid;
261 	u32 *q_idx;
262 	u64 *cgq_len;
263 
264 	__sync_fetch_and_add(&nr_total, 1);
265 
266 	cgrp = scx_bpf_task_cgroup(p);
267 	cgid = cgrp->kn->id;
268 	bpf_cgroup_release(cgrp);
269 
270 	/* find the cgroup's q and push @p into it */
271 	q_idx = bpf_map_lookup_elem(&cgrp_q_idx_hash, &cgid);
272 	if (!q_idx) {
273 		scx_bpf_error("failed to lookup q_idx for cgroup[%llu]", cgid);
274 		return;
275 	}
276 
277 	cgq = bpf_map_lookup_elem(&cgrp_q_arr, q_idx);
278 	if (!cgq) {
279 		scx_bpf_error("failed to lookup q_arr for cgroup[%llu] q_idx[%u]",
280 			      cgid, *q_idx);
281 		return;
282 	}
283 
284 	if (bpf_map_push_elem(cgq, &pid, 0)) {
285 		scx_bpf_error("cgroup[%llu] queue overflow", cgid);
286 		return;
287 	}
288 
289 	/* bump q len, if going 0 -> 1, queue cgroup into the top_q */
290 	cgq_len = MEMBER_VPTR(cgrp_q_len, [*q_idx]);
291 	if (!cgq_len) {
292 		scx_bpf_error("MEMBER_VTPR malfunction");
293 		return;
294 	}
295 
296 	if (!__sync_fetch_and_add(cgq_len, 1) &&
297 	    bpf_map_push_elem(&top_q, &cgid, 0)) {
298 		scx_bpf_error("top_q overflow");
299 		return;
300 	}
301 }
302 
lookup_pairc_and_mask(s32 cpu,struct pair_ctx ** pairc,u32 * mask)303 static int lookup_pairc_and_mask(s32 cpu, struct pair_ctx **pairc, u32 *mask)
304 {
305 	u32 *vptr;
306 
307 	vptr = (u32 *)ARRAY_ELEM_PTR(pair_id, cpu, nr_cpu_ids);
308 	if (!vptr)
309 		return -EINVAL;
310 
311 	*pairc = bpf_map_lookup_elem(&pair_ctx, vptr);
312 	if (!(*pairc))
313 		return -EINVAL;
314 
315 	vptr = (u32 *)ARRAY_ELEM_PTR(in_pair_idx, cpu, nr_cpu_ids);
316 	if (!vptr)
317 		return -EINVAL;
318 
319 	*mask = 1U << *vptr;
320 
321 	return 0;
322 }
323 
324 /*
325  * A task is above SCX whenever its effective priority is in the rt/dl
326  * range. Test p->prio rather than p->policy: rt_mutex_setprio() boosts
327  * a PI beneficiary into the rt/dl classes with its policy left
328  * untouched, so a policy test would misclassify boosted tasks in both
329  * directions. p->prio follows the boost and the deboost.
330  *
331  * This still cannot tell fair and SCX tasks apart. It is complete only
332  * because scx_pair runs in switch-all mode, where no fair class task
333  * exists; in partial mode fair is also above SCX and can take the CPU.
334  */
pair_task_is_highpri(struct task_struct * p)335 static bool pair_task_is_highpri(struct task_struct *p)
336 {
337 	return p->prio < MAX_RT_PRIO;
338 }
339 
pair_cpu_acquire_locked(struct pair_ctx * pairc,u32 in_pair_mask,u32 * kick_flags)340 static void pair_cpu_acquire_locked(struct pair_ctx *pairc, u32 in_pair_mask,
341 					u32 *kick_flags)
342 {
343 	pairc->preempted_mask &= ~in_pair_mask;
344 	/* Kick the pair CPU, unless it was also preempted. */
345 	*kick_flags = !pairc->preempted_mask ? SCX_KICK_PREEMPT : 0;
346 }
347 
pair_cpu_release_locked(struct pair_ctx * pairc,u32 in_pair_mask,u32 * kick_flags)348 static void pair_cpu_release_locked(struct pair_ctx *pairc, u32 in_pair_mask,
349 					u32 *kick_flags)
350 {
351 	pairc->preempted_mask |= in_pair_mask;
352 	pairc->active_mask &= ~in_pair_mask;
353 	/* Kick the pair CPU if it's still running. */
354 	*kick_flags = pairc->active_mask ? SCX_KICK_PREEMPT | SCX_KICK_WAIT : 0;
355 	pairc->draining = true;
356 }
357 
358 __attribute__((noinline))
try_dispatch(s32 cpu)359 static int try_dispatch(s32 cpu)
360 {
361 	struct pair_ctx *pairc;
362 	struct bpf_map *cgq_map;
363 	struct task_struct *p;
364 	u64 now = scx_bpf_now();
365 	bool kick_pair = false;
366 	bool expired, pair_preempted;
367 	u32 *vptr, in_pair_mask;
368 	s32 pid, q_idx;
369 	u64 cgid;
370 	int ret;
371 
372 	ret = lookup_pairc_and_mask(cpu, &pairc, &in_pair_mask);
373 	if (ret) {
374 		scx_bpf_error("failed to lookup pairc and in_pair_mask for cpu[%d]",
375 			      cpu);
376 		return -ENOENT;
377 	}
378 
379 	bpf_spin_lock(&pairc->lock);
380 	pairc->active_mask &= ~in_pair_mask;
381 
382 	expired = time_before(pairc->started_at + pair_batch_dur_ns, now);
383 	if (expired || pairc->draining) {
384 		u64 new_cgid = 0;
385 
386 		__sync_fetch_and_add(&nr_exps, 1);
387 
388 		/*
389 		 * We're done with the current cgid. An obvious optimization
390 		 * would be not draining if the next cgroup is the current one.
391 		 * For now, be dumb and always expire.
392 		 */
393 		pairc->draining = true;
394 
395 		pair_preempted = pairc->preempted_mask;
396 		if (pairc->active_mask || pair_preempted) {
397 			/*
398 			 * The other CPU is still active, or is no longer under
399 			 * our control due to e.g. being preempted by a higher
400 			 * priority sched_class. We want to wait until this
401 			 * cgroup expires, or until control of our pair CPU has
402 			 * been returned to us.
403 			 *
404 			 * If the pair controls its CPU, and the time already
405 			 * expired, kick.  When the other CPU arrives at
406 			 * dispatch and clears its active mask, it'll push the
407 			 * pair to the next cgroup and kick this CPU.
408 			 */
409 			__sync_fetch_and_add(&nr_exp_waits, 1);
410 			bpf_spin_unlock(&pairc->lock);
411 			if (expired && !pair_preempted)
412 				kick_pair = true;
413 			goto out_maybe_kick;
414 		}
415 
416 		bpf_spin_unlock(&pairc->lock);
417 
418 		/*
419 		 * Pick the next cgroup. It'd be easier / cleaner to not drop
420 		 * pairc->lock and use stronger synchronization here especially
421 		 * given that we'll be switching cgroups significantly less
422 		 * frequently than tasks. Unfortunately, bpf_spin_lock can't
423 		 * really protect anything non-trivial. Let's do opportunistic
424 		 * operations instead.
425 		 */
426 		bpf_repeat(BPF_MAX_LOOPS) {
427 			u32 *q_idx;
428 			u64 *cgq_len;
429 
430 			if (bpf_map_pop_elem(&top_q, &new_cgid)) {
431 				/* no active cgroup, go idle */
432 				__sync_fetch_and_add(&nr_exp_empty, 1);
433 				return 0;
434 			}
435 
436 			q_idx = bpf_map_lookup_elem(&cgrp_q_idx_hash, &new_cgid);
437 			if (!q_idx)
438 				continue;
439 
440 			/*
441 			 * This is the only place where empty cgroups are taken
442 			 * off the top_q.
443 			 */
444 			cgq_len = MEMBER_VPTR(cgrp_q_len, [*q_idx]);
445 			if (!cgq_len || !*cgq_len)
446 				continue;
447 
448 			/*
449 			 * If it has any tasks, requeue as we may race and not
450 			 * execute it.
451 			 */
452 			bpf_map_push_elem(&top_q, &new_cgid, 0);
453 			break;
454 		}
455 
456 		bpf_spin_lock(&pairc->lock);
457 
458 		/*
459 		 * The other CPU may already have started on a new cgroup while
460 		 * we dropped the lock. Make sure that we're still draining and
461 		 * start on the new cgroup.
462 		 */
463 		if (pairc->draining && !pairc->active_mask) {
464 			__sync_fetch_and_add(&nr_cgrp_next, 1);
465 			pairc->cgid = new_cgid;
466 			pairc->started_at = now;
467 			pairc->draining = false;
468 			kick_pair = true;
469 		} else {
470 			__sync_fetch_and_add(&nr_cgrp_coll, 1);
471 		}
472 	}
473 
474 	cgid = pairc->cgid;
475 	pairc->active_mask |= in_pair_mask;
476 	bpf_spin_unlock(&pairc->lock);
477 
478 	/* again, it'd be better to do all these with the lock held, oh well */
479 	vptr = bpf_map_lookup_elem(&cgrp_q_idx_hash, &cgid);
480 	if (!vptr) {
481 		scx_bpf_error("failed to lookup q_idx for cgroup[%llu]", cgid);
482 		return -ENOENT;
483 	}
484 	q_idx = *vptr;
485 
486 	/* claim one task from cgrp_q w/ q_idx */
487 	bpf_repeat(BPF_MAX_LOOPS) {
488 		u64 *cgq_len, len;
489 
490 		cgq_len = MEMBER_VPTR(cgrp_q_len, [q_idx]);
491 		if (!cgq_len || !(len = *(volatile u64 *)cgq_len)) {
492 			/* the cgroup must be empty, expire and repeat */
493 			__sync_fetch_and_add(&nr_cgrp_empty, 1);
494 			bpf_spin_lock(&pairc->lock);
495 			pairc->draining = true;
496 			pairc->active_mask &= ~in_pair_mask;
497 			bpf_spin_unlock(&pairc->lock);
498 			return -EAGAIN;
499 		}
500 
501 		if (__sync_val_compare_and_swap(cgq_len, len, len - 1) != len)
502 			continue;
503 
504 		break;
505 	}
506 
507 	cgq_map = bpf_map_lookup_elem(&cgrp_q_arr, &q_idx);
508 	if (!cgq_map) {
509 		scx_bpf_error("failed to lookup cgq_map for cgroup[%llu] q_idx[%d]",
510 			      cgid, q_idx);
511 		return -ENOENT;
512 	}
513 
514 	if (bpf_map_pop_elem(cgq_map, &pid)) {
515 		scx_bpf_error("cgq_map is empty for cgroup[%llu] q_idx[%d]",
516 			      cgid, q_idx);
517 		return -ENOENT;
518 	}
519 
520 	p = bpf_task_from_pid(pid);
521 	if (p) {
522 		__sync_fetch_and_add(&nr_dispatched, 1);
523 		scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0);
524 		bpf_task_release(p);
525 	} else {
526 		/* we don't handle dequeues, retry on lost tasks */
527 		__sync_fetch_and_add(&nr_missing, 1);
528 		return -EAGAIN;
529 	}
530 
531 out_maybe_kick:
532 	if (kick_pair) {
533 		s32 *pair = (s32 *)ARRAY_ELEM_PTR(pair_cpu, cpu, nr_cpu_ids);
534 		if (pair) {
535 			__sync_fetch_and_add(&nr_kicks, 1);
536 			scx_bpf_kick_cpu(*pair, SCX_KICK_PREEMPT);
537 		}
538 	}
539 	return 0;
540 }
541 
BPF_STRUCT_OPS(pair_dispatch,s32 cpu,struct task_struct * prev)542 void BPF_STRUCT_OPS(pair_dispatch, s32 cpu, struct task_struct *prev)
543 {
544 	bpf_repeat(BPF_MAX_LOOPS) {
545 		if (try_dispatch(cpu) != -EAGAIN)
546 			break;
547 	}
548 }
549 
550 SEC("tp_btf/sched_switch")
BPF_PROG(pair_sched_switch,bool preempt,struct task_struct * prev,struct task_struct * next,unsigned int prev_state)551 int BPF_PROG(pair_sched_switch, bool preempt, struct task_struct *prev,
552 	     struct task_struct *next, unsigned int prev_state)
553 {
554 	int ret;
555 	s32 cpu = bpf_get_smp_processor_id();
556 	u32 in_pair_mask;
557 	struct pair_ctx *pairc;
558 	u32 kick_flags = 0;
559 	bool preempted;
560 	bool release, acquire;
561 
562 	ret = lookup_pairc_and_mask(cpu, &pairc, &in_pair_mask);
563 	if (ret)
564 		return 0;
565 
566 	/*
567 	 * This runs on every context switch in the system. A CPU's own
568 	 * preempted_mask bit is only ever written by this tracepoint
569 	 * running on that CPU, so the unlocked read is exact and the
570 	 * pair-shared lock is only taken on actual transitions.
571 	 */
572 	preempted = pairc->preempted_mask & in_pair_mask;
573 	if (next->pid && pair_task_is_highpri(next)) {
574 		/* an SCX task lost the CPU to a higher-priority class */
575 		release = !preempted && prev->pid && !pair_task_is_highpri(prev);
576 		acquire = false;
577 	} else {
578 		/* the CPU is back under SCX control (or idle) */
579 		release = false;
580 		acquire = preempted;
581 	}
582 	if (!release && !acquire)
583 		return 0;
584 
585 	bpf_spin_lock(&pairc->lock);
586 	if (release) {
587 		pair_cpu_release_locked(pairc, in_pair_mask, &kick_flags);
588 		__sync_fetch_and_add(&nr_preemptions, 1);
589 	} else {
590 		pair_cpu_acquire_locked(pairc, in_pair_mask, &kick_flags);
591 	}
592 	bpf_spin_unlock(&pairc->lock);
593 
594 	if (kick_flags) {
595 		s32 *pair = (s32 *)ARRAY_ELEM_PTR(pair_cpu, cpu, nr_cpu_ids);
596 
597 		if (pair) {
598 			__sync_fetch_and_add(&nr_kicks, 1);
599 			scx_bpf_kick_cpu(*pair, kick_flags);
600 		}
601 	}
602 
603 	return 0;
604 }
605 
BPF_STRUCT_OPS(pair_cgroup_init,struct cgroup * cgrp)606 s32 BPF_STRUCT_OPS(pair_cgroup_init, struct cgroup *cgrp)
607 {
608 	u64 cgid = cgrp->kn->id;
609 	s32 i, q_idx;
610 
611 	bpf_for(i, 0, MAX_CGRPS) {
612 		q_idx = __sync_fetch_and_add(&cgrp_q_idx_cursor, 1) % MAX_CGRPS;
613 		if (!__sync_val_compare_and_swap(&cgrp_q_idx_busy[q_idx], 0, 1))
614 			break;
615 	}
616 	if (i == MAX_CGRPS)
617 		return -EBUSY;
618 
619 	if (bpf_map_update_elem(&cgrp_q_idx_hash, &cgid, &q_idx, BPF_ANY)) {
620 		u64 *busy = MEMBER_VPTR(cgrp_q_idx_busy, [q_idx]);
621 		if (busy)
622 			*busy = 0;
623 		return -EBUSY;
624 	}
625 
626 	return 0;
627 }
628 
BPF_STRUCT_OPS(pair_cgroup_exit,struct cgroup * cgrp)629 void BPF_STRUCT_OPS(pair_cgroup_exit, struct cgroup *cgrp)
630 {
631 	u64 cgid = cgrp->kn->id;
632 	s32 *q_idx;
633 
634 	q_idx = bpf_map_lookup_elem(&cgrp_q_idx_hash, &cgid);
635 	if (q_idx) {
636 		u64 *busy = MEMBER_VPTR(cgrp_q_idx_busy, [*q_idx]);
637 		if (busy)
638 			*busy = 0;
639 		bpf_map_delete_elem(&cgrp_q_idx_hash, &cgid);
640 	}
641 }
642 
BPF_STRUCT_OPS(pair_exit,struct scx_exit_info * ei)643 void BPF_STRUCT_OPS(pair_exit, struct scx_exit_info *ei)
644 {
645 	UEI_RECORD(uei, ei);
646 }
647 
648 SCX_OPS_DEFINE(pair_ops,
649 	       .enqueue			= (void *)pair_enqueue,
650 	       .dispatch		= (void *)pair_dispatch,
651 	       .cgroup_init		= (void *)pair_cgroup_init,
652 	       .cgroup_exit		= (void *)pair_cgroup_exit,
653 	       .exit			= (void *)pair_exit,
654 	       .name			= "pair");
655