ext.c

   1/* SPDX-License-Identifier: GPL-2.0 */
   2/*
   3 * BPF extensible scheduler class: Documentation/scheduler/sched-ext.rst
   4 *
   5 * Copyright (c) 2022 Meta Platforms, Inc. and affiliates.
   6 * Copyright (c) 2022 Tejun Heo <tj@kernel.org>
   7 * Copyright (c) 2022 David Vernet <dvernet@meta.com>
   8 */
   9#include <linux/btf_ids.h>
  10#include "ext_idle.h"
  11
  12/*
  13 * NOTE: sched_ext is in the process of growing multiple scheduler support and
  14 * scx_root usage is in a transitional state. Naked dereferences are safe if the
  15 * caller is one of the tasks attached to SCX and explicit RCU dereference is
  16 * necessary otherwise. Naked scx_root dereferences trigger sparse warnings but
  17 * are used as temporary markers to indicate that the dereferences need to be
  18 * updated to point to the associated scheduler instances rather than scx_root.
  19 */
  20static struct scx_sched __rcu *scx_root;
  21
  22/*
  23 * During exit, a task may schedule after losing its PIDs. When disabling the
  24 * BPF scheduler, we need to be able to iterate tasks in every state to
  25 * guarantee system safety. Maintain a dedicated task list which contains every
  26 * task between its fork and eventual free.
  27 */
  28static DEFINE_RAW_SPINLOCK(scx_tasks_lock);
  29static LIST_HEAD(scx_tasks);
  30
  31/* ops enable/disable */
  32static DEFINE_MUTEX(scx_enable_mutex);
  33DEFINE_STATIC_KEY_FALSE(__scx_enabled);
  34DEFINE_STATIC_PERCPU_RWSEM(scx_fork_rwsem);
  35static atomic_t scx_enable_state_var = ATOMIC_INIT(SCX_DISABLED);
  36static int scx_bypass_depth;
  37static cpumask_var_t scx_bypass_lb_donee_cpumask;
  38static cpumask_var_t scx_bypass_lb_resched_cpumask;
  39static bool scx_aborting;
  40static bool scx_init_task_enabled;
  41static bool scx_switching_all;
  42DEFINE_STATIC_KEY_FALSE(__scx_switched_all);
  43
  44/*
  45 * Tracks whether scx_enable() called scx_bypass(true). Used to balance bypass
  46 * depth on enable failure. Will be removed when bypass depth is moved into the
  47 * sched instance.
  48 */
  49static bool scx_bypassed_for_enable;
  50
  51static atomic_long_t scx_nr_rejected = ATOMIC_LONG_INIT(0);
  52static atomic_long_t scx_hotplug_seq = ATOMIC_LONG_INIT(0);
  53
  54/*
  55 * A monotically increasing sequence number that is incremented every time a
  56 * scheduler is enabled. This can be used by to check if any custom sched_ext
  57 * scheduler has ever been used in the system.
  58 */
  59static atomic_long_t scx_enable_seq = ATOMIC_LONG_INIT(0);
  60
  61/*
  62 * The maximum amount of time in jiffies that a task may be runnable without
  63 * being scheduled on a CPU. If this timeout is exceeded, it will trigger
  64 * scx_error().
  65 */
  66static unsigned long scx_watchdog_timeout;
  67
  68/*
  69 * The last time the delayed work was run. This delayed work relies on
  70 * ksoftirqd being able to run to service timer interrupts, so it's possible
  71 * that this work itself could get wedged. To account for this, we check that
  72 * it's not stalled in the timer tick, and trigger an error if it is.
  73 */
  74static unsigned long scx_watchdog_timestamp = INITIAL_JIFFIES;
  75
  76static struct delayed_work scx_watchdog_work;
  77
  78/*
  79 * For %SCX_KICK_WAIT: Each CPU has a pointer to an array of kick_sync sequence
  80 * numbers. The arrays are allocated with kvzalloc() as size can exceed percpu
  81 * allocator limits on large machines. O(nr_cpu_ids^2) allocation, allocated
  82 * lazily when enabling and freed when disabling to avoid waste when sched_ext
  83 * isn't active.
  84 */
  85struct scx_kick_syncs {
  86	struct rcu_head		rcu;
  87	unsigned long		syncs[];
  88};
  89
  90static DEFINE_PER_CPU(struct scx_kick_syncs __rcu *, scx_kick_syncs);
  91
  92/*
  93 * Direct dispatch marker.
  94 *
  95 * Non-NULL values are used for direct dispatch from enqueue path. A valid
  96 * pointer points to the task currently being enqueued. An ERR_PTR value is used
  97 * to indicate that direct dispatch has already happened.
  98 */
  99static DEFINE_PER_CPU(struct task_struct *, direct_dispatch_task);
 100
 101static const struct rhashtable_params dsq_hash_params = {
 102	.key_len		= sizeof_field(struct scx_dispatch_q, id),
 103	.key_offset		= offsetof(struct scx_dispatch_q, id),
 104	.head_offset		= offsetof(struct scx_dispatch_q, hash_node),
 105};
 106
 107static LLIST_HEAD(dsqs_to_free);
 108
 109/* dispatch buf */
 110struct scx_dsp_buf_ent {
 111	struct task_struct	*task;
 112	unsigned long		qseq;
 113	u64			dsq_id;
 114	u64			enq_flags;
 115};
 116
 117static u32 scx_dsp_max_batch;
 118
 119struct scx_dsp_ctx {
 120	struct rq		*rq;
 121	u32			cursor;
 122	u32			nr_tasks;
 123	struct scx_dsp_buf_ent	buf[];
 124};
 125
 126static struct scx_dsp_ctx __percpu *scx_dsp_ctx;
 127
 128/* string formatting from BPF */
 129struct scx_bstr_buf {
 130	u64			data[MAX_BPRINTF_VARARGS];
 131	char			line[SCX_EXIT_MSG_LEN];
 132};
 133
 134static DEFINE_RAW_SPINLOCK(scx_exit_bstr_buf_lock);
 135static struct scx_bstr_buf scx_exit_bstr_buf;
 136
 137/* ops debug dump */
 138struct scx_dump_data {
 139	s32			cpu;
 140	bool			first;
 141	s32			cursor;
 142	struct seq_buf		*s;
 143	const char		*prefix;
 144	struct scx_bstr_buf	buf;
 145};
 146
 147static struct scx_dump_data scx_dump_data = {
 148	.cpu			= -1,
 149};
 150
 151/* /sys/kernel/sched_ext interface */
 152static struct kset *scx_kset;
 153
 154/*
 155 * Parameters that can be adjusted through /sys/module/sched_ext/parameters.
 156 * There usually is no reason to modify these as normal scheduler operation
 157 * shouldn't be affected by them. The knobs are primarily for debugging.
 158 */
 159static u64 scx_slice_dfl = SCX_SLICE_DFL;
 160static unsigned int scx_slice_bypass_us = SCX_SLICE_BYPASS / NSEC_PER_USEC;
 161static unsigned int scx_bypass_lb_intv_us = SCX_BYPASS_LB_DFL_INTV_US;
 162
 163static int set_slice_us(const char *val, const struct kernel_param *kp)
 164{
 165	return param_set_uint_minmax(val, kp, 100, 100 * USEC_PER_MSEC);
 166}
 167
 168static const struct kernel_param_ops slice_us_param_ops = {
 169	.set = set_slice_us,
 170	.get = param_get_uint,
 171};
 172
 173static int set_bypass_lb_intv_us(const char *val, const struct kernel_param *kp)
 174{
 175	return param_set_uint_minmax(val, kp, 0, 10 * USEC_PER_SEC);
 176}
 177
 178static const struct kernel_param_ops bypass_lb_intv_us_param_ops = {
 179	.set = set_bypass_lb_intv_us,
 180	.get = param_get_uint,
 181};
 182
 183#undef MODULE_PARAM_PREFIX
 184#define MODULE_PARAM_PREFIX	"sched_ext."
 185
 186module_param_cb(slice_bypass_us, &slice_us_param_ops, &scx_slice_bypass_us, 0600);
 187MODULE_PARM_DESC(slice_bypass_us, "bypass slice in microseconds, applied on [un]load (100us to 100ms)");
 188module_param_cb(bypass_lb_intv_us, &bypass_lb_intv_us_param_ops, &scx_bypass_lb_intv_us, 0600);
 189MODULE_PARM_DESC(bypass_lb_intv_us, "bypass load balance interval in microseconds (0 (disable) to 10s)");
 190
 191#undef MODULE_PARAM_PREFIX
 192
 193#define CREATE_TRACE_POINTS
 194#include <trace/events/sched_ext.h>
 195
 196static void process_ddsp_deferred_locals(struct rq *rq);
 197static bool task_dead_and_done(struct task_struct *p);
 198static u32 reenq_local(struct rq *rq);
 199static void scx_kick_cpu(struct scx_sched *sch, s32 cpu, u64 flags);
 200static bool scx_vexit(struct scx_sched *sch, enum scx_exit_kind kind,
 201		      s64 exit_code, const char *fmt, va_list args);
 202
 203static __printf(4, 5) bool scx_exit(struct scx_sched *sch,
 204				    enum scx_exit_kind kind, s64 exit_code,
 205				    const char *fmt, ...)
 206{
 207	va_list args;
 208	bool ret;
 209
 210	va_start(args, fmt);
 211	ret = scx_vexit(sch, kind, exit_code, fmt, args);
 212	va_end(args);
 213
 214	return ret;
 215}
 216
 217#define scx_error(sch, fmt, args...)	scx_exit((sch), SCX_EXIT_ERROR, 0, fmt, ##args)
 218#define scx_verror(sch, fmt, args)	scx_vexit((sch), SCX_EXIT_ERROR, 0, fmt, args)
 219
 220#define SCX_HAS_OP(sch, op)	test_bit(SCX_OP_IDX(op), (sch)->has_op)
 221
 222static long jiffies_delta_msecs(unsigned long at, unsigned long now)
 223{
 224	if (time_after(at, now))
 225		return jiffies_to_msecs(at - now);
 226	else
 227		return -(long)jiffies_to_msecs(now - at);
 228}
 229
 230/* if the highest set bit is N, return a mask with bits [N+1, 31] set */
 231static u32 higher_bits(u32 flags)
 232{
 233	return ~((1 << fls(flags)) - 1);
 234}
 235
 236/* return the mask with only the highest bit set */
 237static u32 highest_bit(u32 flags)
 238{
 239	int bit = fls(flags);
 240	return ((u64)1 << bit) >> 1;
 241}
 242
 243static bool u32_before(u32 a, u32 b)
 244{
 245	return (s32)(a - b) < 0;
 246}
 247
 248static struct scx_dispatch_q *find_global_dsq(struct scx_sched *sch,
 249					      struct task_struct *p)
 250{
 251	return sch->global_dsqs[cpu_to_node(task_cpu(p))];
 252}
 253
 254static struct scx_dispatch_q *find_user_dsq(struct scx_sched *sch, u64 dsq_id)
 255{
 256	return rhashtable_lookup(&sch->dsq_hash, &dsq_id, dsq_hash_params);
 257}
 258
 259static const struct sched_class *scx_setscheduler_class(struct task_struct *p)
 260{
 261	if (p->sched_class == &stop_sched_class)
 262		return &stop_sched_class;
 263
 264	return __setscheduler_class(p->policy, p->prio);
 265}
 266
 267/*
 268 * scx_kf_mask enforcement. Some kfuncs can only be called from specific SCX
 269 * ops. When invoking SCX ops, SCX_CALL_OP[_RET]() should be used to indicate
 270 * the allowed kfuncs and those kfuncs should use scx_kf_allowed() to check
 271 * whether it's running from an allowed context.
 272 *
 273 * @mask is constant, always inline to cull the mask calculations.
 274 */
 275static __always_inline void scx_kf_allow(u32 mask)
 276{
 277	/* nesting is allowed only in increasing scx_kf_mask order */
 278	WARN_ONCE((mask | higher_bits(mask)) & current->scx.kf_mask,
 279		  "invalid nesting current->scx.kf_mask=0x%x mask=0x%x\n",
 280		  current->scx.kf_mask, mask);
 281	current->scx.kf_mask |= mask;
 282	barrier();
 283}
 284
 285static void scx_kf_disallow(u32 mask)
 286{
 287	barrier();
 288	current->scx.kf_mask &= ~mask;
 289}
 290
 291/*
 292 * Track the rq currently locked.
 293 *
 294 * This allows kfuncs to safely operate on rq from any scx ops callback,
 295 * knowing which rq is already locked.
 296 */
 297DEFINE_PER_CPU(struct rq *, scx_locked_rq_state);
 298
 299static inline void update_locked_rq(struct rq *rq)
 300{
 301	/*
 302	 * Check whether @rq is actually locked. This can help expose bugs
 303	 * or incorrect assumptions about the context in which a kfunc or
 304	 * callback is executed.
 305	 */
 306	if (rq)
 307		lockdep_assert_rq_held(rq);
 308	__this_cpu_write(scx_locked_rq_state, rq);
 309}
 310
 311#define SCX_CALL_OP(sch, mask, op, rq, args...)					\
 312do {										\
 313	if (rq)									\
 314		update_locked_rq(rq);						\
 315	if (mask) {								\
 316		scx_kf_allow(mask);						\
 317		(sch)->ops.op(args);						\
 318		scx_kf_disallow(mask);						\
 319	} else {								\
 320		(sch)->ops.op(args);						\
 321	}									\
 322	if (rq)									\
 323		update_locked_rq(NULL);						\
 324} while (0)
 325
 326#define SCX_CALL_OP_RET(sch, mask, op, rq, args...)				\
 327({										\
 328	__typeof__((sch)->ops.op(args)) __ret;					\
 329										\
 330	if (rq)									\
 331		update_locked_rq(rq);						\
 332	if (mask) {								\
 333		scx_kf_allow(mask);						\
 334		__ret = (sch)->ops.op(args);					\
 335		scx_kf_disallow(mask);						\
 336	} else {								\
 337		__ret = (sch)->ops.op(args);					\
 338	}									\
 339	if (rq)									\
 340		update_locked_rq(NULL);						\
 341	__ret;									\
 342})
 343
 344/*
 345 * Some kfuncs are allowed only on the tasks that are subjects of the
 346 * in-progress scx_ops operation for, e.g., locking guarantees. To enforce such
 347 * restrictions, the following SCX_CALL_OP_*() variants should be used when
 348 * invoking scx_ops operations that take task arguments. These can only be used
 349 * for non-nesting operations due to the way the tasks are tracked.
 350 *
 351 * kfuncs which can only operate on such tasks can in turn use
 352 * scx_kf_allowed_on_arg_tasks() to test whether the invocation is allowed on
 353 * the specific task.
 354 */
 355#define SCX_CALL_OP_TASK(sch, mask, op, rq, task, args...)			\
 356do {										\
 357	BUILD_BUG_ON((mask) & ~__SCX_KF_TERMINAL);				\
 358	current->scx.kf_tasks[0] = task;					\
 359	SCX_CALL_OP((sch), mask, op, rq, task, ##args);				\
 360	current->scx.kf_tasks[0] = NULL;					\
 361} while (0)
 362
 363#define SCX_CALL_OP_TASK_RET(sch, mask, op, rq, task, args...)			\
 364({										\
 365	__typeof__((sch)->ops.op(task, ##args)) __ret;				\
 366	BUILD_BUG_ON((mask) & ~__SCX_KF_TERMINAL);				\
 367	current->scx.kf_tasks[0] = task;					\
 368	__ret = SCX_CALL_OP_RET((sch), mask, op, rq, task, ##args);		\
 369	current->scx.kf_tasks[0] = NULL;					\
 370	__ret;									\
 371})
 372
 373#define SCX_CALL_OP_2TASKS_RET(sch, mask, op, rq, task0, task1, args...)	\
 374({										\
 375	__typeof__((sch)->ops.op(task0, task1, ##args)) __ret;			\
 376	BUILD_BUG_ON((mask) & ~__SCX_KF_TERMINAL);				\
 377	current->scx.kf_tasks[0] = task0;					\
 378	current->scx.kf_tasks[1] = task1;					\
 379	__ret = SCX_CALL_OP_RET((sch), mask, op, rq, task0, task1, ##args);	\
 380	current->scx.kf_tasks[0] = NULL;					\
 381	current->scx.kf_tasks[1] = NULL;					\
 382	__ret;									\
 383})
 384
 385/* @mask is constant, always inline to cull unnecessary branches */
 386static __always_inline bool scx_kf_allowed(struct scx_sched *sch, u32 mask)
 387{
 388	if (unlikely(!(current->scx.kf_mask & mask))) {
 389		scx_error(sch, "kfunc with mask 0x%x called from an operation only allowing 0x%x",
 390			  mask, current->scx.kf_mask);
 391		return false;
 392	}
 393
 394	/*
 395	 * Enforce nesting boundaries. e.g. A kfunc which can be called from
 396	 * DISPATCH must not be called if we're running DEQUEUE which is nested
 397	 * inside ops.dispatch(). We don't need to check boundaries for any
 398	 * blocking kfuncs as the verifier ensures they're only called from
 399	 * sleepable progs.
 400	 */
 401	if (unlikely(highest_bit(mask) == SCX_KF_CPU_RELEASE &&
 402		     (current->scx.kf_mask & higher_bits(SCX_KF_CPU_RELEASE)))) {
 403		scx_error(sch, "cpu_release kfunc called from a nested operation");
 404		return false;
 405	}
 406
 407	if (unlikely(highest_bit(mask) == SCX_KF_DISPATCH &&
 408		     (current->scx.kf_mask & higher_bits(SCX_KF_DISPATCH)))) {
 409		scx_error(sch, "dispatch kfunc called from a nested operation");
 410		return false;
 411	}
 412
 413	return true;
 414}
 415
 416/* see SCX_CALL_OP_TASK() */
 417static __always_inline bool scx_kf_allowed_on_arg_tasks(struct scx_sched *sch,
 418							u32 mask,
 419							struct task_struct *p)
 420{
 421	if (!scx_kf_allowed(sch, mask))
 422		return false;
 423
 424	if (unlikely((p != current->scx.kf_tasks[0] &&
 425		      p != current->scx.kf_tasks[1]))) {
 426		scx_error(sch, "called on a task not being operated on");
 427		return false;
 428	}
 429
 430	return true;
 431}
 432
 433/**
 434 * nldsq_next_task - Iterate to the next task in a non-local DSQ
 435 * @dsq: user dsq being iterated
 436 * @cur: current position, %NULL to start iteration
 437 * @rev: walk backwards
 438 *
 439 * Returns %NULL when iteration is finished.
 440 */
 441static struct task_struct *nldsq_next_task(struct scx_dispatch_q *dsq,
 442					   struct task_struct *cur, bool rev)
 443{
 444	struct list_head *list_node;
 445	struct scx_dsq_list_node *dsq_lnode;
 446
 447	lockdep_assert_held(&dsq->lock);
 448
 449	if (cur)
 450		list_node = &cur->scx.dsq_list.node;
 451	else
 452		list_node = &dsq->list;
 453
 454	/* find the next task, need to skip BPF iteration cursors */
 455	do {
 456		if (rev)
 457			list_node = list_node->prev;
 458		else
 459			list_node = list_node->next;
 460
 461		if (list_node == &dsq->list)
 462			return NULL;
 463
 464		dsq_lnode = container_of(list_node, struct scx_dsq_list_node,
 465					 node);
 466	} while (dsq_lnode->flags & SCX_DSQ_LNODE_ITER_CURSOR);
 467
 468	return container_of(dsq_lnode, struct task_struct, scx.dsq_list);
 469}
 470
 471#define nldsq_for_each_task(p, dsq)						\
 472	for ((p) = nldsq_next_task((dsq), NULL, false); (p);			\
 473	     (p) = nldsq_next_task((dsq), (p), false))
 474
 475
 476/*
 477 * BPF DSQ iterator. Tasks in a non-local DSQ can be iterated in [reverse]
 478 * dispatch order. BPF-visible iterator is opaque and larger to allow future
 479 * changes without breaking backward compatibility. Can be used with
 480 * bpf_for_each(). See bpf_iter_scx_dsq_*().
 481 */
 482enum scx_dsq_iter_flags {
 483	/* iterate in the reverse dispatch order */
 484	SCX_DSQ_ITER_REV		= 1U << 16,
 485
 486	__SCX_DSQ_ITER_HAS_SLICE	= 1U << 30,
 487	__SCX_DSQ_ITER_HAS_VTIME	= 1U << 31,
 488
 489	__SCX_DSQ_ITER_USER_FLAGS	= SCX_DSQ_ITER_REV,
 490	__SCX_DSQ_ITER_ALL_FLAGS	= __SCX_DSQ_ITER_USER_FLAGS |
 491					  __SCX_DSQ_ITER_HAS_SLICE |
 492					  __SCX_DSQ_ITER_HAS_VTIME,
 493};
 494
 495struct bpf_iter_scx_dsq_kern {
 496	struct scx_dsq_list_node	cursor;
 497	struct scx_dispatch_q		*dsq;
 498	u64				slice;
 499	u64				vtime;
 500} __attribute__((aligned(8)));
 501
 502struct bpf_iter_scx_dsq {
 503	u64				__opaque[6];
 504} __attribute__((aligned(8)));
 505
 506
 507/*
 508 * SCX task iterator.
 509 */
 510struct scx_task_iter {
 511	struct sched_ext_entity		cursor;
 512	struct task_struct		*locked_task;
 513	struct rq			*rq;
 514	struct rq_flags			rf;
 515	u32				cnt;
 516	bool				list_locked;
 517};
 518
 519/**
 520 * scx_task_iter_start - Lock scx_tasks_lock and start a task iteration
 521 * @iter: iterator to init
 522 *
 523 * Initialize @iter and return with scx_tasks_lock held. Once initialized, @iter
 524 * must eventually be stopped with scx_task_iter_stop().
 525 *
 526 * scx_tasks_lock and the rq lock may be released using scx_task_iter_unlock()
 527 * between this and the first next() call or between any two next() calls. If
 528 * the locks are released between two next() calls, the caller is responsible
 529 * for ensuring that the task being iterated remains accessible either through
 530 * RCU read lock or obtaining a reference count.
 531 *
 532 * All tasks which existed when the iteration started are guaranteed to be
 533 * visited as long as they are not dead.
 534 */
 535static void scx_task_iter_start(struct scx_task_iter *iter)
 536{
 537	memset(iter, 0, sizeof(*iter));
 538
 539	raw_spin_lock_irq(&scx_tasks_lock);
 540
 541	iter->cursor = (struct sched_ext_entity){ .flags = SCX_TASK_CURSOR };
 542	list_add(&iter->cursor.tasks_node, &scx_tasks);
 543	iter->list_locked = true;
 544}
 545
 546static void __scx_task_iter_rq_unlock(struct scx_task_iter *iter)
 547{
 548	if (iter->locked_task) {
 549		__balance_callbacks(iter->rq, &iter->rf);
 550		task_rq_unlock(iter->rq, iter->locked_task, &iter->rf);
 551		iter->locked_task = NULL;
 552	}
 553}
 554
 555/**
 556 * scx_task_iter_unlock - Unlock rq and scx_tasks_lock held by a task iterator
 557 * @iter: iterator to unlock
 558 *
 559 * If @iter is in the middle of a locked iteration, it may be locking the rq of
 560 * the task currently being visited in addition to scx_tasks_lock. Unlock both.
 561 * This function can be safely called anytime during an iteration. The next
 562 * iterator operation will automatically restore the necessary locking.
 563 */
 564static void scx_task_iter_unlock(struct scx_task_iter *iter)
 565{
 566	__scx_task_iter_rq_unlock(iter);
 567	if (iter->list_locked) {
 568		iter->list_locked = false;
 569		raw_spin_unlock_irq(&scx_tasks_lock);
 570	}
 571}
 572
 573static void __scx_task_iter_maybe_relock(struct scx_task_iter *iter)
 574{
 575	if (!iter->list_locked) {
 576		raw_spin_lock_irq(&scx_tasks_lock);
 577		iter->list_locked = true;
 578	}
 579}
 580
 581/**
 582 * scx_task_iter_stop - Stop a task iteration and unlock scx_tasks_lock
 583 * @iter: iterator to exit
 584 *
 585 * Exit a previously initialized @iter. Must be called with scx_tasks_lock held
 586 * which is released on return. If the iterator holds a task's rq lock, that rq
 587 * lock is also released. See scx_task_iter_start() for details.
 588 */
 589static void scx_task_iter_stop(struct scx_task_iter *iter)
 590{
 591	__scx_task_iter_maybe_relock(iter);
 592	list_del_init(&iter->cursor.tasks_node);
 593	scx_task_iter_unlock(iter);
 594}
 595
 596/**
 597 * scx_task_iter_next - Next task
 598 * @iter: iterator to walk
 599 *
 600 * Visit the next task. See scx_task_iter_start() for details. Locks are dropped
 601 * and re-acquired every %SCX_TASK_ITER_BATCH iterations to avoid causing stalls
 602 * by holding scx_tasks_lock for too long.
 603 */
 604static struct task_struct *scx_task_iter_next(struct scx_task_iter *iter)
 605{
 606	struct list_head *cursor = &iter->cursor.tasks_node;
 607	struct sched_ext_entity *pos;
 608
 609	if (!(++iter->cnt % SCX_TASK_ITER_BATCH)) {
 610		scx_task_iter_unlock(iter);
 611		cond_resched();
 612	}
 613
 614	__scx_task_iter_maybe_relock(iter);
 615
 616	list_for_each_entry(pos, cursor, tasks_node) {
 617		if (&pos->tasks_node == &scx_tasks)
 618			return NULL;
 619		if (!(pos->flags & SCX_TASK_CURSOR)) {
 620			list_move(cursor, &pos->tasks_node);
 621			return container_of(pos, struct task_struct, scx);
 622		}
 623	}
 624
 625	/* can't happen, should always terminate at scx_tasks above */
 626	BUG();
 627}
 628
 629/**
 630 * scx_task_iter_next_locked - Next non-idle task with its rq locked
 631 * @iter: iterator to walk
 632 *
 633 * Visit the non-idle task with its rq lock held. Allows callers to specify
 634 * whether they would like to filter out dead tasks. See scx_task_iter_start()
 635 * for details.
 636 */
 637static struct task_struct *scx_task_iter_next_locked(struct scx_task_iter *iter)
 638{
 639	struct task_struct *p;
 640
 641	__scx_task_iter_rq_unlock(iter);
 642
 643	while ((p = scx_task_iter_next(iter))) {
 644		/*
 645		 * scx_task_iter is used to prepare and move tasks into SCX
 646		 * while loading the BPF scheduler and vice-versa while
 647		 * unloading. The init_tasks ("swappers") should be excluded
 648		 * from the iteration because:
 649		 *
 650		 * - It's unsafe to use __setschduler_prio() on an init_task to
 651		 *   determine the sched_class to use as it won't preserve its
 652		 *   idle_sched_class.
 653		 *
 654		 * - ops.init/exit_task() can easily be confused if called with
 655		 *   init_tasks as they, e.g., share PID 0.
 656		 *
 657		 * As init_tasks are never scheduled through SCX, they can be
 658		 * skipped safely. Note that is_idle_task() which tests %PF_IDLE
 659		 * doesn't work here:
 660		 *
 661		 * - %PF_IDLE may not be set for an init_task whose CPU hasn't
 662		 *   yet been onlined.
 663		 *
 664		 * - %PF_IDLE can be set on tasks that are not init_tasks. See
 665		 *   play_idle_precise() used by CONFIG_IDLE_INJECT.
 666		 *
 667		 * Test for idle_sched_class as only init_tasks are on it.
 668		 */
 669		if (p->sched_class != &idle_sched_class)
 670			break;
 671	}
 672	if (!p)
 673		return NULL;
 674
 675	iter->rq = task_rq_lock(p, &iter->rf);
 676	iter->locked_task = p;
 677
 678	return p;
 679}
 680
 681/**
 682 * scx_add_event - Increase an event counter for 'name' by 'cnt'
 683 * @sch: scx_sched to account events for
 684 * @name: an event name defined in struct scx_event_stats
 685 * @cnt: the number of the event occurred
 686 *
 687 * This can be used when preemption is not disabled.
 688 */
 689#define scx_add_event(sch, name, cnt) do {					\
 690	this_cpu_add((sch)->pcpu->event_stats.name, (cnt));			\
 691	trace_sched_ext_event(#name, (cnt));					\
 692} while(0)
 693
 694/**
 695 * __scx_add_event - Increase an event counter for 'name' by 'cnt'
 696 * @sch: scx_sched to account events for
 697 * @name: an event name defined in struct scx_event_stats
 698 * @cnt: the number of the event occurred
 699 *
 700 * This should be used only when preemption is disabled.
 701 */
 702#define __scx_add_event(sch, name, cnt) do {					\
 703	__this_cpu_add((sch)->pcpu->event_stats.name, (cnt));			\
 704	trace_sched_ext_event(#name, cnt);					\
 705} while(0)
 706
 707/**
 708 * scx_agg_event - Aggregate an event counter 'kind' from 'src_e' to 'dst_e'
 709 * @dst_e: destination event stats
 710 * @src_e: source event stats
 711 * @kind: a kind of event to be aggregated
 712 */
 713#define scx_agg_event(dst_e, src_e, kind) do {					\
 714	(dst_e)->kind += READ_ONCE((src_e)->kind);				\
 715} while(0)
 716
 717/**
 718 * scx_dump_event - Dump an event 'kind' in 'events' to 's'
 719 * @s: output seq_buf
 720 * @events: event stats
 721 * @kind: a kind of event to dump
 722 */
 723#define scx_dump_event(s, events, kind) do {					\
 724	dump_line(&(s), "%40s: %16lld", #kind, (events)->kind);			\
 725} while (0)
 726
 727
 728static void scx_read_events(struct scx_sched *sch,
 729			    struct scx_event_stats *events);
 730
 731static enum scx_enable_state scx_enable_state(void)
 732{
 733	return atomic_read(&scx_enable_state_var);
 734}
 735
 736static enum scx_enable_state scx_set_enable_state(enum scx_enable_state to)
 737{
 738	return atomic_xchg(&scx_enable_state_var, to);
 739}
 740
 741static bool scx_tryset_enable_state(enum scx_enable_state to,
 742				    enum scx_enable_state from)
 743{
 744	int from_v = from;
 745
 746	return atomic_try_cmpxchg(&scx_enable_state_var, &from_v, to);
 747}
 748
 749/**
 750 * wait_ops_state - Busy-wait the specified ops state to end
 751 * @p: target task
 752 * @opss: state to wait the end of
 753 *
 754 * Busy-wait for @p to transition out of @opss. This can only be used when the
 755 * state part of @opss is %SCX_QUEUEING or %SCX_DISPATCHING. This function also
 756 * has load_acquire semantics to ensure that the caller can see the updates made
 757 * in the enqueueing and dispatching paths.
 758 */
 759static void wait_ops_state(struct task_struct *p, unsigned long opss)
 760{
 761	do {
 762		cpu_relax();
 763	} while (atomic_long_read_acquire(&p->scx.ops_state) == opss);
 764}
 765
 766static inline bool __cpu_valid(s32 cpu)
 767{
 768	return likely(cpu >= 0 && cpu < nr_cpu_ids && cpu_possible(cpu));
 769}
 770
 771/**
 772 * ops_cpu_valid - Verify a cpu number, to be used on ops input args
 773 * @sch: scx_sched to abort on error
 774 * @cpu: cpu number which came from a BPF ops
 775 * @where: extra information reported on error
 776 *
 777 * @cpu is a cpu number which came from the BPF scheduler and can be any value.
 778 * Verify that it is in range and one of the possible cpus. If invalid, trigger
 779 * an ops error.
 780 */
 781static bool ops_cpu_valid(struct scx_sched *sch, s32 cpu, const char *where)
 782{
 783	if (__cpu_valid(cpu)) {
 784		return true;
 785	} else {
 786		scx_error(sch, "invalid CPU %d%s%s", cpu, where ? " " : "", where ?: "");
 787		return false;
 788	}
 789}
 790
 791/**
 792 * ops_sanitize_err - Sanitize a -errno value
 793 * @sch: scx_sched to error out on error
 794 * @ops_name: operation to blame on failure
 795 * @err: -errno value to sanitize
 796 *
 797 * Verify @err is a valid -errno. If not, trigger scx_error() and return
 798 * -%EPROTO. This is necessary because returning a rogue -errno up the chain can
 799 * cause misbehaviors. For an example, a large negative return from
 800 * ops.init_task() triggers an oops when passed up the call chain because the
 801 * value fails IS_ERR() test after being encoded with ERR_PTR() and then is
 802 * handled as a pointer.
 803 */
 804static int ops_sanitize_err(struct scx_sched *sch, const char *ops_name, s32 err)
 805{
 806	if (err < 0 && err >= -MAX_ERRNO)
 807		return err;
 808
 809	scx_error(sch, "ops.%s() returned an invalid errno %d", ops_name, err);
 810	return -EPROTO;
 811}
 812
 813static void run_deferred(struct rq *rq)
 814{
 815	process_ddsp_deferred_locals(rq);
 816
 817	if (local_read(&rq->scx.reenq_local_deferred)) {
 818		local_set(&rq->scx.reenq_local_deferred, 0);
 819		reenq_local(rq);
 820	}
 821}
 822
 823static void deferred_bal_cb_workfn(struct rq *rq)
 824{
 825	run_deferred(rq);
 826}
 827
 828static void deferred_irq_workfn(struct irq_work *irq_work)
 829{
 830	struct rq *rq = container_of(irq_work, struct rq, scx.deferred_irq_work);
 831
 832	raw_spin_rq_lock(rq);
 833	run_deferred(rq);
 834	raw_spin_rq_unlock(rq);
 835}
 836
 837/**
 838 * schedule_deferred - Schedule execution of deferred actions on an rq
 839 * @rq: target rq
 840 *
 841 * Schedule execution of deferred actions on @rq. Deferred actions are executed
 842 * with @rq locked but unpinned, and thus can unlock @rq to e.g. migrate tasks
 843 * to other rqs.
 844 */
 845static void schedule_deferred(struct rq *rq)
 846{
 847	/*
 848	 * Queue an irq work. They are executed on IRQ re-enable which may take
 849	 * a bit longer than the scheduler hook in schedule_deferred_locked().
 850	 */
 851	irq_work_queue(&rq->scx.deferred_irq_work);
 852}
 853
 854/**
 855 * schedule_deferred_locked - Schedule execution of deferred actions on an rq
 856 * @rq: target rq
 857 *
 858 * Schedule execution of deferred actions on @rq. Equivalent to
 859 * schedule_deferred() but requires @rq to be locked and can be more efficient.
 860 */
 861static void schedule_deferred_locked(struct rq *rq)
 862{
 863	lockdep_assert_rq_held(rq);
 864
 865	/*
 866	 * If in the middle of waking up a task, task_woken_scx() will be called
 867	 * afterwards which will then run the deferred actions, no need to
 868	 * schedule anything.
 869	 */
 870	if (rq->scx.flags & SCX_RQ_IN_WAKEUP)
 871		return;
 872
 873	/* Don't do anything if there already is a deferred operation. */
 874	if (rq->scx.flags & SCX_RQ_BAL_CB_PENDING)
 875		return;
 876
 877	/*
 878	 * If in balance, the balance callbacks will be called before rq lock is
 879	 * released. Schedule one.
 880	 *
 881	 *
 882	 * We can't directly insert the callback into the
 883	 * rq's list: The call can drop its lock and make the pending balance
 884	 * callback visible to unrelated code paths that call rq_pin_lock().
 885	 *
 886	 * Just let balance_one() know that it must do it itself.
 887	 */
 888	if (rq->scx.flags & SCX_RQ_IN_BALANCE) {
 889		rq->scx.flags |= SCX_RQ_BAL_CB_PENDING;
 890		return;
 891	}
 892
 893	/*
 894	 * No scheduler hooks available. Use the generic irq_work path. The
 895	 * above WAKEUP and BALANCE paths should cover most of the cases and the
 896	 * time to IRQ re-enable shouldn't be long.
 897	 */
 898	schedule_deferred(rq);
 899}
 900
 901/**
 902 * touch_core_sched - Update timestamp used for core-sched task ordering
 903 * @rq: rq to read clock from, must be locked
 904 * @p: task to update the timestamp for
 905 *
 906 * Update @p->scx.core_sched_at timestamp. This is used by scx_prio_less() to
 907 * implement global or local-DSQ FIFO ordering for core-sched. Should be called
 908 * when a task becomes runnable and its turn on the CPU ends (e.g. slice
 909 * exhaustion).
 910 */
 911static void touch_core_sched(struct rq *rq, struct task_struct *p)
 912{
 913	lockdep_assert_rq_held(rq);
 914
 915#ifdef CONFIG_SCHED_CORE
 916	/*
 917	 * It's okay to update the timestamp spuriously. Use
 918	 * sched_core_disabled() which is cheaper than enabled().
 919	 *
 920	 * As this is used to determine ordering between tasks of sibling CPUs,
 921	 * it may be better to use per-core dispatch sequence instead.
 922	 */
 923	if (!sched_core_disabled())
 924		p->scx.core_sched_at = sched_clock_cpu(cpu_of(rq));
 925#endif
 926}
 927
 928/**
 929 * touch_core_sched_dispatch - Update core-sched timestamp on dispatch
 930 * @rq: rq to read clock from, must be locked
 931 * @p: task being dispatched
 932 *
 933 * If the BPF scheduler implements custom core-sched ordering via
 934 * ops.core_sched_before(), @p->scx.core_sched_at is used to implement FIFO
 935 * ordering within each local DSQ. This function is called from dispatch paths
 936 * and updates @p->scx.core_sched_at if custom core-sched ordering is in effect.
 937 */
 938static void touch_core_sched_dispatch(struct rq *rq, struct task_struct *p)
 939{
 940	lockdep_assert_rq_held(rq);
 941
 942#ifdef CONFIG_SCHED_CORE
 943	if (unlikely(SCX_HAS_OP(scx_root, core_sched_before)))
 944		touch_core_sched(rq, p);
 945#endif
 946}
 947
 948static void update_curr_scx(struct rq *rq)
 949{
 950	struct task_struct *curr = rq->curr;
 951	s64 delta_exec;
 952
 953	delta_exec = update_curr_common(rq);
 954	if (unlikely(delta_exec <= 0))
 955		return;
 956
 957	if (curr->scx.slice != SCX_SLICE_INF) {
 958		curr->scx.slice -= min_t(u64, curr->scx.slice, delta_exec);
 959		if (!curr->scx.slice)
 960			touch_core_sched(rq, curr);
 961	}
 962}
 963
 964static bool scx_dsq_priq_less(struct rb_node *node_a,
 965			      const struct rb_node *node_b)
 966{
 967	const struct task_struct *a =
 968		container_of(node_a, struct task_struct, scx.dsq_priq);
 969	const struct task_struct *b =
 970		container_of(node_b, struct task_struct, scx.dsq_priq);
 971
 972	return time_before64(a->scx.dsq_vtime, b->scx.dsq_vtime);
 973}
 974
 975static void dsq_mod_nr(struct scx_dispatch_q *dsq, s32 delta)
 976{
 977	/* scx_bpf_dsq_nr_queued() reads ->nr without locking, use WRITE_ONCE() */
 978	WRITE_ONCE(dsq->nr, dsq->nr + delta);
 979}
 980
 981static void refill_task_slice_dfl(struct scx_sched *sch, struct task_struct *p)
 982{
 983	p->scx.slice = READ_ONCE(scx_slice_dfl);
 984	__scx_add_event(sch, SCX_EV_REFILL_SLICE_DFL, 1);
 985}
 986
 987static void local_dsq_post_enq(struct scx_dispatch_q *dsq, struct task_struct *p,
 988			       u64 enq_flags)
 989{
 990	struct rq *rq = container_of(dsq, struct rq, scx.local_dsq);
 991	bool preempt = false;
 992
 993	/*
 994	 * If @rq is in balance, the CPU is already vacant and looking for the
 995	 * next task to run. No need to preempt or trigger resched after moving
 996	 * @p into its local DSQ.
 997	 */
 998	if (rq->scx.flags & SCX_RQ_IN_BALANCE)
 999		return;
1000
1001	if ((enq_flags & SCX_ENQ_PREEMPT) && p != rq->curr &&
1002	    rq->curr->sched_class == &ext_sched_class) {
1003		rq->curr->scx.slice = 0;
1004		preempt = true;
1005	}
1006
1007	if (preempt || sched_class_above(&ext_sched_class, rq->curr->sched_class))
1008		resched_curr(rq);
1009}
1010
1011static void dispatch_enqueue(struct scx_sched *sch, struct scx_dispatch_q *dsq,
1012			     struct task_struct *p, u64 enq_flags)
1013{
1014	bool is_local = dsq->id == SCX_DSQ_LOCAL;
1015
1016	WARN_ON_ONCE(p->scx.dsq || !list_empty(&p->scx.dsq_list.node));
1017	WARN_ON_ONCE((p->scx.dsq_flags & SCX_TASK_DSQ_ON_PRIQ) ||
1018		     !RB_EMPTY_NODE(&p->scx.dsq_priq));
1019
1020	if (!is_local) {
1021		raw_spin_lock_nested(&dsq->lock,
1022			(enq_flags & SCX_ENQ_NESTED) ? SINGLE_DEPTH_NESTING : 0);
1023
1024		if (unlikely(dsq->id == SCX_DSQ_INVALID)) {
1025			scx_error(sch, "attempting to dispatch to a destroyed dsq");
1026			/* fall back to the global dsq */
1027			raw_spin_unlock(&dsq->lock);
1028			dsq = find_global_dsq(sch, p);
1029			raw_spin_lock(&dsq->lock);
1030		}
1031	}
1032
1033	if (unlikely((dsq->id & SCX_DSQ_FLAG_BUILTIN) &&
1034		     (enq_flags & SCX_ENQ_DSQ_PRIQ))) {
1035		/*
1036		 * SCX_DSQ_LOCAL and SCX_DSQ_GLOBAL DSQs always consume from
1037		 * their FIFO queues. To avoid confusion and accidentally
1038		 * starving vtime-dispatched tasks by FIFO-dispatched tasks, we
1039		 * disallow any internal DSQ from doing vtime ordering of
1040		 * tasks.
1041		 */
1042		scx_error(sch, "cannot use vtime ordering for built-in DSQs");
1043		enq_flags &= ~SCX_ENQ_DSQ_PRIQ;
1044	}
1045
1046	if (enq_flags & SCX_ENQ_DSQ_PRIQ) {
1047		struct rb_node *rbp;
1048
1049		/*
1050		 * A PRIQ DSQ shouldn't be using FIFO enqueueing. As tasks are
1051		 * linked to both the rbtree and list on PRIQs, this can only be
1052		 * tested easily when adding the first task.
1053		 */
1054		if (unlikely(RB_EMPTY_ROOT(&dsq->priq) &&
1055			     nldsq_next_task(dsq, NULL, false)))
1056			scx_error(sch, "DSQ ID 0x%016llx already had FIFO-enqueued tasks",
1057				  dsq->id);
1058
1059		p->scx.dsq_flags |= SCX_TASK_DSQ_ON_PRIQ;
1060		rb_add(&p->scx.dsq_priq, &dsq->priq, scx_dsq_priq_less);
1061
1062		/*
1063		 * Find the previous task and insert after it on the list so
1064		 * that @dsq->list is vtime ordered.
1065		 */
1066		rbp = rb_prev(&p->scx.dsq_priq);
1067		if (rbp) {
1068			struct task_struct *prev =
1069				container_of(rbp, struct task_struct,
1070					     scx.dsq_priq);
1071			list_add(&p->scx.dsq_list.node, &prev->scx.dsq_list.node);
1072			/* first task unchanged - no update needed */
1073		} else {
1074			list_add(&p->scx.dsq_list.node, &dsq->list);
1075			/* not builtin and new task is at head - use fastpath */
1076			rcu_assign_pointer(dsq->first_task, p);
1077		}
1078	} else {
1079		/* a FIFO DSQ shouldn't be using PRIQ enqueuing */
1080		if (unlikely(!RB_EMPTY_ROOT(&dsq->priq)))
1081			scx_error(sch, "DSQ ID 0x%016llx already had PRIQ-enqueued tasks",
1082				  dsq->id);
1083
1084		if (enq_flags & (SCX_ENQ_HEAD | SCX_ENQ_PREEMPT)) {
1085			list_add(&p->scx.dsq_list.node, &dsq->list);
1086			/* new task inserted at head - use fastpath */
1087			if (!(dsq->id & SCX_DSQ_FLAG_BUILTIN))
1088				rcu_assign_pointer(dsq->first_task, p);
1089		} else {
1090			bool was_empty;
1091
1092			was_empty = list_empty(&dsq->list);
1093			list_add_tail(&p->scx.dsq_list.node, &dsq->list);
1094			if (was_empty && !(dsq->id & SCX_DSQ_FLAG_BUILTIN))
1095				rcu_assign_pointer(dsq->first_task, p);
1096		}
1097	}
1098
1099	/* seq records the order tasks are queued, used by BPF DSQ iterator */
1100	WRITE_ONCE(dsq->seq, dsq->seq + 1);
1101	p->scx.dsq_seq = dsq->seq;
1102
1103	dsq_mod_nr(dsq, 1);
1104	p->scx.dsq = dsq;
1105
1106	/*
1107	 * We're transitioning out of QUEUEING or DISPATCHING. store_release to
1108	 * match waiters' load_acquire.
1109	 */
1110	if (enq_flags & SCX_ENQ_CLEAR_OPSS)
1111		atomic_long_set_release(&p->scx.ops_state, SCX_OPSS_NONE);
1112
1113	if (is_local)
1114		local_dsq_post_enq(dsq, p, enq_flags);
1115	else
1116		raw_spin_unlock(&dsq->lock);
1117}
1118
1119static void task_unlink_from_dsq(struct task_struct *p,
1120				 struct scx_dispatch_q *dsq)
1121{
1122	WARN_ON_ONCE(list_empty(&p->scx.dsq_list.node));
1123
1124	if (p->scx.dsq_flags & SCX_TASK_DSQ_ON_PRIQ) {
1125		rb_erase(&p->scx.dsq_priq, &dsq->priq);
1126		RB_CLEAR_NODE(&p->scx.dsq_priq);
1127		p->scx.dsq_flags &= ~SCX_TASK_DSQ_ON_PRIQ;
1128	}
1129
1130	list_del_init(&p->scx.dsq_list.node);
1131	dsq_mod_nr(dsq, -1);
1132
1133	if (!(dsq->id & SCX_DSQ_FLAG_BUILTIN) && dsq->first_task == p) {
1134		struct task_struct *first_task;
1135
1136		first_task = nldsq_next_task(dsq, NULL, false);
1137		rcu_assign_pointer(dsq->first_task, first_task);
1138	}
1139}
1140
1141static void dispatch_dequeue(struct rq *rq, struct task_struct *p)
1142{
1143	struct scx_dispatch_q *dsq = p->scx.dsq;
1144	bool is_local = dsq == &rq->scx.local_dsq;
1145
1146	lockdep_assert_rq_held(rq);
1147
1148	if (!dsq) {
1149		/*
1150		 * If !dsq && on-list, @p is on @rq's ddsp_deferred_locals.
1151		 * Unlinking is all that's needed to cancel.
1152		 */
1153		if (unlikely(!list_empty(&p->scx.dsq_list.node)))
1154			list_del_init(&p->scx.dsq_list.node);
1155
1156		/*
1157		 * When dispatching directly from the BPF scheduler to a local
1158		 * DSQ, the task isn't associated with any DSQ but
1159		 * @p->scx.holding_cpu may be set under the protection of
1160		 * %SCX_OPSS_DISPATCHING.
1161		 */
1162		if (p->scx.holding_cpu >= 0)
1163			p->scx.holding_cpu = -1;
1164
1165		return;
1166	}
1167
1168	if (!is_local)
1169		raw_spin_lock(&dsq->lock);
1170
1171	/*
1172	 * Now that we hold @dsq->lock, @p->holding_cpu and @p->scx.dsq_* can't
1173	 * change underneath us.
1174	*/
1175	if (p->scx.holding_cpu < 0) {
1176		/* @p must still be on @dsq, dequeue */
1177		task_unlink_from_dsq(p, dsq);
1178	} else {
1179		/*
1180		 * We're racing against dispatch_to_local_dsq() which already
1181		 * removed @p from @dsq and set @p->scx.holding_cpu. Clear the
1182		 * holding_cpu which tells dispatch_to_local_dsq() that it lost
1183		 * the race.
1184		 */
1185		WARN_ON_ONCE(!list_empty(&p->scx.dsq_list.node));
1186		p->scx.holding_cpu = -1;
1187	}
1188	p->scx.dsq = NULL;
1189
1190	if (!is_local)
1191		raw_spin_unlock(&dsq->lock);
1192}
1193
1194/*
1195 * Abbreviated version of dispatch_dequeue() that can be used when both @p's rq
1196 * and dsq are locked.
1197 */
1198static void dispatch_dequeue_locked(struct task_struct *p,
1199				    struct scx_dispatch_q *dsq)
1200{
1201	lockdep_assert_rq_held(task_rq(p));
1202	lockdep_assert_held(&dsq->lock);
1203
1204	task_unlink_from_dsq(p, dsq);
1205	p->scx.dsq = NULL;
1206}
1207
1208static struct scx_dispatch_q *find_dsq_for_dispatch(struct scx_sched *sch,
1209						    struct rq *rq, u64 dsq_id,
1210						    struct task_struct *p)
1211{
1212	struct scx_dispatch_q *dsq;
1213
1214	if (dsq_id == SCX_DSQ_LOCAL)
1215		return &rq->scx.local_dsq;
1216
1217	if ((dsq_id & SCX_DSQ_LOCAL_ON) == SCX_DSQ_LOCAL_ON) {
1218		s32 cpu = dsq_id & SCX_DSQ_LOCAL_CPU_MASK;
1219
1220		if (!ops_cpu_valid(sch, cpu, "in SCX_DSQ_LOCAL_ON dispatch verdict"))
1221			return find_global_dsq(sch, p);
1222
1223		return &cpu_rq(cpu)->scx.local_dsq;
1224	}
1225
1226	if (dsq_id == SCX_DSQ_GLOBAL)
1227		dsq = find_global_dsq(sch, p);
1228	else
1229		dsq = find_user_dsq(sch, dsq_id);
1230
1231	if (unlikely(!dsq)) {
1232		scx_error(sch, "non-existent DSQ 0x%llx for %s[%d]",
1233			  dsq_id, p->comm, p->pid);
1234		return find_global_dsq(sch, p);
1235	}
1236
1237	return dsq;
1238}
1239
1240static void mark_direct_dispatch(struct scx_sched *sch,
1241				 struct task_struct *ddsp_task,
1242				 struct task_struct *p, u64 dsq_id,
1243				 u64 enq_flags)
1244{
1245	/*
1246	 * Mark that dispatch already happened from ops.select_cpu() or
1247	 * ops.enqueue() by spoiling direct_dispatch_task with a non-NULL value
1248	 * which can never match a valid task pointer.
1249	 */
1250	__this_cpu_write(direct_dispatch_task, ERR_PTR(-ESRCH));
1251
1252	/* @p must match the task on the enqueue path */
1253	if (unlikely(p != ddsp_task)) {
1254		if (IS_ERR(ddsp_task))
1255			scx_error(sch, "%s[%d] already direct-dispatched",
1256				  p->comm, p->pid);
1257		else
1258			scx_error(sch, "scheduling for %s[%d] but trying to direct-dispatch %s[%d]",
1259				  ddsp_task->comm, ddsp_task->pid,
1260				  p->comm, p->pid);
1261		return;
1262	}
1263
1264	WARN_ON_ONCE(p->scx.ddsp_dsq_id != SCX_DSQ_INVALID);
1265	WARN_ON_ONCE(p->scx.ddsp_enq_flags);
1266
1267	p->scx.ddsp_dsq_id = dsq_id;
1268	p->scx.ddsp_enq_flags = enq_flags;
1269}
1270
1271/*
1272 * Clear @p direct dispatch state when leaving the scheduler.
1273 *
1274 * Direct dispatch state must be cleared in the following cases:
1275 *  - direct_dispatch(): cleared on the synchronous enqueue path, deferred
1276 *    dispatch keeps the state until consumed
1277 *  - process_ddsp_deferred_locals(): cleared after consuming deferred state,
1278 *  - do_enqueue_task(): cleared on enqueue fallbacks where the dispatch
1279 *    verdict is ignored (local/global/bypass)
1280 *  - dequeue_task_scx(): cleared after dispatch_dequeue(), covering deferred
1281 *    cancellation and holding_cpu races
1282 *  - scx_disable_task(): cleared for queued wakeup tasks, which are excluded by
1283 *    the scx_bypass() loop, so that stale state is not reused by a subsequent
1284 *    scheduler instance
1285 */
1286static inline void clear_direct_dispatch(struct task_struct *p)
1287{
1288	p->scx.ddsp_dsq_id = SCX_DSQ_INVALID;
1289	p->scx.ddsp_enq_flags = 0;
1290}
1291
1292static void direct_dispatch(struct scx_sched *sch, struct task_struct *p,
1293			    u64 enq_flags)
1294{
1295	struct rq *rq = task_rq(p);
1296	struct scx_dispatch_q *dsq =
1297		find_dsq_for_dispatch(sch, rq, p->scx.ddsp_dsq_id, p);
1298	u64 ddsp_enq_flags;
1299
1300	touch_core_sched_dispatch(rq, p);
1301
1302	p->scx.ddsp_enq_flags |= enq_flags;
1303
1304	/*
1305	 * We are in the enqueue path with @rq locked and pinned, and thus can't
1306	 * double lock a remote rq and enqueue to its local DSQ. For
1307	 * DSQ_LOCAL_ON verdicts targeting the local DSQ of a remote CPU, defer
1308	 * the enqueue so that it's executed when @rq can be unlocked.
1309	 */
1310	if (dsq->id == SCX_DSQ_LOCAL && dsq != &rq->scx.local_dsq) {
1311		unsigned long opss;
1312
1313		opss = atomic_long_read(&p->scx.ops_state) & SCX_OPSS_STATE_MASK;
1314
1315		switch (opss & SCX_OPSS_STATE_MASK) {
1316		case SCX_OPSS_NONE:
1317			break;
1318		case SCX_OPSS_QUEUEING:
1319			/*
1320			 * As @p was never passed to the BPF side, _release is
1321			 * not strictly necessary. Still do it for consistency.
1322			 */
1323			atomic_long_set_release(&p->scx.ops_state, SCX_OPSS_NONE);
1324			break;
1325		default:
1326			WARN_ONCE(true, "sched_ext: %s[%d] has invalid ops state 0x%lx in direct_dispatch()",
1327				  p->comm, p->pid, opss);
1328			atomic_long_set_release(&p->scx.ops_state, SCX_OPSS_NONE);
1329			break;
1330		}
1331
1332		WARN_ON_ONCE(p->scx.dsq || !list_empty(&p->scx.dsq_list.node));
1333		list_add_tail(&p->scx.dsq_list.node,
1334			      &rq->scx.ddsp_deferred_locals);
1335		schedule_deferred_locked(rq);
1336		return;
1337	}
1338
1339	ddsp_enq_flags = p->scx.ddsp_enq_flags;
1340	clear_direct_dispatch(p);
1341
1342	dispatch_enqueue(sch, dsq, p, ddsp_enq_flags | SCX_ENQ_CLEAR_OPSS);
1343}
1344
1345static bool scx_rq_online(struct rq *rq)
1346{
1347	/*
1348	 * Test both cpu_active() and %SCX_RQ_ONLINE. %SCX_RQ_ONLINE indicates
1349	 * the online state as seen from the BPF scheduler. cpu_active() test
1350	 * guarantees that, if this function returns %true, %SCX_RQ_ONLINE will
1351	 * stay set until the current scheduling operation is complete even if
1352	 * we aren't locking @rq.
1353	 */
1354	return likely((rq->scx.flags & SCX_RQ_ONLINE) && cpu_active(cpu_of(rq)));
1355}
1356
1357static void do_enqueue_task(struct rq *rq, struct task_struct *p, u64 enq_flags,
1358			    int sticky_cpu)
1359{
1360	struct scx_sched *sch = scx_root;
1361	struct task_struct **ddsp_taskp;
1362	struct scx_dispatch_q *dsq;
1363	unsigned long qseq;
1364
1365	WARN_ON_ONCE(!(p->scx.flags & SCX_TASK_QUEUED));
1366
1367	/* rq migration */
1368	if (sticky_cpu == cpu_of(rq))
1369		goto local_norefill;
1370
1371	/*
1372	 * If !scx_rq_online(), we already told the BPF scheduler that the CPU
1373	 * is offline and are just running the hotplug path. Don't bother the
1374	 * BPF scheduler.
1375	 */
1376	if (!scx_rq_online(rq))
1377		goto local;
1378
1379	if (scx_rq_bypassing(rq)) {
1380		__scx_add_event(sch, SCX_EV_BYPASS_DISPATCH, 1);
1381		goto bypass;
1382	}
1383
1384	if (p->scx.ddsp_dsq_id != SCX_DSQ_INVALID)
1385		goto direct;
1386
1387	/* see %SCX_OPS_ENQ_EXITING */
1388	if (!(sch->ops.flags & SCX_OPS_ENQ_EXITING) &&
1389	    unlikely(p->flags & PF_EXITING)) {
1390		__scx_add_event(sch, SCX_EV_ENQ_SKIP_EXITING, 1);
1391		goto local;
1392	}
1393
1394	/* see %SCX_OPS_ENQ_MIGRATION_DISABLED */
1395	if (!(sch->ops.flags & SCX_OPS_ENQ_MIGRATION_DISABLED) &&
1396	    is_migration_disabled(p)) {
1397		__scx_add_event(sch, SCX_EV_ENQ_SKIP_MIGRATION_DISABLED, 1);
1398		goto local;
1399	}
1400
1401	if (unlikely(!SCX_HAS_OP(sch, enqueue)))
1402		goto global;
1403
1404	/* DSQ bypass didn't trigger, enqueue on the BPF scheduler */
1405	qseq = rq->scx.ops_qseq++ << SCX_OPSS_QSEQ_SHIFT;
1406
1407	WARN_ON_ONCE(atomic_long_read(&p->scx.ops_state) != SCX_OPSS_NONE);
1408	atomic_long_set(&p->scx.ops_state, SCX_OPSS_QUEUEING | qseq);
1409
1410	ddsp_taskp = this_cpu_ptr(&direct_dispatch_task);
1411	WARN_ON_ONCE(*ddsp_taskp);
1412	*ddsp_taskp = p;
1413
1414	SCX_CALL_OP_TASK(sch, SCX_KF_ENQUEUE, enqueue, rq, p, enq_flags);
1415
1416	*ddsp_taskp = NULL;
1417	if (p->scx.ddsp_dsq_id != SCX_DSQ_INVALID)
1418		goto direct;
1419
1420	/*
1421	 * If not directly dispatched, QUEUEING isn't clear yet and dispatch or
1422	 * dequeue may be waiting. The store_release matches their load_acquire.
1423	 */
1424	atomic_long_set_release(&p->scx.ops_state, SCX_OPSS_QUEUED | qseq);
1425	return;
1426
1427direct:
1428	direct_dispatch(sch, p, enq_flags);
1429	return;
1430local_norefill:
1431	dispatch_enqueue(sch, &rq->scx.local_dsq, p, enq_flags);
1432	return;
1433local:
1434	dsq = &rq->scx.local_dsq;
1435	goto enqueue;
1436global:
1437	dsq = find_global_dsq(sch, p);
1438	goto enqueue;
1439bypass:
1440	dsq = &task_rq(p)->scx.bypass_dsq;
1441	goto enqueue;
1442
1443enqueue:
1444	/*
1445	 * For task-ordering, slice refill must be treated as implying the end
1446	 * of the current slice. Otherwise, the longer @p stays on the CPU, the
1447	 * higher priority it becomes from scx_prio_less()'s POV.
1448	 */
1449	touch_core_sched(rq, p);
1450	refill_task_slice_dfl(sch, p);
1451	clear_direct_dispatch(p);
1452	dispatch_enqueue(sch, dsq, p, enq_flags);
1453}
1454
1455static bool task_runnable(const struct task_struct *p)
1456{
1457	return !list_empty(&p->scx.runnable_node);
1458}
1459
1460static void set_task_runnable(struct rq *rq, struct task_struct *p)
1461{
1462	lockdep_assert_rq_held(rq);
1463
1464	if (p->scx.flags & SCX_TASK_RESET_RUNNABLE_AT) {
1465		p->scx.runnable_at = jiffies;
1466		p->scx.flags &= ~SCX_TASK_RESET_RUNNABLE_AT;
1467	}
1468
1469	/*
1470	 * list_add_tail() must be used. scx_bypass() depends on tasks being
1471	 * appended to the runnable_list.
1472	 */
1473	list_add_tail(&p->scx.runnable_node, &rq->scx.runnable_list);
1474}
1475
1476static void clr_task_runnable(struct task_struct *p, bool reset_runnable_at)
1477{
1478	list_del_init(&p->scx.runnable_node);
1479	if (reset_runnable_at)
1480		p->scx.flags |= SCX_TASK_RESET_RUNNABLE_AT;
1481}
1482
1483static void enqueue_task_scx(struct rq *rq, struct task_struct *p, int core_enq_flags)
1484{
1485	struct scx_sched *sch = scx_root;
1486	int sticky_cpu = p->scx.sticky_cpu;
1487	u64 enq_flags = core_enq_flags | rq->scx.extra_enq_flags;
1488
1489	if (enq_flags & ENQUEUE_WAKEUP)
1490		rq->scx.flags |= SCX_RQ_IN_WAKEUP;
1491
1492	if (sticky_cpu >= 0)
1493		p->scx.sticky_cpu = -1;
1494
1495	/*
1496	 * Restoring a running task will be immediately followed by
1497	 * set_next_task_scx() which expects the task to not be on the BPF
1498	 * scheduler as tasks can only start running through local DSQs. Force
1499	 * direct-dispatch into the local DSQ by setting the sticky_cpu.
1500	 */
1501	if (unlikely(enq_flags & ENQUEUE_RESTORE) && task_current(rq, p))
1502		sticky_cpu = cpu_of(rq);
1503
1504	if (p->scx.flags & SCX_TASK_QUEUED) {
1505		WARN_ON_ONCE(!task_runnable(p));
1506		goto out;
1507	}
1508
1509	set_task_runnable(rq, p);
1510	p->scx.flags |= SCX_TASK_QUEUED;
1511	rq->scx.nr_running++;
1512	add_nr_running(rq, 1);
1513
1514	if (SCX_HAS_OP(sch, runnable) && !task_on_rq_migrating(p))
1515		SCX_CALL_OP_TASK(sch, SCX_KF_REST, runnable, rq, p, enq_flags);
1516
1517	if (enq_flags & SCX_ENQ_WAKEUP)
1518		touch_core_sched(rq, p);
1519
1520	do_enqueue_task(rq, p, enq_flags, sticky_cpu);
1521out:
1522	rq->scx.flags &= ~SCX_RQ_IN_WAKEUP;
1523
1524	if ((enq_flags & SCX_ENQ_CPU_SELECTED) &&
1525	    unlikely(cpu_of(rq) != p->scx.selected_cpu))
1526		__scx_add_event(sch, SCX_EV_SELECT_CPU_FALLBACK, 1);
1527}
1528
1529static void ops_dequeue(struct rq *rq, struct task_struct *p, u64 deq_flags)
1530{
1531	struct scx_sched *sch = scx_root;
1532	unsigned long opss;
1533
1534	/* dequeue is always temporary, don't reset runnable_at */
1535	clr_task_runnable(p, false);
1536
1537	/* acquire ensures that we see the preceding updates on QUEUED */
1538	opss = atomic_long_read_acquire(&p->scx.ops_state);
1539
1540	switch (opss & SCX_OPSS_STATE_MASK) {
1541	case SCX_OPSS_NONE:
1542		break;
1543	case SCX_OPSS_QUEUEING:
1544		/*
1545		 * QUEUEING is started and finished while holding @p's rq lock.
1546		 * As we're holding the rq lock now, we shouldn't see QUEUEING.
1547		 */
1548		BUG();
1549	case SCX_OPSS_QUEUED:
1550		if (SCX_HAS_OP(sch, dequeue))
1551			SCX_CALL_OP_TASK(sch, SCX_KF_REST, dequeue, rq,
1552					 p, deq_flags);
1553
1554		if (atomic_long_try_cmpxchg(&p->scx.ops_state, &opss,
1555					    SCX_OPSS_NONE))
1556			break;
1557		fallthrough;
1558	case SCX_OPSS_DISPATCHING:
1559		/*
1560		 * If @p is being dispatched from the BPF scheduler to a DSQ,
1561		 * wait for the transfer to complete so that @p doesn't get
1562		 * added to its DSQ after dequeueing is complete.
1563		 *
1564		 * As we're waiting on DISPATCHING with the rq locked, the
1565		 * dispatching side shouldn't try to lock the rq while
1566		 * DISPATCHING is set. See dispatch_to_local_dsq().
1567		 *
1568		 * DISPATCHING shouldn't have qseq set and control can reach
1569		 * here with NONE @opss from the above QUEUED case block.
1570		 * Explicitly wait on %SCX_OPSS_DISPATCHING instead of @opss.
1571		 */
1572		wait_ops_state(p, SCX_OPSS_DISPATCHING);
1573		BUG_ON(atomic_long_read(&p->scx.ops_state) != SCX_OPSS_NONE);
1574		break;
1575	}
1576}
1577
1578static bool dequeue_task_scx(struct rq *rq, struct task_struct *p, int deq_flags)
1579{
1580	struct scx_sched *sch = scx_root;
1581
1582	if (!(p->scx.flags & SCX_TASK_QUEUED)) {
1583		WARN_ON_ONCE(task_runnable(p));
1584		return true;
1585	}
1586
1587	ops_dequeue(rq, p, deq_flags);
1588
1589	/*
1590	 * A currently running task which is going off @rq first gets dequeued
1591	 * and then stops running. As we want running <-> stopping transitions
1592	 * to be contained within runnable <-> quiescent transitions, trigger
1593	 * ->stopping() early here instead of in put_prev_task_scx().
1594	 *
1595	 * @p may go through multiple stopping <-> running transitions between
1596	 * here and put_prev_task_scx() if task attribute changes occur while
1597	 * balance_one() leaves @rq unlocked. However, they don't contain any
1598	 * information meaningful to the BPF scheduler and can be suppressed by
1599	 * skipping the callbacks if the task is !QUEUED.
1600	 */
1601	if (SCX_HAS_OP(sch, stopping) && task_current(rq, p)) {
1602		update_curr_scx(rq);
1603		SCX_CALL_OP_TASK(sch, SCX_KF_REST, stopping, rq, p, false);
1604	}
1605
1606	if (SCX_HAS_OP(sch, quiescent) && !task_on_rq_migrating(p))
1607		SCX_CALL_OP_TASK(sch, SCX_KF_REST, quiescent, rq, p, deq_flags);
1608
1609	if (deq_flags & SCX_DEQ_SLEEP)
1610		p->scx.flags |= SCX_TASK_DEQD_FOR_SLEEP;
1611	else
1612		p->scx.flags &= ~SCX_TASK_DEQD_FOR_SLEEP;
1613
1614	p->scx.flags &= ~SCX_TASK_QUEUED;
1615	rq->scx.nr_running--;
1616	sub_nr_running(rq, 1);
1617
1618	dispatch_dequeue(rq, p);
1619	clear_direct_dispatch(p);
1620	return true;
1621}
1622
1623static void yield_task_scx(struct rq *rq)
1624{
1625	struct scx_sched *sch = scx_root;
1626	struct task_struct *p = rq->donor;
1627
1628	if (SCX_HAS_OP(sch, yield))
1629		SCX_CALL_OP_2TASKS_RET(sch, SCX_KF_REST, yield, rq, p, NULL);
1630	else
1631		p->scx.slice = 0;
1632}
1633
1634static bool yield_to_task_scx(struct rq *rq, struct task_struct *to)
1635{
1636	struct scx_sched *sch = scx_root;
1637	struct task_struct *from = rq->donor;
1638
1639	if (SCX_HAS_OP(sch, yield))
1640		return SCX_CALL_OP_2TASKS_RET(sch, SCX_KF_REST, yield, rq,
1641					      from, to);
1642	else
1643		return false;
1644}
1645
1646static void move_local_task_to_local_dsq(struct task_struct *p, u64 enq_flags,
1647					 struct scx_dispatch_q *src_dsq,
1648					 struct rq *dst_rq)
1649{
1650	struct scx_dispatch_q *dst_dsq = &dst_rq->scx.local_dsq;
1651
1652	/* @dsq is locked and @p is on @dst_rq */
1653	lockdep_assert_held(&src_dsq->lock);
1654	lockdep_assert_rq_held(dst_rq);
1655
1656	WARN_ON_ONCE(p->scx.holding_cpu >= 0);
1657
1658	if (enq_flags & (SCX_ENQ_HEAD | SCX_ENQ_PREEMPT))
1659		list_add(&p->scx.dsq_list.node, &dst_dsq->list);
1660	else
1661		list_add_tail(&p->scx.dsq_list.node, &dst_dsq->list);
1662
1663	dsq_mod_nr(dst_dsq, 1);
1664	p->scx.dsq = dst_dsq;
1665
1666	local_dsq_post_enq(dst_dsq, p, enq_flags);
1667}
1668
1669/**
1670 * move_remote_task_to_local_dsq - Move a task from a foreign rq to a local DSQ
1671 * @p: task to move
1672 * @enq_flags: %SCX_ENQ_*
1673 * @src_rq: rq to move the task from, locked on entry, released on return
1674 * @dst_rq: rq to move the task into, locked on return
1675 *
1676 * Move @p which is currently on @src_rq to @dst_rq's local DSQ.
1677 */
1678static void move_remote_task_to_local_dsq(struct task_struct *p, u64 enq_flags,
1679					  struct rq *src_rq, struct rq *dst_rq)
1680{
1681	lockdep_assert_rq_held(src_rq);
1682
1683	/* the following marks @p MIGRATING which excludes dequeue */
1684	deactivate_task(src_rq, p, 0);
1685	set_task_cpu(p, cpu_of(dst_rq));
1686	p->scx.sticky_cpu = cpu_of(dst_rq);
1687
1688	raw_spin_rq_unlock(src_rq);
1689	raw_spin_rq_lock(dst_rq);
1690
1691	/*
1692	 * We want to pass scx-specific enq_flags but activate_task() will
1693	 * truncate the upper 32 bit. As we own @rq, we can pass them through
1694	 * @rq->scx.extra_enq_flags instead.
1695	 */
1696	WARN_ON_ONCE(!cpumask_test_cpu(cpu_of(dst_rq), p->cpus_ptr));
1697	WARN_ON_ONCE(dst_rq->scx.extra_enq_flags);
1698	dst_rq->scx.extra_enq_flags = enq_flags;
1699	activate_task(dst_rq, p, 0);
1700	dst_rq->scx.extra_enq_flags = 0;
1701}
1702
1703/*
1704 * Similar to kernel/sched/core.c::is_cpu_allowed(). However, there are two
1705 * differences:
1706 *
1707 * - is_cpu_allowed() asks "Can this task run on this CPU?" while
1708 *   task_can_run_on_remote_rq() asks "Can the BPF scheduler migrate the task to
1709 *   this CPU?".
1710 *
1711 *   While migration is disabled, is_cpu_allowed() has to say "yes" as the task
1712 *   must be allowed to finish on the CPU that it's currently on regardless of
1713 *   the CPU state. However, task_can_run_on_remote_rq() must say "no" as the
1714 *   BPF scheduler shouldn't attempt to migrate a task which has migration
1715 *   disabled.
1716 *
1717 * - The BPF scheduler is bypassed while the rq is offline and we can always say
1718 *   no to the BPF scheduler initiated migrations while offline.
1719 *
1720 * The caller must ensure that @p and @rq are on different CPUs.
1721 */
1722static bool task_can_run_on_remote_rq(struct scx_sched *sch,
1723				      struct task_struct *p, struct rq *rq,
1724				      bool enforce)
1725{
1726	int cpu = cpu_of(rq);
1727
1728	WARN_ON_ONCE(task_cpu(p) == cpu);
1729
1730	/*
1731	 * If @p has migration disabled, @p->cpus_ptr is updated to contain only
1732	 * the pinned CPU in migrate_disable_switch() while @p is being switched
1733	 * out. However, put_prev_task_scx() is called before @p->cpus_ptr is
1734	 * updated and thus another CPU may see @p on a DSQ inbetween leading to
1735	 * @p passing the below task_allowed_on_cpu() check while migration is
1736	 * disabled.
1737	 *
1738	 * Test the migration disabled state first as the race window is narrow
1739	 * and the BPF scheduler failing to check migration disabled state can
1740	 * easily be masked if task_allowed_on_cpu() is done first.
1741	 */
1742	if (unlikely(is_migration_disabled(p))) {
1743		if (enforce)
1744			scx_error(sch, "SCX_DSQ_LOCAL[_ON] cannot move migration disabled %s[%d] from CPU %d to %d",
1745				  p->comm, p->pid, task_cpu(p), cpu);
1746		return false;
1747	}
1748
1749	/*
1750	 * We don't require the BPF scheduler to avoid dispatching to offline
1751	 * CPUs mostly for convenience but also because CPUs can go offline
1752	 * between scx_bpf_dsq_insert() calls and here. Trigger error iff the
1753	 * picked CPU is outside the allowed mask.
1754	 */
1755	if (!task_allowed_on_cpu(p, cpu)) {
1756		if (enforce)
1757			scx_error(sch, "SCX_DSQ_LOCAL[_ON] target CPU %d not allowed for %s[%d]",
1758				  cpu, p->comm, p->pid);
1759		return false;
1760	}
1761
1762	if (!scx_rq_online(rq)) {
1763		if (enforce)
1764			__scx_add_event(sch, SCX_EV_DISPATCH_LOCAL_DSQ_OFFLINE, 1);
1765		return false;
1766	}
1767
1768	return true;
1769}
1770
1771/**
1772 * unlink_dsq_and_lock_src_rq() - Unlink task from its DSQ and lock its task_rq
1773 * @p: target task
1774 * @dsq: locked DSQ @p is currently on
1775 * @src_rq: rq @p is currently on, stable with @dsq locked
1776 *
1777 * Called with @dsq locked but no rq's locked. We want to move @p to a different
1778 * DSQ, including any local DSQ, but are not locking @src_rq. Locking @src_rq is
1779 * required when transferring into a local DSQ. Even when transferring into a
1780 * non-local DSQ, it's better to use the same mechanism to protect against
1781 * dequeues and maintain the invariant that @p->scx.dsq can only change while
1782 * @src_rq is locked, which e.g. scx_dump_task() depends on.
1783 *
1784 * We want to grab @src_rq but that can deadlock if we try while locking @dsq,
1785 * so we want to unlink @p from @dsq, drop its lock and then lock @src_rq. As
1786 * this may race with dequeue, which can't drop the rq lock or fail, do a little
1787 * dancing from our side.
1788 *
1789 * @p->scx.holding_cpu is set to this CPU before @dsq is unlocked. If @p gets
1790 * dequeued after we unlock @dsq but before locking @src_rq, the holding_cpu
1791 * would be cleared to -1. While other cpus may have updated it to different
1792 * values afterwards, as this operation can't be preempted or recurse, the
1793 * holding_cpu can never become this CPU again before we're done. Thus, we can
1794 * tell whether we lost to dequeue by testing whether the holding_cpu still
1795 * points to this CPU. See dispatch_dequeue() for the counterpart.
1796 *
1797 * On return, @dsq is unlocked and @src_rq is locked. Returns %true if @p is
1798 * still valid. %false if lost to dequeue.
1799 */
1800static bool unlink_dsq_and_lock_src_rq(struct task_struct *p,
1801				       struct scx_dispatch_q *dsq,
1802				       struct rq *src_rq)
1803{
1804	s32 cpu = raw_smp_processor_id();
1805
1806	lockdep_assert_held(&dsq->lock);
1807
1808	WARN_ON_ONCE(p->scx.holding_cpu >= 0);
1809	task_unlink_from_dsq(p, dsq);
1810	p->scx.holding_cpu = cpu;
1811
1812	raw_spin_unlock(&dsq->lock);
1813	raw_spin_rq_lock(src_rq);
1814
1815	/* task_rq couldn't have changed if we're still the holding cpu */
1816	return likely(p->scx.holding_cpu == cpu) &&
1817		!WARN_ON_ONCE(src_rq != task_rq(p));
1818}
1819
1820static bool consume_remote_task(struct rq *this_rq, struct task_struct *p,
1821				struct scx_dispatch_q *dsq, struct rq *src_rq)
1822{
1823	raw_spin_rq_unlock(this_rq);
1824
1825	if (unlink_dsq_and_lock_src_rq(p, dsq, src_rq)) {
1826		move_remote_task_to_local_dsq(p, 0, src_rq, this_rq);
1827		return true;
1828	} else {
1829		raw_spin_rq_unlock(src_rq);
1830		raw_spin_rq_lock(this_rq);
1831		return false;
1832	}
1833}
1834
1835/**
1836 * move_task_between_dsqs() - Move a task from one DSQ to another
1837 * @sch: scx_sched being operated on
1838 * @p: target task
1839 * @enq_flags: %SCX_ENQ_*
1840 * @src_dsq: DSQ @p is currently on, must not be a local DSQ
1841 * @dst_dsq: DSQ @p is being moved to, can be any DSQ
1842 *
1843 * Must be called with @p's task_rq and @src_dsq locked. If @dst_dsq is a local
1844 * DSQ and @p is on a different CPU, @p will be migrated and thus its task_rq
1845 * will change. As @p's task_rq is locked, this function doesn't need to use the
1846 * holding_cpu mechanism.
1847 *
1848 * On return, @src_dsq is unlocked and only @p's new task_rq, which is the
1849 * return value, is locked.
1850 */
1851static struct rq *move_task_between_dsqs(struct scx_sched *sch,
1852					 struct task_struct *p, u64 enq_flags,
1853					 struct scx_dispatch_q *src_dsq,
1854					 struct scx_dispatch_q *dst_dsq)
1855{
1856	struct rq *src_rq = task_rq(p), *dst_rq;
1857
1858	BUG_ON(src_dsq->id == SCX_DSQ_LOCAL);
1859	lockdep_assert_held(&src_dsq->lock);
1860	lockdep_assert_rq_held(src_rq);
1861
1862	if (dst_dsq->id == SCX_DSQ_LOCAL) {
1863		dst_rq = container_of(dst_dsq, struct rq, scx.local_dsq);
1864		if (src_rq != dst_rq &&
1865		    unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq, true))) {
1866			dst_dsq = find_global_dsq(sch, p);
1867			dst_rq = src_rq;
1868		}
1869	} else {
1870		/* no need to migrate if destination is a non-local DSQ */
1871		dst_rq = src_rq;
1872	}
1873
1874	/*
1875	 * Move @p into $dst_dsq. If $dst_dsq is the local DSQ of a different
1876	 * CPU, @p will be migrated.
1877	 */
1878	if (dst_dsq->id == SCX_DSQ_LOCAL) {
1879		/* @p is going from a non-local DSQ to a local DSQ */
1880		if (src_rq == dst_rq) {
1881			task_unlink_from_dsq(p, src_dsq);
1882			move_local_task_to_local_dsq(p, enq_flags,
1883						     src_dsq, dst_rq);
1884			raw_spin_unlock(&src_dsq->lock);
1885		} else {
1886			raw_spin_unlock(&src_dsq->lock);
1887			move_remote_task_to_local_dsq(p, enq_flags,
1888						      src_rq, dst_rq);
1889		}
1890	} else {
1891		/*
1892		 * @p is going from a non-local DSQ to a non-local DSQ. As
1893		 * $src_dsq is already locked, do an abbreviated dequeue.
1894		 */
1895		dispatch_dequeue_locked(p, src_dsq);
1896		raw_spin_unlock(&src_dsq->lock);
1897
1898		dispatch_enqueue(sch, dst_dsq, p, enq_flags);
1899	}
1900
1901	return dst_rq;
1902}
1903
1904static bool consume_dispatch_q(struct scx_sched *sch, struct rq *rq,
1905			       struct scx_dispatch_q *dsq)
1906{
1907	struct task_struct *p;
1908retry:
1909	/*
1910	 * The caller can't expect to successfully consume a task if the task's
1911	 * addition to @dsq isn't guaranteed to be visible somehow. Test
1912	 * @dsq->list without locking and skip if it seems empty.
1913	 */
1914	if (list_empty(&dsq->list))
1915		return false;
1916
1917	raw_spin_lock(&dsq->lock);
1918
1919	nldsq_for_each_task(p, dsq) {
1920		struct rq *task_rq = task_rq(p);
1921
1922		/*
1923		 * This loop can lead to multiple lockup scenarios, e.g. the BPF
1924		 * scheduler can put an enormous number of affinitized tasks into
1925		 * a contended DSQ, or the outer retry loop can repeatedly race
1926		 * against scx_bypass() dequeueing tasks from @dsq trying to put
1927		 * the system into the bypass mode. This can easily live-lock the
1928		 * machine. If aborting, exit from all non-bypass DSQs.
1929		 */
1930		if (unlikely(READ_ONCE(scx_aborting)) && dsq->id != SCX_DSQ_BYPASS)
1931			break;
1932
1933		if (rq == task_rq) {
1934			task_unlink_from_dsq(p, dsq);
1935			move_local_task_to_local_dsq(p, 0, dsq, rq);
1936			raw_spin_unlock(&dsq->lock);
1937			return true;
1938		}
1939
1940		if (task_can_run_on_remote_rq(sch, p, rq, false)) {
1941			if (likely(consume_remote_task(rq, p, dsq, task_rq)))
1942				return true;
1943			goto retry;
1944		}
1945	}
1946
1947	raw_spin_unlock(&dsq->lock);
1948	return false;
1949}
1950
1951static bool consume_global_dsq(struct scx_sched *sch, struct rq *rq)
1952{
1953	int node = cpu_to_node(cpu_of(rq));
1954
1955	return consume_dispatch_q(sch, rq, sch->global_dsqs[node]);
1956}
1957
1958/**
1959 * dispatch_to_local_dsq - Dispatch a task to a local dsq
1960 * @sch: scx_sched being operated on
1961 * @rq: current rq which is locked
1962 * @dst_dsq: destination DSQ
1963 * @p: task to dispatch
1964 * @enq_flags: %SCX_ENQ_*
1965 *
1966 * We're holding @rq lock and want to dispatch @p to @dst_dsq which is a local
1967 * DSQ. This function performs all the synchronization dancing needed because
1968 * local DSQs are protected with rq locks.
1969 *
1970 * The caller must have exclusive ownership of @p (e.g. through
1971 * %SCX_OPSS_DISPATCHING).
1972 */
1973static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
1974				  struct scx_dispatch_q *dst_dsq,
1975				  struct task_struct *p, u64 enq_flags)
1976{
1977	struct rq *src_rq = task_rq(p);
1978	struct rq *dst_rq = container_of(dst_dsq, struct rq, scx.local_dsq);
1979	struct rq *locked_rq = rq;
1980
1981	/*
1982	 * We're synchronized against dequeue through DISPATCHING. As @p can't
1983	 * be dequeued, its task_rq and cpus_allowed are stable too.
1984	 *
1985	 * If dispatching to @rq that @p is already on, no lock dancing needed.
1986	 */
1987	if (rq == src_rq && rq == dst_rq) {
1988		dispatch_enqueue(sch, dst_dsq, p,
1989				 enq_flags | SCX_ENQ_CLEAR_OPSS);
1990		return;
1991	}
1992
1993	if (src_rq != dst_rq &&
1994	    unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq, true))) {
1995		dispatch_enqueue(sch, find_global_dsq(sch, p), p,
1996				 enq_flags | SCX_ENQ_CLEAR_OPSS);
1997		return;
1998	}
1999
2000	/*
2001	 * @p is on a possibly remote @src_rq which we need to lock to move the
2002	 * task. If dequeue is in progress, it'd be locking @src_rq and waiting
2003	 * on DISPATCHING, so we can't grab @src_rq lock while holding
2004	 * DISPATCHING.
2005	 *
2006	 * As DISPATCHING guarantees that @p is wholly ours, we can pretend that
2007	 * we're moving from a DSQ and use the same mechanism - mark the task
2008	 * under transfer with holding_cpu, release DISPATCHING and then follow
2009	 * the same protocol. See unlink_dsq_and_lock_src_rq().
2010	 */
2011	p->scx.holding_cpu = raw_smp_processor_id();
2012
2013	/* store_release ensures that dequeue sees the above */
2014	atomic_long_set_release(&p->scx.ops_state, SCX_OPSS_NONE);
2015
2016	/* switch to @src_rq lock */
2017	if (locked_rq != src_rq) {
2018		raw_spin_rq_unlock(locked_rq);
2019		locked_rq = src_rq;
2020		raw_spin_rq_lock(src_rq);
2021	}
2022
2023	/* task_rq couldn't have changed if we're still the holding cpu */
2024	if (likely(p->scx.holding_cpu == raw_smp_processor_id()) &&
2025	    !WARN_ON_ONCE(src_rq != task_rq(p))) {
2026		/*
2027		 * If @p is staying on the same rq, there's no need to go
2028		 * through the full deactivate/activate cycle. Optimize by
2029		 * abbreviating move_remote_task_to_local_dsq().
2030		 */
2031		if (src_rq == dst_rq) {
2032			p->scx.holding_cpu = -1;
2033			dispatch_enqueue(sch, &dst_rq->scx.local_dsq, p,
2034					 enq_flags);
2035		} else {
2036			move_remote_task_to_local_dsq(p, enq_flags,
2037						      src_rq, dst_rq);
2038			/* task has been moved to dst_rq, which is now locked */
2039			locked_rq = dst_rq;
2040		}
2041
2042		/* if the destination CPU is idle, wake it up */
2043		if (sched_class_above(p->sched_class, dst_rq->curr->sched_class))
2044			resched_curr(dst_rq);
2045	}
2046
2047	/* switch back to @rq lock */
2048	if (locked_rq != rq) {
2049		raw_spin_rq_unlock(locked_rq);
2050		raw_spin_rq_lock(rq);
2051	}
2052}
2053
2054/**
2055 * finish_dispatch - Asynchronously finish dispatching a task
2056 * @rq: current rq which is locked
2057 * @p: task to finish dispatching
2058 * @qseq_at_dispatch: qseq when @p started getting dispatched
2059 * @dsq_id: destination DSQ ID
2060 * @enq_flags: %SCX_ENQ_*
2061 *
2062 * Dispatching to local DSQs may need to wait for queueing to complete or
2063 * require rq lock dancing. As we don't wanna do either while inside
2064 * ops.dispatch() to avoid locking order inversion, we split dispatching into
2065 * two parts. scx_bpf_dsq_insert() which is called by ops.dispatch() records the
2066 * task and its qseq. Once ops.dispatch() returns, this function is called to
2067 * finish up.
2068 *
2069 * There is no guarantee that @p is still valid for dispatching or even that it
2070 * was valid in the first place. Make sure that the task is still owned by the
2071 * BPF scheduler and claim the ownership before dispatching.
2072 */
2073static void finish_dispatch(struct scx_sched *sch, struct rq *rq,
2074			    struct task_struct *p,
2075			    unsigned long qseq_at_dispatch,
2076			    u64 dsq_id, u64 enq_flags)
2077{
2078	struct scx_dispatch_q *dsq;
2079	unsigned long opss;
2080
2081	touch_core_sched_dispatch(rq, p);
2082retry:
2083	/*
2084	 * No need for _acquire here. @p is accessed only after a successful
2085	 * try_cmpxchg to DISPATCHING.
2086	 */
2087	opss = atomic_long_read(&p->scx.ops_state);
2088
2089	switch (opss & SCX_OPSS_STATE_MASK) {
2090	case SCX_OPSS_DISPATCHING:
2091	case SCX_OPSS_NONE:
2092		/* someone else already got to it */
2093		return;
2094	case SCX_OPSS_QUEUED:
2095		/*
2096		 * If qseq doesn't match, @p has gone through at least one
2097		 * dispatch/dequeue and re-enqueue cycle between
2098		 * scx_bpf_dsq_insert() and here and we have no claim on it.
2099		 */
2100		if ((opss & SCX_OPSS_QSEQ_MASK) != qseq_at_dispatch)
2101			return;
2102
2103		/*
2104		 * While we know @p is accessible, we don't yet have a claim on
2105		 * it - the BPF scheduler is allowed to dispatch tasks
2106		 * spuriously and there can be a racing dequeue attempt. Let's
2107		 * claim @p by atomically transitioning it from QUEUED to
2108		 * DISPATCHING.
2109		 */
2110		if (likely(atomic_long_try_cmpxchg(&p->scx.ops_state, &opss,
2111						   SCX_OPSS_DISPATCHING)))
2112			break;
2113		goto retry;
2114	case SCX_OPSS_QUEUEING:
2115		/*
2116		 * do_enqueue_task() is in the process of transferring the task
2117		 * to the BPF scheduler while holding @p's rq lock. As we aren't
2118		 * holding any kernel or BPF resource that the enqueue path may
2119		 * depend upon, it's safe to wait.
2120		 */
2121		wait_ops_state(p, opss);
2122		goto retry;
2123	}
2124
2125	BUG_ON(!(p->scx.flags & SCX_TASK_QUEUED));
2126
2127	dsq = find_dsq_for_dispatch(sch, this_rq(), dsq_id, p);
2128
2129	if (dsq->id == SCX_DSQ_LOCAL)
2130		dispatch_to_local_dsq(sch, rq, dsq, p, enq_flags);
2131	else
2132		dispatch_enqueue(sch, dsq, p, enq_flags | SCX_ENQ_CLEAR_OPSS);
2133}
2134
2135static void flush_dispatch_buf(struct scx_sched *sch, struct rq *rq)
2136{
2137	struct scx_dsp_ctx *dspc = this_cpu_ptr(scx_dsp_ctx);
2138	u32 u;
2139
2140	for (u = 0; u < dspc->cursor; u++) {
2141		struct scx_dsp_buf_ent *ent = &dspc->buf[u];
2142
2143		finish_dispatch(sch, rq, ent->task, ent->qseq, ent->dsq_id,
2144				ent->enq_flags);
2145	}
2146
2147	dspc->nr_tasks += dspc->cursor;
2148	dspc->cursor = 0;
2149}
2150
2151static inline void maybe_queue_balance_callback(struct rq *rq)
2152{
2153	lockdep_assert_rq_held(rq);
2154
2155	if (!(rq->scx.flags & SCX_RQ_BAL_CB_PENDING))
2156		return;
2157
2158	queue_balance_callback(rq, &rq->scx.deferred_bal_cb,
2159				deferred_bal_cb_workfn);
2160
2161	rq->scx.flags &= ~SCX_RQ_BAL_CB_PENDING;
2162}
2163
2164static int balance_one(struct rq *rq, struct task_struct *prev)
2165{
2166	struct scx_sched *sch = scx_root;
2167	struct scx_dsp_ctx *dspc = this_cpu_ptr(scx_dsp_ctx);
2168	bool prev_on_scx = prev->sched_class == &ext_sched_class;
2169	bool prev_on_rq = prev->scx.flags & SCX_TASK_QUEUED;
2170	int nr_loops = SCX_DSP_MAX_LOOPS;
2171
2172	lockdep_assert_rq_held(rq);
2173	rq->scx.flags |= SCX_RQ_IN_BALANCE;
2174	rq->scx.flags &= ~SCX_RQ_BAL_KEEP;
2175
2176	if ((sch->ops.flags & SCX_OPS_HAS_CPU_PREEMPT) &&
2177	    unlikely(rq->scx.cpu_released)) {
2178		/*
2179		 * If the previous sched_class for the current CPU was not SCX,
2180		 * notify the BPF scheduler that it again has control of the
2181		 * core. This callback complements ->cpu_release(), which is
2182		 * emitted in switch_class().
2183		 */
2184		if (SCX_HAS_OP(sch, cpu_acquire))
2185			SCX_CALL_OP(sch, SCX_KF_REST, cpu_acquire, rq,
2186				    cpu_of(rq), NULL);
2187		rq->scx.cpu_released = false;
2188	}
2189
2190	if (prev_on_scx) {
2191		update_curr_scx(rq);
2192
2193		/*
2194		 * If @prev is runnable & has slice left, it has priority and
2195		 * fetching more just increases latency for the fetched tasks.
2196		 * Tell pick_task_scx() to keep running @prev. If the BPF
2197		 * scheduler wants to handle this explicitly, it should
2198		 * implement ->cpu_release().
2199		 *
2200		 * See scx_disable_workfn() for the explanation on the bypassing
2201		 * test.
2202		 */
2203		if (prev_on_rq && prev->scx.slice && !scx_rq_bypassing(rq)) {
2204			rq->scx.flags |= SCX_RQ_BAL_KEEP;
2205			goto has_tasks;
2206		}
2207	}
2208
2209	/* if there already are tasks to run, nothing to do */
2210	if (rq->scx.local_dsq.nr)
2211		goto has_tasks;
2212
2213	if (consume_global_dsq(sch, rq))
2214		goto has_tasks;
2215
2216	if (scx_rq_bypassing(rq)) {
2217		if (consume_dispatch_q(sch, rq, &rq->scx.bypass_dsq))
2218			goto has_tasks;
2219		else
2220			goto no_tasks;
2221	}
2222
2223	if (unlikely(!SCX_HAS_OP(sch, dispatch)) || !scx_rq_online(rq))
2224		goto no_tasks;
2225
2226	dspc->rq = rq;
2227
2228	/*
2229	 * The dispatch loop. Because flush_dispatch_buf() may drop the rq lock,
2230	 * the local DSQ might still end up empty after a successful
2231	 * ops.dispatch(). If the local DSQ is empty even after ops.dispatch()
2232	 * produced some tasks, retry. The BPF scheduler may depend on this
2233	 * looping behavior to simplify its implementation.
2234	 */
2235	do {
2236		dspc->nr_tasks = 0;
2237
2238		SCX_CALL_OP(sch, SCX_KF_DISPATCH, dispatch, rq,
2239			    cpu_of(rq), prev_on_scx ? prev : NULL);
2240
2241		flush_dispatch_buf(sch, rq);
2242
2243		if (prev_on_rq && prev->scx.slice) {
2244			rq->scx.flags |= SCX_RQ_BAL_KEEP;
2245			goto has_tasks;
2246		}
2247		if (rq->scx.local_dsq.nr)
2248			goto has_tasks;
2249		if (consume_global_dsq(sch, rq))
2250			goto has_tasks;
2251
2252		/*
2253		 * ops.dispatch() can trap us in this loop by repeatedly
2254		 * dispatching ineligible tasks. Break out once in a while to
2255		 * allow the watchdog to run. As IRQ can't be enabled in
2256		 * balance(), we want to complete this scheduling cycle and then
2257		 * start a new one. IOW, we want to call resched_curr() on the
2258		 * next, most likely idle, task, not the current one. Use
2259		 * scx_kick_cpu() for deferred kicking.
2260		 */
2261		if (unlikely(!--nr_loops)) {
2262			scx_kick_cpu(sch, cpu_of(rq), 0);
2263			break;
2264		}
2265	} while (dspc->nr_tasks);
2266
2267no_tasks:
2268	/*
2269	 * Didn't find another task to run. Keep running @prev unless
2270	 * %SCX_OPS_ENQ_LAST is in effect.
2271	 */
2272	if (prev_on_rq &&
2273	    (!(sch->ops.flags & SCX_OPS_ENQ_LAST) || scx_rq_bypassing(rq))) {
2274		rq->scx.flags |= SCX_RQ_BAL_KEEP;
2275		__scx_add_event(sch, SCX_EV_DISPATCH_KEEP_LAST, 1);
2276		goto has_tasks;
2277	}
2278	rq->scx.flags &= ~SCX_RQ_IN_BALANCE;
2279	return false;
2280
2281has_tasks:
2282	rq->scx.flags &= ~SCX_RQ_IN_BALANCE;
2283	return true;
2284}
2285
2286static void process_ddsp_deferred_locals(struct rq *rq)
2287{
2288	struct task_struct *p;
2289
2290	lockdep_assert_rq_held(rq);
2291
2292	/*
2293	 * Now that @rq can be unlocked, execute the deferred enqueueing of
2294	 * tasks directly dispatched to the local DSQs of other CPUs. See
2295	 * direct_dispatch(). Keep popping from the head instead of using
2296	 * list_for_each_entry_safe() as dispatch_local_dsq() may unlock @rq
2297	 * temporarily.
2298	 */
2299	while ((p = list_first_entry_or_null(&rq->scx.ddsp_deferred_locals,
2300				struct task_struct, scx.dsq_list.node))) {
2301		struct scx_sched *sch = scx_root;
2302		struct scx_dispatch_q *dsq;
2303		u64 dsq_id = p->scx.ddsp_dsq_id;
2304		u64 enq_flags = p->scx.ddsp_enq_flags;
2305
2306		list_del_init(&p->scx.dsq_list.node);
2307		clear_direct_dispatch(p);
2308
2309		dsq = find_dsq_for_dispatch(sch, rq, dsq_id, p);
2310		if (!WARN_ON_ONCE(dsq->id != SCX_DSQ_LOCAL))
2311			dispatch_to_local_dsq(sch, rq, dsq, p, enq_flags);
2312	}
2313}
2314
2315static void set_next_task_scx(struct rq *rq, struct task_struct *p, bool first)
2316{
2317	struct scx_sched *sch = scx_root;
2318
2319	if (p->scx.flags & SCX_TASK_QUEUED) {
2320		/*
2321		 * Core-sched might decide to execute @p before it is
2322		 * dispatched. Call ops_dequeue() to notify the BPF scheduler.
2323		 */
2324		ops_dequeue(rq, p, SCX_DEQ_CORE_SCHED_EXEC);
2325		dispatch_dequeue(rq, p);
2326	}
2327
2328	p->se.exec_start = rq_clock_task(rq);
2329
2330	/* see dequeue_task_scx() on why we skip when !QUEUED */
2331	if (SCX_HAS_OP(sch, running) && (p->scx.flags & SCX_TASK_QUEUED))
2332		SCX_CALL_OP_TASK(sch, SCX_KF_REST, running, rq, p);
2333
2334	clr_task_runnable(p, true);
2335
2336	/*
2337	 * @p is getting newly scheduled or got kicked after someone updated its
2338	 * slice. Refresh whether tick can be stopped. See scx_can_stop_tick().
2339	 */
2340	if ((p->scx.slice == SCX_SLICE_INF) !=
2341	    (bool)(rq->scx.flags & SCX_RQ_CAN_STOP_TICK)) {
2342		if (p->scx.slice == SCX_SLICE_INF)
2343			rq->scx.flags |= SCX_RQ_CAN_STOP_TICK;
2344		else
2345			rq->scx.flags &= ~SCX_RQ_CAN_STOP_TICK;
2346
2347		sched_update_tick_dependency(rq);
2348
2349		/*
2350		 * For now, let's refresh the load_avgs just when transitioning
2351		 * in and out of nohz. In the future, we might want to add a
2352		 * mechanism which calls the following periodically on
2353		 * tick-stopped CPUs.
2354		 */
2355		update_other_load_avgs(rq);
2356	}
2357}
2358
2359static enum scx_cpu_preempt_reason
2360preempt_reason_from_class(const struct sched_class *class)
2361{
2362	if (class == &stop_sched_class)
2363		return SCX_CPU_PREEMPT_STOP;
2364	if (class == &dl_sched_class)
2365		return SCX_CPU_PREEMPT_DL;
2366	if (class == &rt_sched_class)
2367		return SCX_CPU_PREEMPT_RT;
2368	return SCX_CPU_PREEMPT_UNKNOWN;
2369}
2370
2371static void switch_class(struct rq *rq, struct task_struct *next)
2372{
2373	struct scx_sched *sch = scx_root;
2374	const struct sched_class *next_class = next->sched_class;
2375
2376	if (!(sch->ops.flags & SCX_OPS_HAS_CPU_PREEMPT))
2377		return;
2378
2379	/*
2380	 * The callback is conceptually meant to convey that the CPU is no
2381	 * longer under the control of SCX. Therefore, don't invoke the callback
2382	 * if the next class is below SCX (in which case the BPF scheduler has
2383	 * actively decided not to schedule any tasks on the CPU).
2384	 */
2385	if (sched_class_above(&ext_sched_class, next_class))
2386		return;
2387
2388	/*
2389	 * At this point we know that SCX was preempted by a higher priority
2390	 * sched_class, so invoke the ->cpu_release() callback if we have not
2391	 * done so already. We only send the callback once between SCX being
2392	 * preempted, and it regaining control of the CPU.
2393	 *
2394	 * ->cpu_release() complements ->cpu_acquire(), which is emitted the
2395	 *  next time that balance_one() is invoked.
2396	 */
2397	if (!rq->scx.cpu_released) {
2398		if (SCX_HAS_OP(sch, cpu_release)) {
2399			struct scx_cpu_release_args args = {
2400				.reason = preempt_reason_from_class(next_class),
2401				.task = next,
2402			};
2403
2404			SCX_CALL_OP(sch, SCX_KF_CPU_RELEASE, cpu_release, rq,
2405				    cpu_of(rq), &args);
2406		}
2407		rq->scx.cpu_released = true;
2408	}
2409}
2410
2411static void put_prev_task_scx(struct rq *rq, struct task_struct *p,
2412			      struct task_struct *next)
2413{
2414	struct scx_sched *sch = scx_root;
2415
2416	/* see kick_sync_wait_bal_cb() */
2417	smp_store_release(&rq->scx.kick_sync, rq->scx.kick_sync + 1);
2418
2419	update_curr_scx(rq);
2420
2421	/* see dequeue_task_scx() on why we skip when !QUEUED */
2422	if (SCX_HAS_OP(sch, stopping) && (p->scx.flags & SCX_TASK_QUEUED))
2423		SCX_CALL_OP_TASK(sch, SCX_KF_REST, stopping, rq, p, true);
2424
2425	if (p->scx.flags & SCX_TASK_QUEUED) {
2426		set_task_runnable(rq, p);
2427
2428		/*
2429		 * If @p has slice left and is being put, @p is getting
2430		 * preempted by a higher priority scheduler class or core-sched
2431		 * forcing a different task. Leave it at the head of the local
2432		 * DSQ.
2433		 */
2434		if (p->scx.slice && !scx_rq_bypassing(rq)) {
2435			dispatch_enqueue(sch, &rq->scx.local_dsq, p,
2436					 SCX_ENQ_HEAD);
2437			goto switch_class;
2438		}
2439
2440		/*
2441		 * If @p is runnable but we're about to enter a lower
2442		 * sched_class, %SCX_OPS_ENQ_LAST must be set. Tell
2443		 * ops.enqueue() that @p is the only one available for this cpu,
2444		 * which should trigger an explicit follow-up scheduling event.
2445		 */
2446		if (next && sched_class_above(&ext_sched_class, next->sched_class)) {
2447			WARN_ON_ONCE(!(sch->ops.flags & SCX_OPS_ENQ_LAST));
2448			do_enqueue_task(rq, p, SCX_ENQ_LAST, -1);
2449		} else {
2450			do_enqueue_task(rq, p, 0, -1);
2451		}
2452	}
2453
2454switch_class:
2455	if (next && next->sched_class != &ext_sched_class)
2456		switch_class(rq, next);
2457}
2458
2459static void kick_sync_wait_bal_cb(struct rq *rq)
2460{
2461	struct scx_kick_syncs __rcu *ks = __this_cpu_read(scx_kick_syncs);
2462	unsigned long *ksyncs = rcu_dereference_sched(ks)->syncs;
2463	bool waited;
2464	s32 cpu;
2465
2466	/*
2467	 * Drop rq lock and enable IRQs while waiting. IRQs must be enabled
2468	 * — a target CPU may be waiting for us to process an IPI (e.g. TLB
2469	 * flush) while we wait for its kick_sync to advance.
2470	 *
2471	 * Also, keep advancing our own kick_sync so that new kick_sync waits
2472	 * targeting us, which can start after we drop the lock, cannot form
2473	 * cyclic dependencies.
2474	 */
2475retry:
2476	waited = false;
2477	for_each_cpu(cpu, rq->scx.cpus_to_sync) {
2478		/*
2479		 * smp_load_acquire() pairs with smp_store_release() on
2480		 * kick_sync updates on the target CPUs.
2481		 */
2482		if (cpu == cpu_of(rq) ||
2483		    smp_load_acquire(&cpu_rq(cpu)->scx.kick_sync) != ksyncs[cpu]) {
2484			cpumask_clear_cpu(cpu, rq->scx.cpus_to_sync);
2485			continue;
2486		}
2487
2488		raw_spin_rq_unlock_irq(rq);
2489		while (READ_ONCE(cpu_rq(cpu)->scx.kick_sync) == ksyncs[cpu]) {
2490			smp_store_release(&rq->scx.kick_sync, rq->scx.kick_sync + 1);
2491			cpu_relax();
2492		}
2493		raw_spin_rq_lock_irq(rq);
2494		waited = true;
2495	}
2496
2497	if (waited)
2498		goto retry;
2499}
2500
2501static struct task_struct *first_local_task(struct rq *rq)
2502{
2503	return list_first_entry_or_null(&rq->scx.local_dsq.list,
2504					struct task_struct, scx.dsq_list.node);
2505}
2506
2507static struct task_struct *
2508do_pick_task_scx(struct rq *rq, struct rq_flags *rf, bool force_scx)
2509{
2510	struct task_struct *prev = rq->curr;
2511	bool keep_prev;
2512	struct task_struct *p;
2513
2514	/* see kick_sync_wait_bal_cb() */
2515	smp_store_release(&rq->scx.kick_sync, rq->scx.kick_sync + 1);
2516
2517	rq_modified_clear(rq);
2518
2519	rq_unpin_lock(rq, rf);
2520	balance_one(rq, prev);
2521	rq_repin_lock(rq, rf);
2522	maybe_queue_balance_callback(rq);
2523
2524	/*
2525	 * Defer to a balance callback which can drop rq lock and enable
2526	 * IRQs. Waiting directly in the pick path would deadlock against
2527	 * CPUs sending us IPIs (e.g. TLB flushes) while we wait for them.
2528	 */
2529	if (unlikely(rq->scx.kick_sync_pending)) {
2530		rq->scx.kick_sync_pending = false;
2531		queue_balance_callback(rq, &rq->scx.kick_sync_bal_cb,
2532				       kick_sync_wait_bal_cb);
2533	}
2534
2535	/*
2536	 * If any higher-priority sched class enqueued a runnable task on
2537	 * this rq during balance_one(), abort and return RETRY_TASK, so
2538	 * that the scheduler loop can restart.
2539	 *
2540	 * If @force_scx is true, always try to pick a SCHED_EXT task,
2541	 * regardless of any higher-priority sched classes activity.
2542	 */
2543	if (!force_scx && rq_modified_above(rq, &ext_sched_class))
2544		return RETRY_TASK;
2545
2546	keep_prev = rq->scx.flags & SCX_RQ_BAL_KEEP;
2547	if (unlikely(keep_prev &&
2548		     prev->sched_class != &ext_sched_class)) {
2549		WARN_ON_ONCE(scx_enable_state() == SCX_ENABLED);
2550		keep_prev = false;
2551	}
2552
2553	/*
2554	 * If balance_one() is telling us to keep running @prev, replenish slice
2555	 * if necessary and keep running @prev. Otherwise, pop the first one
2556	 * from the local DSQ.
2557	 */
2558	if (keep_prev) {
2559		p = prev;
2560		if (!p->scx.slice)
2561			refill_task_slice_dfl(rcu_dereference_sched(scx_root), p);
2562	} else {
2563		p = first_local_task(rq);
2564		if (!p)
2565			return NULL;
2566
2567		if (unlikely(!p->scx.slice)) {
2568			struct scx_sched *sch = rcu_dereference_sched(scx_root);
2569
2570			if (!scx_rq_bypassing(rq) && !sch->warned_zero_slice) {
2571				printk_deferred(KERN_WARNING "sched_ext: %s[%d] has zero slice in %s()\n",
2572						p->comm, p->pid, __func__);
2573				sch->warned_zero_slice = true;
2574			}
2575			refill_task_slice_dfl(sch, p);
2576		}
2577	}
2578
2579	return p;
2580}
2581
2582static struct task_struct *pick_task_scx(struct rq *rq, struct rq_flags *rf)
2583{
2584	return do_pick_task_scx(rq, rf, false);
2585}
2586
2587#ifdef CONFIG_SCHED_CORE
2588/**
2589 * scx_prio_less - Task ordering for core-sched
2590 * @a: task A
2591 * @b: task B
2592 * @in_fi: in forced idle state
2593 *
2594 * Core-sched is implemented as an additional scheduling layer on top of the
2595 * usual sched_class'es and needs to find out the expected task ordering. For
2596 * SCX, core-sched calls this function to interrogate the task ordering.
2597 *
2598 * Unless overridden by ops.core_sched_before(), @p->scx.core_sched_at is used
2599 * to implement the default task ordering. The older the timestamp, the higher
2600 * priority the task - the global FIFO ordering matching the default scheduling
2601 * behavior.
2602 *
2603 * When ops.core_sched_before() is enabled, @p->scx.core_sched_at is used to
2604 * implement FIFO ordering within each local DSQ. See pick_task_scx().
2605 */
2606bool scx_prio_less(const struct task_struct *a, const struct task_struct *b,
2607		   bool in_fi)
2608{
2609	struct scx_sched *sch = scx_root;
2610
2611	/*
2612	 * The const qualifiers are dropped from task_struct pointers when
2613	 * calling ops.core_sched_before(). Accesses are controlled by the
2614	 * verifier.
2615	 */
2616	if (SCX_HAS_OP(sch, core_sched_before) &&
2617	    !scx_rq_bypassing(task_rq(a)))
2618		return SCX_CALL_OP_2TASKS_RET(sch, SCX_KF_REST, core_sched_before,
2619					      NULL,
2620					      (struct task_struct *)a,
2621					      (struct task_struct *)b);
2622	else
2623		return time_after64(a->scx.core_sched_at, b->scx.core_sched_at);
2624}
2625#endif	/* CONFIG_SCHED_CORE */
2626
2627static int select_task_rq_scx(struct task_struct *p, int prev_cpu, int wake_flags)
2628{
2629	struct scx_sched *sch = scx_root;
2630	bool rq_bypass;
2631
2632	/*
2633	 * sched_exec() calls with %WF_EXEC when @p is about to exec(2) as it
2634	 * can be a good migration opportunity with low cache and memory
2635	 * footprint. Returning a CPU different than @prev_cpu triggers
2636	 * immediate rq migration. However, for SCX, as the current rq
2637	 * association doesn't dictate where the task is going to run, this
2638	 * doesn't fit well. If necessary, we can later add a dedicated method
2639	 * which can decide to preempt self to force it through the regular
2640	 * scheduling path.
2641	 */
2642	if (unlikely(wake_flags & WF_EXEC))
2643		return prev_cpu;
2644
2645	rq_bypass = scx_rq_bypassing(task_rq(p));
2646	if (likely(SCX_HAS_OP(sch, select_cpu)) && !rq_bypass) {
2647		s32 cpu;
2648		struct task_struct **ddsp_taskp;
2649
2650		ddsp_taskp = this_cpu_ptr(&direct_dispatch_task);
2651		WARN_ON_ONCE(*ddsp_taskp);
2652		*ddsp_taskp = p;
2653
2654		cpu = SCX_CALL_OP_TASK_RET(sch,
2655					   SCX_KF_ENQUEUE | SCX_KF_SELECT_CPU,
2656					   select_cpu, NULL, p, prev_cpu,
2657					   wake_flags);
2658		p->scx.selected_cpu = cpu;
2659		*ddsp_taskp = NULL;
2660		if (ops_cpu_valid(sch, cpu, "from ops.select_cpu()"))
2661			return cpu;
2662		else
2663			return prev_cpu;
2664	} else {
2665		s32 cpu;
2666
2667		cpu = scx_select_cpu_dfl(p, prev_cpu, wake_flags, NULL, 0);
2668		if (cpu >= 0) {
2669			refill_task_slice_dfl(sch, p);
2670			p->scx.ddsp_dsq_id = SCX_DSQ_LOCAL;
2671		} else {
2672			cpu = prev_cpu;
2673		}
2674		p->scx.selected_cpu = cpu;
2675
2676		if (rq_bypass)
2677			__scx_add_event(sch, SCX_EV_BYPASS_DISPATCH, 1);
2678		return cpu;
2679	}
2680}
2681
2682static void task_woken_scx(struct rq *rq, struct task_struct *p)
2683{
2684	run_deferred(rq);
2685}
2686
2687static void set_cpus_allowed_scx(struct task_struct *p,
2688				 struct affinity_context *ac)
2689{
2690	struct scx_sched *sch = scx_root;
2691
2692	set_cpus_allowed_common(p, ac);
2693
2694	if (task_dead_and_done(p))
2695		return;
2696
2697	/*
2698	 * The effective cpumask is stored in @p->cpus_ptr which may temporarily
2699	 * differ from the configured one in @p->cpus_mask. Always tell the bpf
2700	 * scheduler the effective one.
2701	 *
2702	 * Fine-grained memory write control is enforced by BPF making the const
2703	 * designation pointless. Cast it away when calling the operation.
2704	 */
2705	if (SCX_HAS_OP(sch, set_cpumask))
2706		SCX_CALL_OP_TASK(sch, SCX_KF_REST, set_cpumask, NULL,
2707				 p, (struct cpumask *)p->cpus_ptr);
2708}
2709
2710static void handle_hotplug(struct rq *rq, bool online)
2711{
2712	struct scx_sched *sch = scx_root;
2713	int cpu = cpu_of(rq);
2714
2715	atomic_long_inc(&scx_hotplug_seq);
2716
2717	/*
2718	 * scx_root updates are protected by cpus_read_lock() and will stay
2719	 * stable here. Note that we can't depend on scx_enabled() test as the
2720	 * hotplug ops need to be enabled before __scx_enabled is set.
2721	 */
2722	if (unlikely(!sch))
2723		return;
2724
2725	if (scx_enabled())
2726		scx_idle_update_selcpu_topology(&sch->ops);
2727
2728	if (online && SCX_HAS_OP(sch, cpu_online))
2729		SCX_CALL_OP(sch, SCX_KF_UNLOCKED, cpu_online, NULL, cpu);
2730	else if (!online && SCX_HAS_OP(sch, cpu_offline))
2731		SCX_CALL_OP(sch, SCX_KF_UNLOCKED, cpu_offline, NULL, cpu);
2732	else
2733		scx_exit(sch, SCX_EXIT_UNREG_KERN,
2734			 SCX_ECODE_ACT_RESTART | SCX_ECODE_RSN_HOTPLUG,
2735			 "cpu %d going %s, exiting scheduler", cpu,
2736			 online ? "online" : "offline");
2737}
2738
2739void scx_rq_activate(struct rq *rq)
2740{
2741	handle_hotplug(rq, true);
2742}
2743
2744void scx_rq_deactivate(struct rq *rq)
2745{
2746	handle_hotplug(rq, false);
2747}
2748
2749static void rq_online_scx(struct rq *rq)
2750{
2751	rq->scx.flags |= SCX_RQ_ONLINE;
2752}
2753
2754static void rq_offline_scx(struct rq *rq)
2755{
2756	rq->scx.flags &= ~SCX_RQ_ONLINE;
2757}
2758
2759
2760static bool check_rq_for_timeouts(struct rq *rq)
2761{
2762	struct scx_sched *sch;
2763	struct task_struct *p;
2764	struct rq_flags rf;
2765	bool timed_out = false;
2766
2767	rq_lock_irqsave(rq, &rf);
2768	sch = rcu_dereference_bh(scx_root);
2769	if (unlikely(!sch))
2770		goto out_unlock;
2771
2772	list_for_each_entry(p, &rq->scx.runnable_list, scx.runnable_node) {
2773		unsigned long last_runnable = p->scx.runnable_at;
2774
2775		if (unlikely(time_after(jiffies,
2776					last_runnable + scx_watchdog_timeout))) {
2777			u32 dur_ms = jiffies_to_msecs(jiffies - last_runnable);
2778
2779			scx_exit(sch, SCX_EXIT_ERROR_STALL, 0,
2780				 "%s[%d] failed to run for %u.%03us",
2781				 p->comm, p->pid, dur_ms / 1000, dur_ms % 1000);
2782			timed_out = true;
2783			break;
2784		}
2785	}
2786out_unlock:
2787	rq_unlock_irqrestore(rq, &rf);
2788	return timed_out;
2789}
2790
2791static void scx_watchdog_workfn(struct work_struct *work)
2792{
2793	int cpu;
2794
2795	WRITE_ONCE(scx_watchdog_timestamp, jiffies);
2796
2797	for_each_online_cpu(cpu) {
2798		if (unlikely(check_rq_for_timeouts(cpu_rq(cpu))))
2799			break;
2800
2801		cond_resched();
2802	}
2803	queue_delayed_work(system_unbound_wq, to_delayed_work(work),
2804			   scx_watchdog_timeout / 2);
2805}
2806
2807void scx_tick(struct rq *rq)
2808{
2809	struct scx_sched *sch;
2810	unsigned long last_check;
2811
2812	if (!scx_enabled())
2813		return;
2814
2815	sch = rcu_dereference_bh(scx_root);
2816	if (unlikely(!sch))
2817		return;
2818
2819	last_check = READ_ONCE(scx_watchdog_timestamp);
2820	if (unlikely(time_after(jiffies,
2821				last_check + READ_ONCE(scx_watchdog_timeout)))) {
2822		u32 dur_ms = jiffies_to_msecs(jiffies - last_check);
2823
2824		scx_exit(sch, SCX_EXIT_ERROR_STALL, 0,
2825			 "watchdog failed to check in for %u.%03us",
2826			 dur_ms / 1000, dur_ms % 1000);
2827	}
2828
2829	update_other_load_avgs(rq);
2830}
2831
2832static void task_tick_scx(struct rq *rq, struct task_struct *curr, int queued)
2833{
2834	struct scx_sched *sch = scx_root;
2835
2836	update_curr_scx(rq);
2837
2838	/*
2839	 * While disabling, always resched and refresh core-sched timestamp as
2840	 * we can't trust the slice management or ops.core_sched_before().
2841	 */
2842	if (scx_rq_bypassing(rq)) {
2843		curr->scx.slice = 0;
2844		touch_core_sched(rq, curr);
2845	} else if (SCX_HAS_OP(sch, tick)) {
2846		SCX_CALL_OP_TASK(sch, SCX_KF_REST, tick, rq, curr);
2847	}
2848
2849	if (!curr->scx.slice)
2850		resched_curr(rq);
2851}
2852
2853#ifdef CONFIG_EXT_GROUP_SCHED
2854static struct cgroup *tg_cgrp(struct task_group *tg)
2855{
2856	/*
2857	 * If CGROUP_SCHED is disabled, @tg is NULL. If @tg is an autogroup,
2858	 * @tg->css.cgroup is NULL. In both cases, @tg can be treated as the
2859	 * root cgroup.
2860	 */
2861	if (tg && tg->css.cgroup)
2862		return tg->css.cgroup;
2863	else
2864		return &cgrp_dfl_root.cgrp;
2865}
2866
2867#define SCX_INIT_TASK_ARGS_CGROUP(tg)		.cgroup = tg_cgrp(tg),
2868
2869#else	/* CONFIG_EXT_GROUP_SCHED */
2870
2871#define SCX_INIT_TASK_ARGS_CGROUP(tg)
2872
2873#endif	/* CONFIG_EXT_GROUP_SCHED */
2874
2875static enum scx_task_state scx_get_task_state(const struct task_struct *p)
2876{
2877	return (p->scx.flags & SCX_TASK_STATE_MASK) >> SCX_TASK_STATE_SHIFT;
2878}
2879
2880static void scx_set_task_state(struct task_struct *p, enum scx_task_state state)
2881{
2882	enum scx_task_state prev_state = scx_get_task_state(p);
2883	bool warn = false;
2884
2885	BUILD_BUG_ON(SCX_TASK_NR_STATES > (1 << SCX_TASK_STATE_BITS));
2886
2887	switch (state) {
2888	case SCX_TASK_NONE:
2889		break;
2890	case SCX_TASK_INIT:
2891		warn = prev_state != SCX_TASK_NONE;
2892		break;
2893	case SCX_TASK_READY:
2894		warn = prev_state == SCX_TASK_NONE;
2895		break;
2896	case SCX_TASK_ENABLED:
2897		warn = prev_state != SCX_TASK_READY;
2898		break;
2899	default:
2900		warn = true;
2901		return;
2902	}
2903
2904	WARN_ONCE(warn, "sched_ext: Invalid task state transition %d -> %d for %s[%d]",
2905		  prev_state, state, p->comm, p->pid);
2906
2907	p->scx.flags &= ~SCX_TASK_STATE_MASK;
2908	p->scx.flags |= state << SCX_TASK_STATE_SHIFT;
2909}
2910
2911static int scx_init_task(struct task_struct *p, struct task_group *tg, bool fork)
2912{
2913	struct scx_sched *sch = scx_root;
2914	int ret;
2915
2916	p->scx.disallow = false;
2917
2918	if (SCX_HAS_OP(sch, init_task)) {
2919		struct scx_init_task_args args = {
2920			SCX_INIT_TASK_ARGS_CGROUP(tg)
2921			.fork = fork,
2922		};
2923
2924		ret = SCX_CALL_OP_RET(sch, SCX_KF_UNLOCKED, init_task, NULL,
2925				      p, &args);
2926		if (unlikely(ret)) {
2927			ret = ops_sanitize_err(sch, "init_task", ret);
2928			return ret;
2929		}
2930	}
2931
2932	scx_set_task_state(p, SCX_TASK_INIT);
2933
2934	if (p->scx.disallow) {
2935		if (!fork) {
2936			struct rq *rq;
2937			struct rq_flags rf;
2938
2939			rq = task_rq_lock(p, &rf);
2940
2941			/*
2942			 * We're in the load path and @p->policy will be applied
2943			 * right after. Reverting @p->policy here and rejecting
2944			 * %SCHED_EXT transitions from scx_check_setscheduler()
2945			 * guarantees that if ops.init_task() sets @p->disallow,
2946			 * @p can never be in SCX.
2947			 */
2948			if (p->policy == SCHED_EXT) {
2949				p->policy = SCHED_NORMAL;
2950				atomic_long_inc(&scx_nr_rejected);
2951			}
2952
2953			task_rq_unlock(rq, p, &rf);
2954		} else if (p->policy == SCHED_EXT) {
2955			scx_error(sch, "ops.init_task() set task->scx.disallow for %s[%d] during fork",
2956				  p->comm, p->pid);
2957		}
2958	}
2959
2960	p->scx.flags |= SCX_TASK_RESET_RUNNABLE_AT;
2961	return 0;
2962}
2963
2964static void scx_enable_task(struct task_struct *p)
2965{
2966	struct scx_sched *sch = scx_root;
2967	struct rq *rq = task_rq(p);
2968	u32 weight;
2969
2970	lockdep_assert_rq_held(rq);
2971
2972	/*
2973	 * Set the weight before calling ops.enable() so that the scheduler
2974	 * doesn't see a stale value if they inspect the task struct.
2975	 */
2976	if (task_has_idle_policy(p))
2977		weight = WEIGHT_IDLEPRIO;
2978	else
2979		weight = sched_prio_to_weight[p->static_prio - MAX_RT_PRIO];
2980
2981	p->scx.weight = sched_weight_to_cgroup(weight);
2982
2983	if (SCX_HAS_OP(sch, enable))
2984		SCX_CALL_OP_TASK(sch, SCX_KF_REST, enable, rq, p);
2985	scx_set_task_state(p, SCX_TASK_ENABLED);
2986
2987	if (SCX_HAS_OP(sch, set_weight))
2988		SCX_CALL_OP_TASK(sch, SCX_KF_REST, set_weight, rq,
2989				 p, p->scx.weight);
2990}
2991
2992static void scx_disable_task(struct task_struct *p)
2993{
2994	struct scx_sched *sch = scx_root;
2995	struct rq *rq = task_rq(p);
2996
2997	lockdep_assert_rq_held(rq);
2998	WARN_ON_ONCE(scx_get_task_state(p) != SCX_TASK_ENABLED);
2999
3000	clear_direct_dispatch(p);
3001
3002	if (SCX_HAS_OP(sch, disable))
3003		SCX_CALL_OP_TASK(sch, SCX_KF_REST, disable, rq, p);
3004	scx_set_task_state(p, SCX_TASK_READY);
3005}
3006
3007static void scx_exit_task(struct task_struct *p)
3008{
3009	struct scx_sched *sch = scx_root;
3010	struct scx_exit_task_args args = {
3011		.cancelled = false,
3012	};
3013
3014	lockdep_assert_rq_held(task_rq(p));
3015
3016	switch (scx_get_task_state(p)) {
3017	case SCX_TASK_NONE:
3018		return;
3019	case SCX_TASK_INIT:
3020		args.cancelled = true;
3021		break;
3022	case SCX_TASK_READY:
3023		break;
3024	case SCX_TASK_ENABLED:
3025		scx_disable_task(p);
3026		break;
3027	default:
3028		WARN_ON_ONCE(true);
3029		return;
3030	}
3031
3032	if (SCX_HAS_OP(sch, exit_task))
3033		SCX_CALL_OP_TASK(sch, SCX_KF_REST, exit_task, task_rq(p),
3034				 p, &args);
3035	scx_set_task_state(p, SCX_TASK_NONE);
3036}
3037
3038void init_scx_entity(struct sched_ext_entity *scx)
3039{
3040	memset(scx, 0, sizeof(*scx));
3041	INIT_LIST_HEAD(&scx->dsq_list.node);
3042	RB_CLEAR_NODE(&scx->dsq_priq);
3043	scx->sticky_cpu = -1;
3044	scx->holding_cpu = -1;
3045	INIT_LIST_HEAD(&scx->runnable_node);
3046	scx->runnable_at = jiffies;
3047	scx->ddsp_dsq_id = SCX_DSQ_INVALID;
3048	scx->slice = READ_ONCE(scx_slice_dfl);
3049}
3050
3051void scx_pre_fork(struct task_struct *p)
3052{
3053	/*
3054	 * BPF scheduler enable/disable paths want to be able to iterate and
3055	 * update all tasks which can become complex when racing forks. As
3056	 * enable/disable are very cold paths, let's use a percpu_rwsem to
3057	 * exclude forks.
3058	 */
3059	percpu_down_read(&scx_fork_rwsem);
3060}
3061
3062int scx_fork(struct task_struct *p)
3063{
3064	percpu_rwsem_assert_held(&scx_fork_rwsem);
3065
3066	if (scx_init_task_enabled)
3067		return scx_init_task(p, task_group(p), true);
3068	else
3069		return 0;
3070}
3071
3072void scx_post_fork(struct task_struct *p)
3073{
3074	if (scx_init_task_enabled) {
3075		scx_set_task_state(p, SCX_TASK_READY);
3076
3077		/*
3078		 * Enable the task immediately if it's running on sched_ext.
3079		 * Otherwise, it'll be enabled in switching_to_scx() if and
3080		 * when it's ever configured to run with a SCHED_EXT policy.
3081		 */
3082		if (p->sched_class == &ext_sched_class) {
3083			struct rq_flags rf;
3084			struct rq *rq;
3085
3086			rq = task_rq_lock(p, &rf);
3087			scx_enable_task(p);
3088			task_rq_unlock(rq, p, &rf);
3089		}
3090	}
3091
3092	raw_spin_lock_irq(&scx_tasks_lock);
3093	list_add_tail(&p->scx.tasks_node, &scx_tasks);
3094	raw_spin_unlock_irq(&scx_tasks_lock);
3095
3096	percpu_up_read(&scx_fork_rwsem);
3097}
3098
3099void scx_cancel_fork(struct task_struct *p)
3100{
3101	if (scx_enabled()) {
3102		struct rq *rq;
3103		struct rq_flags rf;
3104
3105		rq = task_rq_lock(p, &rf);
3106		WARN_ON_ONCE(scx_get_task_state(p) >= SCX_TASK_READY);
3107		scx_exit_task(p);
3108		task_rq_unlock(rq, p, &rf);
3109	}
3110
3111	percpu_up_read(&scx_fork_rwsem);
3112}
3113
3114/**
3115 * task_dead_and_done - Is a task dead and done running?
3116 * @p: target task
3117 *
3118 * Once sched_ext_dead() removes the dead task from scx_tasks and exits it, the
3119 * task no longer exists from SCX's POV. However, certain sched_class ops may be
3120 * invoked on these dead tasks leading to failures - e.g. sched_setscheduler()
3121 * may try to switch a task which finished sched_ext_dead() back into SCX
3122 * triggering invalid SCX task state transitions and worse.
3123 *
3124 * Once a task has finished the final switch, sched_ext_dead() is the only thing
3125 * that needs to happen on the task. Use this test to short-circuit sched_class
3126 * operations which may be called on dead tasks.
3127 */
3128static bool task_dead_and_done(struct task_struct *p)
3129{
3130	struct rq *rq = task_rq(p);
3131
3132	lockdep_assert_rq_held(rq);
3133
3134	/*
3135	 * In do_task_dead(), a dying task sets %TASK_DEAD with preemption
3136	 * disabled and __schedule(). If @p has %TASK_DEAD set and off CPU, @p
3137	 * won't ever run again.
3138	 */
3139	return unlikely(READ_ONCE(p->__state) == TASK_DEAD) &&
3140		!task_on_cpu(rq, p);
3141}
3142
3143void sched_ext_dead(struct task_struct *p)
3144{
3145	unsigned long flags;
3146
3147	/*
3148	 * By the time control reaches here, @p has %TASK_DEAD set, switched out
3149	 * for the last time and then dropped the rq lock - task_dead_and_done()
3150	 * should be returning %true nullifying the straggling sched_class ops.
3151	 * Remove from scx_tasks and exit @p.
3152	 */
3153	raw_spin_lock_irqsave(&scx_tasks_lock, flags);
3154	list_del_init(&p->scx.tasks_node);
3155	raw_spin_unlock_irqrestore(&scx_tasks_lock, flags);
3156
3157	/*
3158	 * @p is off scx_tasks and wholly ours. scx_enable()'s READY -> ENABLED
3159	 * transitions can't race us. Disable ops for @p.
3160	 */
3161	if (scx_get_task_state(p) != SCX_TASK_NONE) {
3162		struct rq_flags rf;
3163		struct rq *rq;
3164
3165		rq = task_rq_lock(p, &rf);
3166		scx_exit_task(p);
3167		task_rq_unlock(rq, p, &rf);
3168	}
3169}
3170
3171static void reweight_task_scx(struct rq *rq, struct task_struct *p,
3172			      const struct load_weight *lw)
3173{
3174	struct scx_sched *sch = scx_root;
3175
3176	lockdep_assert_rq_held(task_rq(p));
3177
3178	if (task_dead_and_done(p))
3179		return;
3180
3181	p->scx.weight = sched_weight_to_cgroup(scale_load_down(lw->weight));
3182	if (SCX_HAS_OP(sch, set_weight))
3183		SCX_CALL_OP_TASK(sch, SCX_KF_REST, set_weight, rq,
3184				 p, p->scx.weight);
3185}
3186
3187static void prio_changed_scx(struct rq *rq, struct task_struct *p, u64 oldprio)
3188{
3189}
3190
3191static void switching_to_scx(struct rq *rq, struct task_struct *p)
3192{
3193	struct scx_sched *sch = scx_root;
3194
3195	if (task_dead_and_done(p))
3196		return;
3197
3198	scx_enable_task(p);
3199
3200	/*
3201	 * set_cpus_allowed_scx() is not called while @p is associated with a
3202	 * different scheduler class. Keep the BPF scheduler up-to-date.
3203	 */
3204	if (SCX_HAS_OP(sch, set_cpumask))
3205		SCX_CALL_OP_TASK(sch, SCX_KF_REST, set_cpumask, rq,
3206				 p, (struct cpumask *)p->cpus_ptr);
3207}
3208
3209static void switched_from_scx(struct rq *rq, struct task_struct *p)
3210{
3211	if (task_dead_and_done(p))
3212		return;
3213
3214	scx_disable_task(p);
3215}
3216
3217static void wakeup_preempt_scx(struct rq *rq, struct task_struct *p,int wake_flags) {}
3218static void switched_to_scx(struct rq *rq, struct task_struct *p) {}
3219
3220int scx_check_setscheduler(struct task_struct *p, int policy)
3221{
3222	lockdep_assert_rq_held(task_rq(p));
3223
3224	/* if disallow, reject transitioning into SCX */
3225	if (scx_enabled() && READ_ONCE(p->scx.disallow) &&
3226	    p->policy != policy && policy == SCHED_EXT)
3227		return -EACCES;
3228
3229	return 0;
3230}
3231
3232#ifdef CONFIG_NO_HZ_FULL
3233bool scx_can_stop_tick(struct rq *rq)
3234{
3235	struct task_struct *p = rq->curr;
3236
3237	if (scx_rq_bypassing(rq))
3238		return false;
3239
3240	if (p->sched_class != &ext_sched_class)
3241		return true;
3242
3243	/*
3244	 * @rq can dispatch from different DSQs, so we can't tell whether it
3245	 * needs the tick or not by looking at nr_running. Allow stopping ticks
3246	 * iff the BPF scheduler indicated so. See set_next_task_scx().
3247	 */
3248	return rq->scx.flags & SCX_RQ_CAN_STOP_TICK;
3249}
3250#endif
3251
3252#ifdef CONFIG_EXT_GROUP_SCHED
3253
3254DEFINE_STATIC_PERCPU_RWSEM(scx_cgroup_ops_rwsem);
3255static bool scx_cgroup_enabled;
3256
3257void scx_tg_init(struct task_group *tg)
3258{
3259	tg->scx.weight = CGROUP_WEIGHT_DFL;
3260	tg->scx.bw_period_us = default_bw_period_us();
3261	tg->scx.bw_quota_us = RUNTIME_INF;
3262	tg->scx.idle = false;
3263}
3264
3265int scx_tg_online(struct task_group *tg)
3266{
3267	struct scx_sched *sch = scx_root;
3268	int ret = 0;
3269
3270	WARN_ON_ONCE(tg->scx.flags & (SCX_TG_ONLINE | SCX_TG_INITED));
3271
3272	if (scx_cgroup_enabled) {
3273		if (SCX_HAS_OP(sch, cgroup_init)) {
3274			struct scx_cgroup_init_args args =
3275				{ .weight = tg->scx.weight,
3276				  .bw_period_us = tg->scx.bw_period_us,
3277				  .bw_quota_us = tg->scx.bw_quota_us,
3278				  .bw_burst_us = tg->scx.bw_burst_us };
3279
3280			ret = SCX_CALL_OP_RET(sch, SCX_KF_UNLOCKED, cgroup_init,
3281					      NULL, tg->css.cgroup, &args);
3282			if (ret)
3283				ret = ops_sanitize_err(sch, "cgroup_init", ret);
3284		}
3285		if (ret == 0)
3286			tg->scx.flags |= SCX_TG_ONLINE | SCX_TG_INITED;
3287	} else {
3288		tg->scx.flags |= SCX_TG_ONLINE;
3289	}
3290
3291	return ret;
3292}
3293
3294void scx_tg_offline(struct task_group *tg)
3295{
3296	struct scx_sched *sch = scx_root;
3297
3298	WARN_ON_ONCE(!(tg->scx.flags & SCX_TG_ONLINE));
3299
3300	if (scx_cgroup_enabled && SCX_HAS_OP(sch, cgroup_exit) &&
3301	    (tg->scx.flags & SCX_TG_INITED))
3302		SCX_CALL_OP(sch, SCX_KF_UNLOCKED, cgroup_exit, NULL,
3303			    tg->css.cgroup);
3304	tg->scx.flags &= ~(SCX_TG_ONLINE | SCX_TG_INITED);
3305}
3306
3307int scx_cgroup_can_attach(struct cgroup_taskset *tset)
3308{
3309	struct scx_sched *sch = scx_root;
3310	struct cgroup_subsys_state *css;
3311	struct task_struct *p;
3312	int ret;
3313
3314	if (!scx_cgroup_enabled)
3315		return 0;
3316
3317	cgroup_taskset_for_each(p, css, tset) {
3318		struct cgroup *from = tg_cgrp(task_group(p));
3319		struct cgroup *to = tg_cgrp(css_tg(css));
3320
3321		WARN_ON_ONCE(p->scx.cgrp_moving_from);
3322
3323		/*
3324		 * sched_move_task() omits identity migrations. Let's match the
3325		 * behavior so that ops.cgroup_prep_move() and ops.cgroup_move()
3326		 * always match one-to-one.
3327		 */
3328		if (from == to)
3329			continue;
3330
3331		if (SCX_HAS_OP(sch, cgroup_prep_move)) {
3332			ret = SCX_CALL_OP_RET(sch, SCX_KF_UNLOCKED,
3333					      cgroup_prep_move, NULL,
3334					      p, from, css->cgroup);
3335			if (ret)
3336				goto err;
3337		}
3338
3339		p->scx.cgrp_moving_from = from;
3340	}
3341
3342	return 0;
3343
3344err:
3345	cgroup_taskset_for_each(p, css, tset) {
3346		if (SCX_HAS_OP(sch, cgroup_cancel_move) &&
3347		    p->scx.cgrp_moving_from)
3348			SCX_CALL_OP(sch, SCX_KF_UNLOCKED, cgroup_cancel_move, NULL,
3349				    p, p->scx.cgrp_moving_from, css->cgroup);
3350		p->scx.cgrp_moving_from = NULL;
3351	}
3352
3353	return ops_sanitize_err(sch, "cgroup_prep_move", ret);
3354}
3355
3356void scx_cgroup_move_task(struct task_struct *p)
3357{
3358	struct scx_sched *sch = scx_root;
3359
3360	if (!scx_cgroup_enabled)
3361		return;
3362
3363	/*
3364	 * @p must have ops.cgroup_prep_move() called on it and thus
3365	 * cgrp_moving_from set.
3366	 */
3367	if (SCX_HAS_OP(sch, cgroup_move) &&
3368	    !WARN_ON_ONCE(!p->scx.cgrp_moving_from))
3369		SCX_CALL_OP_TASK(sch, SCX_KF_UNLOCKED, cgroup_move, NULL,
3370				 p, p->scx.cgrp_moving_from,
3371				 tg_cgrp(task_group(p)));
3372	p->scx.cgrp_moving_from = NULL;
3373}
3374
3375void scx_cgroup_cancel_attach(struct cgroup_taskset *tset)
3376{
3377	struct scx_sched *sch = scx_root;
3378	struct cgroup_subsys_state *css;
3379	struct task_struct *p;
3380
3381	if (!scx_cgroup_enabled)
3382		return;
3383
3384	cgroup_taskset_for_each(p, css, tset) {
3385		if (SCX_HAS_OP(sch, cgroup_cancel_move) &&
3386		    p->scx.cgrp_moving_from)
3387			SCX_CALL_OP(sch, SCX_KF_UNLOCKED, cgroup_cancel_move, NULL,
3388				    p, p->scx.cgrp_moving_from, css->cgroup);
3389		p->scx.cgrp_moving_from = NULL;
3390	}
3391}
3392
3393void scx_group_set_weight(struct task_group *tg, unsigned long weight)
3394{
3395	struct scx_sched *sch = scx_root;
3396
3397	percpu_down_read(&scx_cgroup_ops_rwsem);
3398
3399	if (scx_cgroup_enabled && SCX_HAS_OP(sch, cgroup_set_weight) &&
3400	    tg->scx.weight != weight)
3401		SCX_CALL_OP(sch, SCX_KF_UNLOCKED, cgroup_set_weight, NULL,
3402			    tg_cgrp(tg), weight);
3403
3404	tg->scx.weight = weight;
3405
3406	percpu_up_read(&scx_cgroup_ops_rwsem);
3407}
3408
3409void scx_group_set_idle(struct task_group *tg, bool idle)
3410{
3411	struct scx_sched *sch = scx_root;
3412
3413	percpu_down_read(&scx_cgroup_ops_rwsem);
3414
3415	if (scx_cgroup_enabled && SCX_HAS_OP(sch, cgroup_set_idle))
3416		SCX_CALL_OP(sch, SCX_KF_UNLOCKED, cgroup_set_idle, NULL,
3417			    tg_cgrp(tg), idle);
3418
3419	/* Update the task group's idle state */
3420	tg->scx.idle = idle;
3421
3422	percpu_up_read(&scx_cgroup_ops_rwsem);
3423}
3424
3425void scx_group_set_bandwidth(struct task_group *tg,
3426			     u64 period_us, u64 quota_us, u64 burst_us)
3427{
3428	struct scx_sched *sch = scx_root;
3429
3430	percpu_down_read(&scx_cgroup_ops_rwsem);
3431
3432	if (scx_cgroup_enabled && SCX_HAS_OP(sch, cgroup_set_bandwidth) &&
3433	    (tg->scx.bw_period_us != period_us ||
3434	     tg->scx.bw_quota_us != quota_us ||
3435	     tg->scx.bw_burst_us != burst_us))
3436		SCX_CALL_OP(sch, SCX_KF_UNLOCKED, cgroup_set_bandwidth, NULL,
3437			    tg_cgrp(tg), period_us, quota_us, burst_us);
3438
3439	tg->scx.bw_period_us = period_us;
3440	tg->scx.bw_quota_us = quota_us;
3441	tg->scx.bw_burst_us = burst_us;
3442
3443	percpu_up_read(&scx_cgroup_ops_rwsem);
3444}
3445
3446static void scx_cgroup_lock(void)
3447{
3448	percpu_down_write(&scx_cgroup_ops_rwsem);
3449	cgroup_lock();
3450}
3451
3452static void scx_cgroup_unlock(void)
3453{
3454	cgroup_unlock();
3455	percpu_up_write(&scx_cgroup_ops_rwsem);
3456}
3457
3458#else	/* CONFIG_EXT_GROUP_SCHED */
3459
3460static void scx_cgroup_lock(void) {}
3461static void scx_cgroup_unlock(void) {}
3462
3463#endif	/* CONFIG_EXT_GROUP_SCHED */
3464
3465/*
3466 * Omitted operations:
3467 *
3468 * - wakeup_preempt: NOOP as it isn't useful in the wakeup path because the task
3469 *   isn't tied to the CPU at that point. Preemption is implemented by resetting
3470 *   the victim task's slice to 0 and triggering reschedule on the target CPU.
3471 *
3472 * - migrate_task_rq: Unnecessary as task to cpu mapping is transient.
3473 *
3474 * - task_fork/dead: We need fork/dead notifications for all tasks regardless of
3475 *   their current sched_class. Call them directly from sched core instead.
3476 */
3477DEFINE_SCHED_CLASS(ext) = {
3478	.queue_mask		= 1,
3479
3480	.enqueue_task		= enqueue_task_scx,
3481	.dequeue_task		= dequeue_task_scx,
3482	.yield_task		= yield_task_scx,
3483	.yield_to_task		= yield_to_task_scx,
3484
3485	.wakeup_preempt		= wakeup_preempt_scx,
3486
3487	.pick_task		= pick_task_scx,
3488
3489	.put_prev_task		= put_prev_task_scx,
3490	.set_next_task		= set_next_task_scx,
3491
3492	.select_task_rq		= select_task_rq_scx,
3493	.task_woken		= task_woken_scx,
3494	.set_cpus_allowed	= set_cpus_allowed_scx,
3495
3496	.rq_online		= rq_online_scx,
3497	.rq_offline		= rq_offline_scx,
3498
3499	.task_tick		= task_tick_scx,
3500
3501	.switching_to		= switching_to_scx,
3502	.switched_from		= switched_from_scx,
3503	.switched_to		= switched_to_scx,
3504	.reweight_task		= reweight_task_scx,
3505	.prio_changed		= prio_changed_scx,
3506
3507	.update_curr		= update_curr_scx,
3508
3509#ifdef CONFIG_UCLAMP_TASK
3510	.uclamp_enabled		= 1,
3511#endif
3512};
3513
3514static void init_dsq(struct scx_dispatch_q *dsq, u64 dsq_id)
3515{
3516	memset(dsq, 0, sizeof(*dsq));
3517
3518	raw_spin_lock_init(&dsq->lock);
3519	INIT_LIST_HEAD(&dsq->list);
3520	dsq->id = dsq_id;
3521}
3522
3523static void free_dsq_irq_workfn(struct irq_work *irq_work)
3524{
3525	struct llist_node *to_free = llist_del_all(&dsqs_to_free);
3526	struct scx_dispatch_q *dsq, *tmp_dsq;
3527
3528	llist_for_each_entry_safe(dsq, tmp_dsq, to_free, free_node)
3529		kfree_rcu(dsq, rcu);
3530}
3531
3532static DEFINE_IRQ_WORK(free_dsq_irq_work, free_dsq_irq_workfn);
3533
3534static void destroy_dsq(struct scx_sched *sch, u64 dsq_id)
3535{
3536	struct scx_dispatch_q *dsq;
3537	unsigned long flags;
3538
3539	rcu_read_lock();
3540
3541	dsq = find_user_dsq(sch, dsq_id);
3542	if (!dsq)
3543		goto out_unlock_rcu;
3544
3545	raw_spin_lock_irqsave(&dsq->lock, flags);
3546
3547	if (dsq->nr) {
3548		scx_error(sch, "attempting to destroy in-use dsq 0x%016llx (nr=%u)",
3549			  dsq->id, dsq->nr);
3550		goto out_unlock_dsq;
3551	}
3552
3553	if (rhashtable_remove_fast(&sch->dsq_hash, &dsq->hash_node,
3554				   dsq_hash_params))
3555		goto out_unlock_dsq;
3556
3557	/*
3558	 * Mark dead by invalidating ->id to prevent dispatch_enqueue() from
3559	 * queueing more tasks. As this function can be called from anywhere,
3560	 * freeing is bounced through an irq work to avoid nesting RCU
3561	 * operations inside scheduler locks.
3562	 */
3563	dsq->id = SCX_DSQ_INVALID;
3564	llist_add(&dsq->free_node, &dsqs_to_free);
3565	irq_work_queue(&free_dsq_irq_work);
3566
3567out_unlock_dsq:
3568	raw_spin_unlock_irqrestore(&dsq->lock, flags);
3569out_unlock_rcu:
3570	rcu_read_unlock();
3571}
3572
3573#ifdef CONFIG_EXT_GROUP_SCHED
3574static void scx_cgroup_exit(struct scx_sched *sch)
3575{
3576	struct cgroup_subsys_state *css;
3577
3578	scx_cgroup_enabled = false;
3579
3580	/*
3581	 * scx_tg_on/offline() are excluded through cgroup_lock(). If we walk
3582	 * cgroups and exit all the inited ones, all online cgroups are exited.
3583	 */
3584	css_for_each_descendant_post(css, &root_task_group.css) {
3585		struct task_group *tg = css_tg(css);
3586
3587		if (!(tg->scx.flags & SCX_TG_INITED))
3588			continue;
3589		tg->scx.flags &= ~SCX_TG_INITED;
3590
3591		if (!sch->ops.cgroup_exit)
3592			continue;
3593
3594		SCX_CALL_OP(sch, SCX_KF_UNLOCKED, cgroup_exit, NULL,
3595			    css->cgroup);
3596	}
3597}
3598
3599static int scx_cgroup_init(struct scx_sched *sch)
3600{
3601	struct cgroup_subsys_state *css;
3602	int ret;
3603
3604	/*
3605	 * scx_tg_on/offline() are excluded through cgroup_lock(). If we walk
3606	 * cgroups and init, all online cgroups are initialized.
3607	 */
3608	css_for_each_descendant_pre(css, &root_task_group.css) {
3609		struct task_group *tg = css_tg(css);
3610		struct scx_cgroup_init_args args = {
3611			.weight = tg->scx.weight,
3612			.bw_period_us = tg->scx.bw_period_us,
3613			.bw_quota_us = tg->scx.bw_quota_us,
3614			.bw_burst_us = tg->scx.bw_burst_us,
3615		};
3616
3617		if ((tg->scx.flags &
3618		     (SCX_TG_ONLINE | SCX_TG_INITED)) != SCX_TG_ONLINE)
3619			continue;
3620
3621		if (!sch->ops.cgroup_init) {
3622			tg->scx.flags |= SCX_TG_INITED;
3623			continue;
3624		}
3625
3626		ret = SCX_CALL_OP_RET(sch, SCX_KF_UNLOCKED, cgroup_init, NULL,
3627				      css->cgroup, &args);
3628		if (ret) {
3629			scx_error(sch, "ops.cgroup_init() failed (%d)", ret);
3630			return ret;
3631		}
3632		tg->scx.flags |= SCX_TG_INITED;
3633	}
3634
3635	WARN_ON_ONCE(scx_cgroup_enabled);
3636	scx_cgroup_enabled = true;
3637
3638	return 0;
3639}
3640
3641#else
3642static void scx_cgroup_exit(struct scx_sched *sch) {}
3643static int scx_cgroup_init(struct scx_sched *sch) { return 0; }
3644#endif
3645
3646
3647/********************************************************************************
3648 * Sysfs interface and ops enable/disable.
3649 */
3650
3651#define SCX_ATTR(_name)								\
3652	static struct kobj_attribute scx_attr_##_name = {			\
3653		.attr = { .name = __stringify(_name), .mode = 0444 },		\
3654		.show = scx_attr_##_name##_show,				\
3655	}
3656
3657static ssize_t scx_attr_state_show(struct kobject *kobj,
3658				   struct kobj_attribute *ka, char *buf)
3659{
3660	return sysfs_emit(buf, "%s\n", scx_enable_state_str[scx_enable_state()]);
3661}
3662SCX_ATTR(state);
3663
3664static ssize_t scx_attr_switch_all_show(struct kobject *kobj,
3665					struct kobj_attribute *ka, char *buf)
3666{
3667	return sysfs_emit(buf, "%d\n", READ_ONCE(scx_switching_all));
3668}
3669SCX_ATTR(switch_all);
3670
3671static ssize_t scx_attr_nr_rejected_show(struct kobject *kobj,
3672					 struct kobj_attribute *ka, char *buf)
3673{
3674	return sysfs_emit(buf, "%ld\n", atomic_long_read(&scx_nr_rejected));
3675}
3676SCX_ATTR(nr_rejected);
3677
3678static ssize_t scx_attr_hotplug_seq_show(struct kobject *kobj,
3679					 struct kobj_attribute *ka, char *buf)
3680{
3681	return sysfs_emit(buf, "%ld\n", atomic_long_read(&scx_hotplug_seq));
3682}
3683SCX_ATTR(hotplug_seq);
3684
3685static ssize_t scx_attr_enable_seq_show(struct kobject *kobj,
3686					struct kobj_attribute *ka, char *buf)
3687{
3688	return sysfs_emit(buf, "%ld\n", atomic_long_read(&scx_enable_seq));
3689}
3690SCX_ATTR(enable_seq);
3691
3692static struct attribute *scx_global_attrs[] = {
3693	&scx_attr_state.attr,
3694	&scx_attr_switch_all.attr,
3695	&scx_attr_nr_rejected.attr,
3696	&scx_attr_hotplug_seq.attr,
3697	&scx_attr_enable_seq.attr,
3698	NULL,
3699};
3700
3701static const struct attribute_group scx_global_attr_group = {
3702	.attrs = scx_global_attrs,
3703};
3704
3705static void free_exit_info(struct scx_exit_info *ei);
3706
3707static void scx_sched_free_rcu_work(struct work_struct *work)
3708{
3709	struct rcu_work *rcu_work = to_rcu_work(work);
3710	struct scx_sched *sch = container_of(rcu_work, struct scx_sched, rcu_work);
3711	struct rhashtable_iter rht_iter;
3712	struct scx_dispatch_q *dsq;
3713	int node;
3714
3715	irq_work_sync(&sch->error_irq_work);
3716	kthread_destroy_worker(sch->helper);
3717
3718	free_percpu(sch->pcpu);
3719
3720	for_each_node_state(node, N_POSSIBLE)
3721		kfree(sch->global_dsqs[node]);
3722	kfree(sch->global_dsqs);
3723
3724	rhashtable_walk_enter(&sch->dsq_hash, &rht_iter);
3725	do {
3726		rhashtable_walk_start(&rht_iter);
3727
3728		while ((dsq = rhashtable_walk_next(&rht_iter)) && !IS_ERR(dsq))
3729			destroy_dsq(sch, dsq->id);
3730
3731		rhashtable_walk_stop(&rht_iter);
3732	} while (dsq == ERR_PTR(-EAGAIN));
3733	rhashtable_walk_exit(&rht_iter);
3734
3735	rhashtable_free_and_destroy(&sch->dsq_hash, NULL, NULL);
3736	free_exit_info(sch->exit_info);
3737	kfree(sch);
3738}
3739
3740static void scx_kobj_release(struct kobject *kobj)
3741{
3742	struct scx_sched *sch = container_of(kobj, struct scx_sched, kobj);
3743
3744	INIT_RCU_WORK(&sch->rcu_work, scx_sched_free_rcu_work);
3745	queue_rcu_work(system_unbound_wq, &sch->rcu_work);
3746}
3747
3748static ssize_t scx_attr_ops_show(struct kobject *kobj,
3749				 struct kobj_attribute *ka, char *buf)
3750{
3751	return sysfs_emit(buf, "%s\n", scx_root->ops.name);
3752}
3753SCX_ATTR(ops);
3754
3755#define scx_attr_event_show(buf, at, events, kind) ({				\
3756	sysfs_emit_at(buf, at, "%s %llu\n", #kind, (events)->kind);		\
3757})
3758
3759static ssize_t scx_attr_events_show(struct kobject *kobj,
3760				    struct kobj_attribute *ka, char *buf)
3761{
3762	struct scx_sched *sch = container_of(kobj, struct scx_sched, kobj);
3763	struct scx_event_stats events;
3764	int at = 0;
3765
3766	scx_read_events(sch, &events);
3767	at += scx_attr_event_show(buf, at, &events, SCX_EV_SELECT_CPU_FALLBACK);
3768	at += scx_attr_event_show(buf, at, &events, SCX_EV_DISPATCH_LOCAL_DSQ_OFFLINE);
3769	at += scx_attr_event_show(buf, at, &events, SCX_EV_DISPATCH_KEEP_LAST);
3770	at += scx_attr_event_show(buf, at, &events, SCX_EV_ENQ_SKIP_EXITING);
3771	at += scx_attr_event_show(buf, at, &events, SCX_EV_ENQ_SKIP_MIGRATION_DISABLED);
3772	at += scx_attr_event_show(buf, at, &events, SCX_EV_REFILL_SLICE_DFL);
3773	at += scx_attr_event_show(buf, at, &events, SCX_EV_BYPASS_DURATION);
3774	at += scx_attr_event_show(buf, at, &events, SCX_EV_BYPASS_DISPATCH);
3775	at += scx_attr_event_show(buf, at, &events, SCX_EV_BYPASS_ACTIVATE);
3776	return at;
3777}
3778SCX_ATTR(events);
3779
3780static struct attribute *scx_sched_attrs[] = {
3781	&scx_attr_ops.attr,
3782	&scx_attr_events.attr,
3783	NULL,
3784};
3785ATTRIBUTE_GROUPS(scx_sched);
3786
3787static const struct kobj_type scx_ktype = {
3788	.release = scx_kobj_release,
3789	.sysfs_ops = &kobj_sysfs_ops,
3790	.default_groups = scx_sched_groups,
3791};
3792
3793static int scx_uevent(const struct kobject *kobj, struct kobj_uevent_env *env)
3794{
3795	return add_uevent_var(env, "SCXOPS=%s", scx_root->ops.name);
3796}
3797
3798static const struct kset_uevent_ops scx_uevent_ops = {
3799	.uevent = scx_uevent,
3800};
3801
3802/*
3803 * Used by sched_fork() and __setscheduler_prio() to pick the matching
3804 * sched_class. dl/rt are already handled.
3805 */
3806bool task_should_scx(int policy)
3807{
3808	if (!scx_enabled() || unlikely(scx_enable_state() == SCX_DISABLING))
3809		return false;
3810	if (READ_ONCE(scx_switching_all))
3811		return true;
3812	return policy == SCHED_EXT;
3813}
3814
3815bool scx_allow_ttwu_queue(const struct task_struct *p)
3816{
3817	struct scx_sched *sch;
3818
3819	if (!scx_enabled())
3820		return true;
3821
3822	sch = rcu_dereference_sched(scx_root);
3823	if (unlikely(!sch))
3824		return true;
3825
3826	if (sch->ops.flags & SCX_OPS_ALLOW_QUEUED_WAKEUP)
3827		return true;
3828
3829	if (unlikely(p->sched_class != &ext_sched_class))
3830		return true;
3831
3832	return false;
3833}
3834
3835/**
3836 * handle_lockup - sched_ext common lockup handler
3837 * @fmt: format string
3838 *
3839 * Called on system stall or lockup condition and initiates abort of sched_ext
3840 * if enabled, which may resolve the reported lockup.
3841 *
3842 * Returns %true if sched_ext is enabled and abort was initiated, which may
3843 * resolve the lockup. %false if sched_ext is not enabled or abort was already
3844 * initiated by someone else.
3845 */
3846static __printf(1, 2) bool handle_lockup(const char *fmt, ...)
3847{
3848	struct scx_sched *sch;
3849	va_list args;
3850	bool ret;
3851
3852	guard(rcu)();
3853
3854	sch = rcu_dereference(scx_root);
3855	if (unlikely(!sch))
3856		return false;
3857
3858	switch (scx_enable_state()) {
3859	case SCX_ENABLING:
3860	case SCX_ENABLED:
3861		va_start(args, fmt);
3862		ret = scx_verror(sch, fmt, args);
3863		va_end(args);
3864		return ret;
3865	default:
3866		return false;
3867	}
3868}
3869
3870/**
3871 * scx_rcu_cpu_stall - sched_ext RCU CPU stall handler
3872 *
3873 * While there are various reasons why RCU CPU stalls can occur on a system
3874 * that may not be caused by the current BPF scheduler, try kicking out the
3875 * current scheduler in an attempt to recover the system to a good state before
3876 * issuing panics.
3877 *
3878 * Returns %true if sched_ext is enabled and abort was initiated, which may
3879 * resolve the reported RCU stall. %false if sched_ext is not enabled or someone
3880 * else already initiated abort.
3881 */
3882bool scx_rcu_cpu_stall(void)
3883{
3884	return handle_lockup("RCU CPU stall detected!");
3885}
3886
3887/**
3888 * scx_softlockup - sched_ext softlockup handler
3889 * @dur_s: number of seconds of CPU stuck due to soft lockup
3890 *
3891 * On some multi-socket setups (e.g. 2x Intel 8480c), the BPF scheduler can
3892 * live-lock the system by making many CPUs target the same DSQ to the point
3893 * where soft-lockup detection triggers. This function is called from
3894 * soft-lockup watchdog when the triggering point is close and tries to unjam
3895 * the system and aborting the BPF scheduler.
3896 */
3897void scx_softlockup(u32 dur_s)
3898{
3899	if (!handle_lockup("soft lockup - CPU %d stuck for %us", smp_processor_id(), dur_s))
3900		return;
3901
3902	printk_deferred(KERN_ERR "sched_ext: Soft lockup - CPU %d stuck for %us, disabling BPF scheduler\n",
3903			smp_processor_id(), dur_s);
3904}
3905
3906/**
3907 * scx_hardlockup - sched_ext hardlockup handler
3908 *
3909 * A poorly behaving BPF scheduler can trigger hard lockup by e.g. putting
3910 * numerous affinitized tasks in a single queue and directing all CPUs at it.
3911 * Try kicking out the current scheduler in an attempt to recover the system to
3912 * a good state before taking more drastic actions.
3913 *
3914 * Returns %true if sched_ext is enabled and abort was initiated, which may
3915 * resolve the reported hardlockdup. %false if sched_ext is not enabled or
3916 * someone else already initiated abort.
3917 */
3918bool scx_hardlockup(int cpu)
3919{
3920	if (!handle_lockup("hard lockup - CPU %d", cpu))
3921		return false;
3922
3923	printk_deferred(KERN_ERR "sched_ext: Hard lockup - CPU %d, disabling BPF scheduler\n",
3924			cpu);
3925	return true;
3926}
3927
3928static u32 bypass_lb_cpu(struct scx_sched *sch, struct rq *rq,
3929			 struct cpumask *donee_mask, struct cpumask *resched_mask,
3930			 u32 nr_donor_target, u32 nr_donee_target)
3931{
3932	struct scx_dispatch_q *donor_dsq = &rq->scx.bypass_dsq;
3933	struct task_struct *p, *n;
3934	struct scx_dsq_list_node cursor = INIT_DSQ_LIST_CURSOR(cursor, 0, 0);
3935	s32 delta = READ_ONCE(donor_dsq->nr) - nr_donor_target;
3936	u32 nr_balanced = 0, min_delta_us;
3937
3938	/*
3939	 * All we want to guarantee is reasonable forward progress. No reason to
3940	 * fine tune. Assuming every task on @donor_dsq runs their full slice,
3941	 * consider offloading iff the total queued duration is over the
3942	 * threshold.
3943	 */
3944	min_delta_us = scx_bypass_lb_intv_us / SCX_BYPASS_LB_MIN_DELTA_DIV;
3945	if (delta < DIV_ROUND_UP(min_delta_us, scx_slice_bypass_us))
3946		return 0;
3947
3948	raw_spin_rq_lock_irq(rq);
3949	raw_spin_lock(&donor_dsq->lock);
3950	list_add(&cursor.node, &donor_dsq->list);
3951resume:
3952	n = container_of(&cursor, struct task_struct, scx.dsq_list);
3953	n = nldsq_next_task(donor_dsq, n, false);
3954
3955	while ((p = n)) {
3956		struct rq *donee_rq;
3957		struct scx_dispatch_q *donee_dsq;
3958		int donee;
3959
3960		n = nldsq_next_task(donor_dsq, n, false);
3961
3962		if (donor_dsq->nr <= nr_donor_target)
3963			break;
3964
3965		if (cpumask_empty(donee_mask))
3966			break;
3967
3968		donee = cpumask_any_and_distribute(donee_mask, p->cpus_ptr);
3969		if (donee >= nr_cpu_ids)
3970			continue;
3971
3972		donee_rq = cpu_rq(donee);
3973		donee_dsq = &donee_rq->scx.bypass_dsq;
3974
3975		/*
3976		 * $p's rq is not locked but $p's DSQ lock protects its
3977		 * scheduling properties making this test safe.
3978		 */
3979		if (!task_can_run_on_remote_rq(sch, p, donee_rq, false))
3980			continue;
3981
3982		/*
3983		 * Moving $p from one non-local DSQ to another. The source rq
3984		 * and DSQ are already locked. Do an abbreviated dequeue and
3985		 * then perform enqueue without unlocking $donor_dsq.
3986		 *
3987		 * We don't want to drop and reacquire the lock on each
3988		 * iteration as @donor_dsq can be very long and potentially
3989		 * highly contended. Donee DSQs are less likely to be contended.
3990		 * The nested locking is safe as only this LB moves tasks
3991		 * between bypass DSQs.
3992		 */
3993		dispatch_dequeue_locked(p, donor_dsq);
3994		dispatch_enqueue(sch, donee_dsq, p, SCX_ENQ_NESTED);
3995
3996		/*
3997		 * $donee might have been idle and need to be woken up. No need
3998		 * to be clever. Kick every CPU that receives tasks.
3999		 */
4000		cpumask_set_cpu(donee, resched_mask);
4001
4002		if (READ_ONCE(donee_dsq->nr) >= nr_donee_target)
4003			cpumask_clear_cpu(donee, donee_mask);
4004
4005		nr_balanced++;
4006		if (!(nr_balanced % SCX_BYPASS_LB_BATCH) && n) {
4007			list_move_tail(&cursor.node, &n->scx.dsq_list.node);
4008			raw_spin_unlock(&donor_dsq->lock);
4009			raw_spin_rq_unlock_irq(rq);
4010			cpu_relax();
4011			raw_spin_rq_lock_irq(rq);
4012			raw_spin_lock(&donor_dsq->lock);
4013			goto resume;
4014		}
4015	}
4016
4017	list_del_init(&cursor.node);
4018	raw_spin_unlock(&donor_dsq->lock);
4019	raw_spin_rq_unlock_irq(rq);
4020
4021	return nr_balanced;
4022}
4023
4024static void bypass_lb_node(struct scx_sched *sch, int node)
4025{
4026	const struct cpumask *node_mask = cpumask_of_node(node);
4027	struct cpumask *donee_mask = scx_bypass_lb_donee_cpumask;
4028	struct cpumask *resched_mask = scx_bypass_lb_resched_cpumask;
4029	u32 nr_tasks = 0, nr_cpus = 0, nr_balanced = 0;
4030	u32 nr_target, nr_donor_target;
4031	u32 before_min = U32_MAX, before_max = 0;
4032	u32 after_min = U32_MAX, after_max = 0;
4033	int cpu;
4034
4035	/* count the target tasks and CPUs */
4036	for_each_cpu_and(cpu, cpu_online_mask, node_mask) {
4037		u32 nr = READ_ONCE(cpu_rq(cpu)->scx.bypass_dsq.nr);
4038
4039		nr_tasks += nr;
4040		nr_cpus++;
4041
4042		before_min = min(nr, before_min);
4043		before_max = max(nr, before_max);
4044	}
4045
4046	if (!nr_cpus)
4047		return;
4048
4049	/*
4050	 * We don't want CPUs to have more than $nr_donor_target tasks and
4051	 * balancing to fill donee CPUs upto $nr_target. Once targets are
4052	 * calculated, find the donee CPUs.
4053	 */
4054	nr_target = DIV_ROUND_UP(nr_tasks, nr_cpus);
4055	nr_donor_target = DIV_ROUND_UP(nr_target * SCX_BYPASS_LB_DONOR_PCT, 100);
4056
4057	cpumask_clear(donee_mask);
4058	for_each_cpu_and(cpu, cpu_online_mask, node_mask) {
4059		if (READ_ONCE(cpu_rq(cpu)->scx.bypass_dsq.nr) < nr_target)
4060			cpumask_set_cpu(cpu, donee_mask);
4061	}
4062
4063	/* iterate !donee CPUs and see if they should be offloaded */
4064	cpumask_clear(resched_mask);
4065	for_each_cpu_and(cpu, cpu_online_mask, node_mask) {
4066		struct rq *rq = cpu_rq(cpu);
4067		struct scx_dispatch_q *donor_dsq = &rq->scx.bypass_dsq;
4068
4069		if (cpumask_empty(donee_mask))
4070			break;
4071		if (cpumask_test_cpu(cpu, donee_mask))
4072			continue;
4073		if (READ_ONCE(donor_dsq->nr) <= nr_donor_target)
4074			continue;
4075
4076		nr_balanced += bypass_lb_cpu(sch, rq, donee_mask, resched_mask,
4077					     nr_donor_target, nr_target);
4078	}
4079
4080	for_each_cpu(cpu, resched_mask)
4081		resched_cpu(cpu);
4082
4083	for_each_cpu_and(cpu, cpu_online_mask, node_mask) {
4084		u32 nr = READ_ONCE(cpu_rq(cpu)->scx.bypass_dsq.nr);
4085
4086		after_min = min(nr, after_min);
4087		after_max = max(nr, after_max);
4088
4089	}
4090
4091	trace_sched_ext_bypass_lb(node, nr_cpus, nr_tasks, nr_balanced,
4092				  before_min, before_max, after_min, after_max);
4093}
4094
4095/*
4096 * In bypass mode, all tasks are put on the per-CPU bypass DSQs. If the machine
4097 * is over-saturated and the BPF scheduler skewed tasks into few CPUs, some
4098 * bypass DSQs can be overloaded. If there are enough tasks to saturate other
4099 * lightly loaded CPUs, such imbalance can lead to very high execution latency
4100 * on the overloaded CPUs and thus to hung tasks and RCU stalls. To avoid such
4101 * outcomes, a simple load balancing mechanism is implemented by the following
4102 * timer which runs periodically while bypass mode is in effect.
4103 */
4104static void scx_bypass_lb_timerfn(struct timer_list *timer)
4105{
4106	struct scx_sched *sch;
4107	int node;
4108	u32 intv_us;
4109
4110	sch = rcu_dereference_all(scx_root);
4111	if (unlikely(!sch) || !READ_ONCE(scx_bypass_depth))
4112		return;
4113
4114	for_each_node_with_cpus(node)
4115		bypass_lb_node(sch, node);
4116
4117	intv_us = READ_ONCE(scx_bypass_lb_intv_us);
4118	if (intv_us)
4119		mod_timer(timer, jiffies + usecs_to_jiffies(intv_us));
4120}
4121
4122static DEFINE_TIMER(scx_bypass_lb_timer, scx_bypass_lb_timerfn);
4123
4124/**
4125 * scx_bypass - [Un]bypass scx_ops and guarantee forward progress
4126 * @bypass: true for bypass, false for unbypass
4127 *
4128 * Bypassing guarantees that all runnable tasks make forward progress without
4129 * trusting the BPF scheduler. We can't grab any mutexes or rwsems as they might
4130 * be held by tasks that the BPF scheduler is forgetting to run, which
4131 * unfortunately also excludes toggling the static branches.
4132 *
4133 * Let's work around by overriding a couple ops and modifying behaviors based on
4134 * the DISABLING state and then cycling the queued tasks through dequeue/enqueue
4135 * to force global FIFO scheduling.
4136 *
4137 * - ops.select_cpu() is ignored and the default select_cpu() is used.
4138 *
4139 * - ops.enqueue() is ignored and tasks are queued in simple global FIFO order.
4140 *   %SCX_OPS_ENQ_LAST is also ignored.
4141 *
4142 * - ops.dispatch() is ignored.
4143 *
4144 * - balance_one() does not set %SCX_RQ_BAL_KEEP on non-zero slice as slice
4145 *   can't be trusted. Whenever a tick triggers, the running task is rotated to
4146 *   the tail of the queue with core_sched_at touched.
4147 *
4148 * - pick_next_task() suppresses zero slice warning.
4149 *
4150 * - scx_kick_cpu() is disabled to avoid irq_work malfunction during PM
4151 *   operations.
4152 *
4153 * - scx_prio_less() reverts to the default core_sched_at order.
4154 */
4155static void scx_bypass(bool bypass)
4156{
4157	static DEFINE_RAW_SPINLOCK(bypass_lock);
4158	static unsigned long bypass_timestamp;
4159	struct scx_sched *sch;
4160	unsigned long flags;
4161	int cpu;
4162
4163	raw_spin_lock_irqsave(&bypass_lock, flags);
4164	sch = rcu_dereference_bh(scx_root);
4165
4166	if (bypass) {
4167		u32 intv_us;
4168
4169		WRITE_ONCE(scx_bypass_depth, scx_bypass_depth + 1);
4170		WARN_ON_ONCE(scx_bypass_depth <= 0);
4171		if (scx_bypass_depth != 1)
4172			goto unlock;
4173		WRITE_ONCE(scx_slice_dfl, scx_slice_bypass_us * NSEC_PER_USEC);
4174		bypass_timestamp = ktime_get_ns();
4175		if (sch)
4176			scx_add_event(sch, SCX_EV_BYPASS_ACTIVATE, 1);
4177
4178		intv_us = READ_ONCE(scx_bypass_lb_intv_us);
4179		if (intv_us && !timer_pending(&scx_bypass_lb_timer)) {
4180			scx_bypass_lb_timer.expires =
4181				jiffies + usecs_to_jiffies(intv_us);
4182			add_timer_global(&scx_bypass_lb_timer);
4183		}
4184	} else {
4185		WRITE_ONCE(scx_bypass_depth, scx_bypass_depth - 1);
4186		WARN_ON_ONCE(scx_bypass_depth < 0);
4187		if (scx_bypass_depth != 0)
4188			goto unlock;
4189		WRITE_ONCE(scx_slice_dfl, SCX_SLICE_DFL);
4190		if (sch)
4191			scx_add_event(sch, SCX_EV_BYPASS_DURATION,
4192				      ktime_get_ns() - bypass_timestamp);
4193	}
4194
4195	/*
4196	 * No task property is changing. We just need to make sure all currently
4197	 * queued tasks are re-queued according to the new scx_rq_bypassing()
4198	 * state. As an optimization, walk each rq's runnable_list instead of
4199	 * the scx_tasks list.
4200	 *
4201	 * This function can't trust the scheduler and thus can't use
4202	 * cpus_read_lock(). Walk all possible CPUs instead of online.
4203	 */
4204	for_each_possible_cpu(cpu) {
4205		struct rq *rq = cpu_rq(cpu);
4206		struct task_struct *p, *n;
4207
4208		raw_spin_rq_lock(rq);
4209
4210		if (bypass) {
4211			WARN_ON_ONCE(rq->scx.flags & SCX_RQ_BYPASSING);
4212			rq->scx.flags |= SCX_RQ_BYPASSING;
4213		} else {
4214			WARN_ON_ONCE(!(rq->scx.flags & SCX_RQ_BYPASSING));
4215			rq->scx.flags &= ~SCX_RQ_BYPASSING;
4216		}
4217
4218		/*
4219		 * We need to guarantee that no tasks are on the BPF scheduler
4220		 * while bypassing. Either we see enabled or the enable path
4221		 * sees scx_rq_bypassing() before moving tasks to SCX.
4222		 */
4223		if (!scx_enabled()) {
4224			raw_spin_rq_unlock(rq);
4225			continue;
4226		}
4227
4228		/*
4229		 * The use of list_for_each_entry_safe_reverse() is required
4230		 * because each task is going to be removed from and added back
4231		 * to the runnable_list during iteration. Because they're added
4232		 * to the tail of the list, safe reverse iteration can still
4233		 * visit all nodes.
4234		 */
4235		list_for_each_entry_safe_reverse(p, n, &rq->scx.runnable_list,
4236						 scx.runnable_node) {
4237			/* cycling deq/enq is enough, see the function comment */
4238			scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
4239				/* nothing */ ;
4240			}
4241		}
4242
4243		/* resched to restore ticks and idle state */
4244		if (cpu_online(cpu) || cpu == smp_processor_id())
4245			resched_curr(rq);
4246
4247		raw_spin_rq_unlock(rq);
4248	}
4249
4250unlock:
4251	raw_spin_unlock_irqrestore(&bypass_lock, flags);
4252}
4253
4254static void free_exit_info(struct scx_exit_info *ei)
4255{
4256	kvfree(ei->dump);
4257	kfree(ei->msg);
4258	kfree(ei->bt);
4259	kfree(ei);
4260}
4261
4262static struct scx_exit_info *alloc_exit_info(size_t exit_dump_len)
4263{
4264	struct scx_exit_info *ei;
4265
4266	ei = kzalloc(sizeof(*ei), GFP_KERNEL);
4267	if (!ei)
4268		return NULL;
4269
4270	ei->bt = kcalloc(SCX_EXIT_BT_LEN, sizeof(ei->bt[0]), GFP_KERNEL);
4271	ei->msg = kzalloc(SCX_EXIT_MSG_LEN, GFP_KERNEL);
4272	ei->dump = kvzalloc(exit_dump_len, GFP_KERNEL);
4273
4274	if (!ei->bt || !ei->msg || !ei->dump) {
4275		free_exit_info(ei);
4276		return NULL;
4277	}
4278
4279	return ei;
4280}
4281
4282static const char *scx_exit_reason(enum scx_exit_kind kind)
4283{
4284	switch (kind) {
4285	case SCX_EXIT_UNREG:
4286		return "unregistered from user space";
4287	case SCX_EXIT_UNREG_BPF:
4288		return "unregistered from BPF";
4289	case SCX_EXIT_UNREG_KERN:
4290		return "unregistered from the main kernel";
4291	case SCX_EXIT_SYSRQ:
4292		return "disabled by sysrq-S";
4293	case SCX_EXIT_ERROR:
4294		return "runtime error";
4295	case SCX_EXIT_ERROR_BPF:
4296		return "scx_bpf_error";
4297	case SCX_EXIT_ERROR_STALL:
4298		return "runnable task stall";
4299	default:
4300		return "<UNKNOWN>";
4301	}
4302}
4303
4304static void free_kick_syncs(void)
4305{
4306	int cpu;
4307
4308	for_each_possible_cpu(cpu) {
4309		struct scx_kick_syncs **ksyncs = per_cpu_ptr(&scx_kick_syncs, cpu);
4310		struct scx_kick_syncs *to_free;
4311
4312		to_free = rcu_replace_pointer(*ksyncs, NULL, true);
4313		if (to_free)
4314			kvfree_rcu(to_free, rcu);
4315	}
4316}
4317
4318static void scx_disable_workfn(struct kthread_work *work)
4319{
4320	struct scx_sched *sch = container_of(work, struct scx_sched, disable_work);
4321	struct scx_exit_info *ei = sch->exit_info;
4322	struct scx_task_iter sti;
4323	struct task_struct *p;
4324	int kind, cpu;
4325
4326	kind = atomic_read(&sch->exit_kind);
4327	while (true) {
4328		if (kind == SCX_EXIT_DONE)	/* already disabled? */
4329			return;
4330		WARN_ON_ONCE(kind == SCX_EXIT_NONE);
4331		if (atomic_try_cmpxchg(&sch->exit_kind, &kind, SCX_EXIT_DONE))
4332			break;
4333	}
4334	ei->kind = kind;
4335	ei->reason = scx_exit_reason(ei->kind);
4336
4337	/* guarantee forward progress by bypassing scx_ops */
4338	scx_bypass(true);
4339	WRITE_ONCE(scx_aborting, false);
4340
4341	switch (scx_set_enable_state(SCX_DISABLING)) {
4342	case SCX_DISABLING:
4343		WARN_ONCE(true, "sched_ext: duplicate disabling instance?");
4344		break;
4345	case SCX_DISABLED:
4346		pr_warn("sched_ext: ops error detected without ops (%s)\n",
4347			sch->exit_info->msg);
4348		WARN_ON_ONCE(scx_set_enable_state(SCX_DISABLED) != SCX_DISABLING);
4349		goto done;
4350	default:
4351		break;
4352	}
4353
4354	/*
4355	 * Here, every runnable task is guaranteed to make forward progress and
4356	 * we can safely use blocking synchronization constructs. Actually
4357	 * disable ops.
4358	 */
4359	mutex_lock(&scx_enable_mutex);
4360
4361	static_branch_disable(&__scx_switched_all);
4362	WRITE_ONCE(scx_switching_all, false);
4363
4364	/*
4365	 * Shut down cgroup support before tasks so that the cgroup attach path
4366	 * doesn't race against scx_exit_task().
4367	 */
4368	scx_cgroup_lock();
4369	scx_cgroup_exit(sch);
4370	scx_cgroup_unlock();
4371
4372	/*
4373	 * The BPF scheduler is going away. All tasks including %TASK_DEAD ones
4374	 * must be switched out and exited synchronously.
4375	 */
4376	percpu_down_write(&scx_fork_rwsem);
4377
4378	scx_init_task_enabled = false;
4379
4380	scx_task_iter_start(&sti);
4381	while ((p = scx_task_iter_next_locked(&sti))) {
4382		unsigned int queue_flags = DEQUEUE_SAVE | DEQUEUE_MOVE | DEQUEUE_NOCLOCK;
4383		const struct sched_class *old_class = p->sched_class;
4384		const struct sched_class *new_class = scx_setscheduler_class(p);
4385
4386		update_rq_clock(task_rq(p));
4387
4388		if (old_class != new_class)
4389			queue_flags |= DEQUEUE_CLASS;
4390
4391		scoped_guard (sched_change, p, queue_flags) {
4392			p->sched_class = new_class;
4393		}
4394
4395		scx_exit_task(p);
4396	}
4397	scx_task_iter_stop(&sti);
4398	percpu_up_write(&scx_fork_rwsem);
4399
4400	/*
4401	 * Invalidate all the rq clocks to prevent getting outdated
4402	 * rq clocks from a previous scx scheduler.
4403	 */
4404	for_each_possible_cpu(cpu) {
4405		struct rq *rq = cpu_rq(cpu);
4406		scx_rq_clock_invalidate(rq);
4407	}
4408
4409	/* no task is on scx, turn off all the switches and flush in-progress calls */
4410	static_branch_disable(&__scx_enabled);
4411	bitmap_zero(sch->has_op, SCX_OPI_END);
4412	scx_idle_disable();
4413	synchronize_rcu();
4414
4415	if (ei->kind >= SCX_EXIT_ERROR) {
4416		pr_err("sched_ext: BPF scheduler \"%s\" disabled (%s)\n",
4417		       sch->ops.name, ei->reason);
4418
4419		if (ei->msg[0] != '\0')
4420			pr_err("sched_ext: %s: %s\n", sch->ops.name, ei->msg);
4421#ifdef CONFIG_STACKTRACE
4422		stack_trace_print(ei->bt, ei->bt_len, 2);
4423#endif
4424	} else {
4425		pr_info("sched_ext: BPF scheduler \"%s\" disabled (%s)\n",
4426			sch->ops.name, ei->reason);
4427	}
4428
4429	if (sch->ops.exit)
4430		SCX_CALL_OP(sch, SCX_KF_UNLOCKED, exit, NULL, ei);
4431
4432	cancel_delayed_work_sync(&scx_watchdog_work);
4433
4434	/*
4435	 * scx_root clearing must be inside cpus_read_lock(). See
4436	 * handle_hotplug().
4437	 */
4438	cpus_read_lock();
4439	RCU_INIT_POINTER(scx_root, NULL);
4440	cpus_read_unlock();
4441
4442	/*
4443	 * Delete the kobject from the hierarchy synchronously. Otherwise, sysfs
4444	 * could observe an object of the same name still in the hierarchy when
4445	 * the next scheduler is loaded.
4446	 */
4447	kobject_del(&sch->kobj);
4448
4449	free_percpu(scx_dsp_ctx);
4450	scx_dsp_ctx = NULL;
4451	scx_dsp_max_batch = 0;
4452	free_kick_syncs();
4453
4454	if (scx_bypassed_for_enable) {
4455		scx_bypassed_for_enable = false;
4456		scx_bypass(false);
4457	}
4458
4459	mutex_unlock(&scx_enable_mutex);
4460
4461	WARN_ON_ONCE(scx_set_enable_state(SCX_DISABLED) != SCX_DISABLING);
4462done:
4463	scx_bypass(false);
4464}
4465
4466/*
4467 * Claim the exit on @sch. The caller must ensure that the helper kthread work
4468 * is kicked before the current task can be preempted. Once exit_kind is
4469 * claimed, scx_error() can no longer trigger, so if the current task gets
4470 * preempted and the BPF scheduler fails to schedule it back, the helper work
4471 * will never be kicked and the whole system can wedge.
4472 */
4473static bool scx_claim_exit(struct scx_sched *sch, enum scx_exit_kind kind)
4474{
4475	int none = SCX_EXIT_NONE;
4476
4477	lockdep_assert_preemption_disabled();
4478
4479	if (!atomic_try_cmpxchg(&sch->exit_kind, &none, kind))
4480		return false;
4481
4482	/*
4483	 * Some CPUs may be trapped in the dispatch paths. Set the aborting
4484	 * flag to break potential live-lock scenarios, ensuring we can
4485	 * successfully reach scx_bypass().
4486	 */
4487	WRITE_ONCE(scx_aborting, true);
4488	return true;
4489}
4490
4491static void scx_disable(enum scx_exit_kind kind)
4492{
4493	struct scx_sched *sch;
4494
4495	if (WARN_ON_ONCE(kind == SCX_EXIT_NONE || kind == SCX_EXIT_DONE))
4496		kind = SCX_EXIT_ERROR;
4497
4498	rcu_read_lock();
4499	sch = rcu_dereference(scx_root);
4500	if (sch) {
4501		guard(preempt)();
4502		scx_claim_exit(sch, kind);
4503		kthread_queue_work(sch->helper, &sch->disable_work);
4504	}
4505	rcu_read_unlock();
4506}
4507
4508static void dump_newline(struct seq_buf *s)
4509{
4510	trace_sched_ext_dump("");
4511
4512	/* @s may be zero sized and seq_buf triggers WARN if so */
4513	if (s->size)
4514		seq_buf_putc(s, '\n');
4515}
4516
4517static __printf(2, 3) void dump_line(struct seq_buf *s, const char *fmt, ...)
4518{
4519	va_list args;
4520
4521#ifdef CONFIG_TRACEPOINTS
4522	if (trace_sched_ext_dump_enabled()) {
4523		/* protected by scx_dump_state()::dump_lock */
4524		static char line_buf[SCX_EXIT_MSG_LEN];
4525
4526		va_start(args, fmt);
4527		vscnprintf(line_buf, sizeof(line_buf), fmt, args);
4528		va_end(args);
4529
4530		trace_sched_ext_dump(line_buf);
4531	}
4532#endif
4533	/* @s may be zero sized and seq_buf triggers WARN if so */
4534	if (s->size) {
4535		va_start(args, fmt);
4536		seq_buf_vprintf(s, fmt, args);
4537		va_end(args);
4538
4539		seq_buf_putc(s, '\n');
4540	}
4541}
4542
4543static void dump_stack_trace(struct seq_buf *s, const char *prefix,
4544			     const unsigned long *bt, unsigned int len)
4545{
4546	unsigned int i;
4547
4548	for (i = 0; i < len; i++)
4549		dump_line(s, "%s%pS", prefix, (void *)bt[i]);
4550}
4551
4552static void ops_dump_init(struct seq_buf *s, const char *prefix)
4553{
4554	struct scx_dump_data *dd = &scx_dump_data;
4555
4556	lockdep_assert_irqs_disabled();
4557
4558	dd->cpu = smp_processor_id();		/* allow scx_bpf_dump() */
4559	dd->first = true;
4560	dd->cursor = 0;
4561	dd->s = s;
4562	dd->prefix = prefix;
4563}
4564
4565static void ops_dump_flush(void)
4566{
4567	struct scx_dump_data *dd = &scx_dump_data;
4568	char *line = dd->buf.line;
4569
4570	if (!dd->cursor)
4571		return;
4572
4573	/*
4574	 * There's something to flush and this is the first line. Insert a blank
4575	 * line to distinguish ops dump.
4576	 */
4577	if (dd->first) {
4578		dump_newline(dd->s);
4579		dd->first = false;
4580	}
4581
4582	/*
4583	 * There may be multiple lines in $line. Scan and emit each line
4584	 * separately.
4585	 */
4586	while (true) {
4587		char *end = line;
4588		char c;
4589
4590		while (*end != '\n' && *end != '\0')
4591			end++;
4592
4593		/*
4594		 * If $line overflowed, it may not have newline at the end.
4595		 * Always emit with a newline.
4596		 */
4597		c = *end;
4598		*end = '\0';
4599		dump_line(dd->s, "%s%s", dd->prefix, line);
4600		if (c == '\0')
4601			break;
4602
4603		/* move to the next line */
4604		end++;
4605		if (*end == '\0')
4606			break;
4607		line = end;
4608	}
4609
4610	dd->cursor = 0;
4611}
4612
4613static void ops_dump_exit(void)
4614{
4615	ops_dump_flush();
4616	scx_dump_data.cpu = -1;
4617}
4618
4619static void scx_dump_task(struct seq_buf *s, struct scx_dump_ctx *dctx,
4620			  struct task_struct *p, char marker)
4621{
4622	static unsigned long bt[SCX_EXIT_BT_LEN];
4623	struct scx_sched *sch = scx_root;
4624	char dsq_id_buf[19] = "(n/a)";
4625	unsigned long ops_state = atomic_long_read(&p->scx.ops_state);
4626	unsigned int bt_len = 0;
4627
4628	if (p->scx.dsq)
4629		scnprintf(dsq_id_buf, sizeof(dsq_id_buf), "0x%llx",
4630			  (unsigned long long)p->scx.dsq->id);
4631
4632	dump_newline(s);
4633	dump_line(s, " %c%c %s[%d] %+ldms",
4634		  marker, task_state_to_char(p), p->comm, p->pid,
4635		  jiffies_delta_msecs(p->scx.runnable_at, dctx->at_jiffies));
4636	dump_line(s, "      scx_state/flags=%u/0x%x dsq_flags=0x%x ops_state/qseq=%lu/%lu",
4637		  scx_get_task_state(p), p->scx.flags & ~SCX_TASK_STATE_MASK,
4638		  p->scx.dsq_flags, ops_state & SCX_OPSS_STATE_MASK,
4639		  ops_state >> SCX_OPSS_QSEQ_SHIFT);
4640	dump_line(s, "      sticky/holding_cpu=%d/%d dsq_id=%s",
4641		  p->scx.sticky_cpu, p->scx.holding_cpu, dsq_id_buf);
4642	dump_line(s, "      dsq_vtime=%llu slice=%llu weight=%u",
4643		  p->scx.dsq_vtime, p->scx.slice, p->scx.weight);
4644	dump_line(s, "      cpus=%*pb no_mig=%u", cpumask_pr_args(p->cpus_ptr),
4645		  p->migration_disabled);
4646
4647	if (SCX_HAS_OP(sch, dump_task)) {
4648		ops_dump_init(s, "    ");
4649		SCX_CALL_OP(sch, SCX_KF_REST, dump_task, NULL, dctx, p);
4650		ops_dump_exit();
4651	}
4652
4653#ifdef CONFIG_STACKTRACE
4654	bt_len = stack_trace_save_tsk(p, bt, SCX_EXIT_BT_LEN, 1);
4655#endif
4656	if (bt_len) {
4657		dump_newline(s);
4658		dump_stack_trace(s, "    ", bt, bt_len);
4659	}
4660}
4661
4662static void scx_dump_state(struct scx_exit_info *ei, size_t dump_len)
4663{
4664	static DEFINE_SPINLOCK(dump_lock);
4665	static const char trunc_marker[] = "\n\n~~~~ TRUNCATED ~~~~\n";
4666	struct scx_sched *sch = scx_root;
4667	struct scx_dump_ctx dctx = {
4668		.kind = ei->kind,
4669		.exit_code = ei->exit_code,
4670		.reason = ei->reason,
4671		.at_ns = ktime_get_ns(),
4672		.at_jiffies = jiffies,
4673	};
4674	struct seq_buf s;
4675	struct scx_event_stats events;
4676	unsigned long flags;
4677	char *buf;
4678	int cpu;
4679
4680	spin_lock_irqsave(&dump_lock, flags);
4681
4682	seq_buf_init(&s, ei->dump, dump_len);
4683
4684	if (ei->kind == SCX_EXIT_NONE) {
4685		dump_line(&s, "Debug dump triggered by %s", ei->reason);
4686	} else {
4687		dump_line(&s, "%s[%d] triggered exit kind %d:",
4688			  current->comm, current->pid, ei->kind);
4689		dump_line(&s, "  %s (%s)", ei->reason, ei->msg);
4690		dump_newline(&s);
4691		dump_line(&s, "Backtrace:");
4692		dump_stack_trace(&s, "  ", ei->bt, ei->bt_len);
4693	}
4694
4695	if (SCX_HAS_OP(sch, dump)) {
4696		ops_dump_init(&s, "");
4697		SCX_CALL_OP(sch, SCX_KF_UNLOCKED, dump, NULL, &dctx);
4698		ops_dump_exit();
4699	}
4700
4701	dump_newline(&s);
4702	dump_line(&s, "CPU states");
4703	dump_line(&s, "----------");
4704
4705	for_each_possible_cpu(cpu) {
4706		struct rq *rq = cpu_rq(cpu);
4707		struct rq_flags rf;
4708		struct task_struct *p;
4709		struct seq_buf ns;
4710		size_t avail, used;
4711		bool idle;
4712
4713		rq_lock_irqsave(rq, &rf);
4714
4715		idle = list_empty(&rq->scx.runnable_list) &&
4716			rq->curr->sched_class == &idle_sched_class;
4717
4718		if (idle && !SCX_HAS_OP(sch, dump_cpu))
4719			goto next;
4720
4721		/*
4722		 * We don't yet know whether ops.dump_cpu() will produce output
4723		 * and we may want to skip the default CPU dump if it doesn't.
4724		 * Use a nested seq_buf to generate the standard dump so that we
4725		 * can decide whether to commit later.
4726		 */
4727		avail = seq_buf_get_buf(&s, &buf);
4728		seq_buf_init(&ns, buf, avail);
4729
4730		dump_newline(&ns);
4731		dump_line(&ns, "CPU %-4d: nr_run=%u flags=0x%x cpu_rel=%d ops_qseq=%lu ksync=%lu",
4732			  cpu, rq->scx.nr_running, rq->scx.flags,
4733			  rq->scx.cpu_released, rq->scx.ops_qseq,
4734			  rq->scx.kick_sync);
4735		dump_line(&ns, "          curr=%s[%d] class=%ps",
4736			  rq->curr->comm, rq->curr->pid,
4737			  rq->curr->sched_class);
4738		if (!cpumask_empty(rq->scx.cpus_to_kick))
4739			dump_line(&ns, "  cpus_to_kick   : %*pb",
4740				  cpumask_pr_args(rq->scx.cpus_to_kick));
4741		if (!cpumask_empty(rq->scx.cpus_to_kick_if_idle))
4742			dump_line(&ns, "  idle_to_kick   : %*pb",
4743				  cpumask_pr_args(rq->scx.cpus_to_kick_if_idle));
4744		if (!cpumask_empty(rq->scx.cpus_to_preempt))
4745			dump_line(&ns, "  cpus_to_preempt: %*pb",
4746				  cpumask_pr_args(rq->scx.cpus_to_preempt));
4747		if (!cpumask_empty(rq->scx.cpus_to_wait))
4748			dump_line(&ns, "  cpus_to_wait   : %*pb",
4749				  cpumask_pr_args(rq->scx.cpus_to_wait));
4750		if (!cpumask_empty(rq->scx.cpus_to_sync))
4751			dump_line(&ns, "  cpus_to_sync   : %*pb",
4752				  cpumask_pr_args(rq->scx.cpus_to_sync));
4753
4754		used = seq_buf_used(&ns);
4755		if (SCX_HAS_OP(sch, dump_cpu)) {
4756			ops_dump_init(&ns, "  ");
4757			SCX_CALL_OP(sch, SCX_KF_REST, dump_cpu, NULL,
4758				    &dctx, cpu, idle);
4759			ops_dump_exit();
4760		}
4761
4762		/*
4763		 * If idle && nothing generated by ops.dump_cpu(), there's
4764		 * nothing interesting. Skip.
4765		 */
4766		if (idle && used == seq_buf_used(&ns))
4767			goto next;
4768
4769		/*
4770		 * $s may already have overflowed when $ns was created. If so,
4771		 * calling commit on it will trigger BUG.
4772		 */
4773		if (avail) {
4774			seq_buf_commit(&s, seq_buf_used(&ns));
4775			if (seq_buf_has_overflowed(&ns))
4776				seq_buf_set_overflow(&s);
4777		}
4778
4779		if (rq->curr->sched_class == &ext_sched_class)
4780			scx_dump_task(&s, &dctx, rq->curr, '*');
4781
4782		list_for_each_entry(p, &rq->scx.runnable_list, scx.runnable_node)
4783			scx_dump_task(&s, &dctx, p, ' ');
4784	next:
4785		rq_unlock_irqrestore(rq, &rf);
4786	}
4787
4788	dump_newline(&s);
4789	dump_line(&s, "Event counters");
4790	dump_line(&s, "--------------");
4791
4792	scx_read_events(sch, &events);
4793	scx_dump_event(s, &events, SCX_EV_SELECT_CPU_FALLBACK);
4794	scx_dump_event(s, &events, SCX_EV_DISPATCH_LOCAL_DSQ_OFFLINE);
4795	scx_dump_event(s, &events, SCX_EV_DISPATCH_KEEP_LAST);
4796	scx_dump_event(s, &events, SCX_EV_ENQ_SKIP_EXITING);
4797	scx_dump_event(s, &events, SCX_EV_ENQ_SKIP_MIGRATION_DISABLED);
4798	scx_dump_event(s, &events, SCX_EV_REFILL_SLICE_DFL);
4799	scx_dump_event(s, &events, SCX_EV_BYPASS_DURATION);
4800	scx_dump_event(s, &events, SCX_EV_BYPASS_DISPATCH);
4801	scx_dump_event(s, &events, SCX_EV_BYPASS_ACTIVATE);
4802
4803	if (seq_buf_has_overflowed(&s) && dump_len >= sizeof(trunc_marker))
4804		memcpy(ei->dump + dump_len - sizeof(trunc_marker),
4805		       trunc_marker, sizeof(trunc_marker));
4806
4807	spin_unlock_irqrestore(&dump_lock, flags);
4808}
4809
4810static void scx_error_irq_workfn(struct irq_work *irq_work)
4811{
4812	struct scx_sched *sch = container_of(irq_work, struct scx_sched, error_irq_work);
4813	struct scx_exit_info *ei = sch->exit_info;
4814
4815	if (ei->kind >= SCX_EXIT_ERROR)
4816		scx_dump_state(ei, sch->ops.exit_dump_len);
4817
4818	kthread_queue_work(sch->helper, &sch->disable_work);
4819}
4820
4821static bool scx_vexit(struct scx_sched *sch,
4822		      enum scx_exit_kind kind, s64 exit_code,
4823		      const char *fmt, va_list args)
4824{
4825	struct scx_exit_info *ei = sch->exit_info;
4826
4827	guard(preempt)();
4828
4829	if (!scx_claim_exit(sch, kind))
4830		return false;
4831
4832	ei->exit_code = exit_code;
4833#ifdef CONFIG_STACKTRACE
4834	if (kind >= SCX_EXIT_ERROR)
4835		ei->bt_len = stack_trace_save(ei->bt, SCX_EXIT_BT_LEN, 1);
4836#endif
4837	vscnprintf(ei->msg, SCX_EXIT_MSG_LEN, fmt, args);
4838
4839	/*
4840	 * Set ei->kind and ->reason for scx_dump_state(). They'll be set again
4841	 * in scx_disable_workfn().
4842	 */
4843	ei->kind = kind;
4844	ei->reason = scx_exit_reason(ei->kind);
4845
4846	irq_work_queue(&sch->error_irq_work);
4847	return true;
4848}
4849
4850static int alloc_kick_syncs(void)
4851{
4852	int cpu;
4853
4854	/*
4855	 * Allocate per-CPU arrays sized by nr_cpu_ids. Use kvzalloc as size
4856	 * can exceed percpu allocator limits on large machines.
4857	 */
4858	for_each_possible_cpu(cpu) {
4859		struct scx_kick_syncs **ksyncs = per_cpu_ptr(&scx_kick_syncs, cpu);
4860		struct scx_kick_syncs *new_ksyncs;
4861
4862		WARN_ON_ONCE(rcu_access_pointer(*ksyncs));
4863
4864		new_ksyncs = kvzalloc_node(struct_size(new_ksyncs, syncs, nr_cpu_ids),
4865					   GFP_KERNEL, cpu_to_node(cpu));
4866		if (!new_ksyncs) {
4867			free_kick_syncs();
4868			return -ENOMEM;
4869		}
4870
4871		rcu_assign_pointer(*ksyncs, new_ksyncs);
4872	}
4873
4874	return 0;
4875}
4876
4877static struct scx_sched *scx_alloc_and_add_sched(struct sched_ext_ops *ops)
4878{
4879	struct scx_sched *sch;
4880	int node, ret;
4881
4882	sch = kzalloc(sizeof(*sch), GFP_KERNEL);
4883	if (!sch)
4884		return ERR_PTR(-ENOMEM);
4885
4886	sch->exit_info = alloc_exit_info(ops->exit_dump_len);
4887	if (!sch->exit_info) {
4888		ret = -ENOMEM;
4889		goto err_free_sch;
4890	}
4891
4892	ret = rhashtable_init(&sch->dsq_hash, &dsq_hash_params);
4893	if (ret < 0)
4894		goto err_free_ei;
4895
4896	sch->global_dsqs = kcalloc(nr_node_ids, sizeof(sch->global_dsqs[0]),
4897				   GFP_KERNEL);
4898	if (!sch->global_dsqs) {
4899		ret = -ENOMEM;
4900		goto err_free_hash;
4901	}
4902
4903	for_each_node_state(node, N_POSSIBLE) {
4904		struct scx_dispatch_q *dsq;
4905
4906		dsq = kzalloc_node(sizeof(*dsq), GFP_KERNEL, node);
4907		if (!dsq) {
4908			ret = -ENOMEM;
4909			goto err_free_gdsqs;
4910		}
4911
4912		init_dsq(dsq, SCX_DSQ_GLOBAL);
4913		sch->global_dsqs[node] = dsq;
4914	}
4915
4916	sch->pcpu = alloc_percpu(struct scx_sched_pcpu);
4917	if (!sch->pcpu) {
4918		ret = -ENOMEM;
4919		goto err_free_gdsqs;
4920	}
4921
4922	sch->helper = kthread_run_worker(0, "sched_ext_helper");
4923	if (IS_ERR(sch->helper)) {
4924		ret = PTR_ERR(sch->helper);
4925		goto err_free_pcpu;
4926	}
4927
4928	sched_set_fifo(sch->helper->task);
4929
4930	atomic_set(&sch->exit_kind, SCX_EXIT_NONE);
4931	init_irq_work(&sch->error_irq_work, scx_error_irq_workfn);
4932	kthread_init_work(&sch->disable_work, scx_disable_workfn);
4933	sch->ops = *ops;
4934	ops->priv = sch;
4935
4936	sch->kobj.kset = scx_kset;
4937	ret = kobject_init_and_add(&sch->kobj, &scx_ktype, NULL, "root");
4938	if (ret < 0)
4939		goto err_stop_helper;
4940
4941	return sch;
4942
4943err_stop_helper:
4944	kthread_destroy_worker(sch->helper);
4945err_free_pcpu:
4946	free_percpu(sch->pcpu);
4947err_free_gdsqs:
4948	for_each_node_state(node, N_POSSIBLE)
4949		kfree(sch->global_dsqs[node]);
4950	kfree(sch->global_dsqs);
4951err_free_hash:
4952	rhashtable_free_and_destroy(&sch->dsq_hash, NULL, NULL);
4953err_free_ei:
4954	free_exit_info(sch->exit_info);
4955err_free_sch:
4956	kfree(sch);
4957	return ERR_PTR(ret);
4958}
4959
4960static int check_hotplug_seq(struct scx_sched *sch,
4961			      const struct sched_ext_ops *ops)
4962{
4963	unsigned long long global_hotplug_seq;
4964
4965	/*
4966	 * If a hotplug event has occurred between when a scheduler was
4967	 * initialized, and when we were able to attach, exit and notify user
4968	 * space about it.
4969	 */
4970	if (ops->hotplug_seq) {
4971		global_hotplug_seq = atomic_long_read(&scx_hotplug_seq);
4972		if (ops->hotplug_seq != global_hotplug_seq) {
4973			scx_exit(sch, SCX_EXIT_UNREG_KERN,
4974				 SCX_ECODE_ACT_RESTART | SCX_ECODE_RSN_HOTPLUG,
4975				 "expected hotplug seq %llu did not match actual %llu",
4976				 ops->hotplug_seq, global_hotplug_seq);
4977			return -EBUSY;
4978		}
4979	}
4980
4981	return 0;
4982}
4983
4984static int validate_ops(struct scx_sched *sch, const struct sched_ext_ops *ops)
4985{
4986	/*
4987	 * It doesn't make sense to specify the SCX_OPS_ENQ_LAST flag if the
4988	 * ops.enqueue() callback isn't implemented.
4989	 */
4990	if ((ops->flags & SCX_OPS_ENQ_LAST) && !ops->enqueue) {
4991		scx_error(sch, "SCX_OPS_ENQ_LAST requires ops.enqueue() to be implemented");
4992		return -EINVAL;
4993	}
4994
4995	/*
4996	 * SCX_OPS_BUILTIN_IDLE_PER_NODE requires built-in CPU idle
4997	 * selection policy to be enabled.
4998	 */
4999	if ((ops->flags & SCX_OPS_BUILTIN_IDLE_PER_NODE) &&
5000	    (ops->update_idle && !(ops->flags & SCX_OPS_KEEP_BUILTIN_IDLE))) {
5001		scx_error(sch, "SCX_OPS_BUILTIN_IDLE_PER_NODE requires CPU idle selection enabled");
5002		return -EINVAL;
5003	}
5004
5005	if (ops->flags & SCX_OPS_HAS_CGROUP_WEIGHT)
5006		pr_warn("SCX_OPS_HAS_CGROUP_WEIGHT is deprecated and a noop\n");
5007
5008	if (ops->cpu_acquire || ops->cpu_release)
5009		pr_warn("ops->cpu_acquire/release() are deprecated, use sched_switch TP instead\n");
5010
5011	return 0;
5012}
5013
5014/*
5015 * scx_enable() is offloaded to a dedicated system-wide RT kthread to avoid
5016 * starvation. During the READY -> ENABLED task switching loop, the calling
5017 * thread's sched_class gets switched from fair to ext. As fair has higher
5018 * priority than ext, the calling thread can be indefinitely starved under
5019 * fair-class saturation, leading to a system hang.
5020 */
5021struct scx_enable_cmd {
5022	struct kthread_work	work;
5023	struct sched_ext_ops	*ops;
5024	int			ret;
5025};
5026
5027static void scx_enable_workfn(struct kthread_work *work)
5028{
5029	struct scx_enable_cmd *cmd =
5030		container_of(work, struct scx_enable_cmd, work);
5031	struct sched_ext_ops *ops = cmd->ops;
5032	struct scx_sched *sch;
5033	struct scx_task_iter sti;
5034	struct task_struct *p;
5035	unsigned long timeout;
5036	int i, cpu, ret;
5037
5038	mutex_lock(&scx_enable_mutex);
5039
5040	if (scx_enable_state() != SCX_DISABLED) {
5041		ret = -EBUSY;
5042		goto err_unlock;
5043	}
5044
5045	ret = alloc_kick_syncs();
5046	if (ret)
5047		goto err_unlock;
5048
5049	sch = scx_alloc_and_add_sched(ops);
5050	if (IS_ERR(sch)) {
5051		ret = PTR_ERR(sch);
5052		goto err_free_ksyncs;
5053	}
5054
5055	/*
5056	 * Transition to ENABLING and clear exit info to arm the disable path.
5057	 * Failure triggers full disabling from here on.
5058	 */
5059	WARN_ON_ONCE(scx_set_enable_state(SCX_ENABLING) != SCX_DISABLED);
5060	WARN_ON_ONCE(scx_root);
5061	if (WARN_ON_ONCE(READ_ONCE(scx_aborting)))
5062		WRITE_ONCE(scx_aborting, false);
5063
5064	atomic_long_set(&scx_nr_rejected, 0);
5065
5066	for_each_possible_cpu(cpu)
5067		cpu_rq(cpu)->scx.cpuperf_target = SCX_CPUPERF_ONE;
5068
5069	/*
5070	 * Keep CPUs stable during enable so that the BPF scheduler can track
5071	 * online CPUs by watching ->on/offline_cpu() after ->init().
5072	 */
5073	cpus_read_lock();
5074
5075	/*
5076	 * Make the scheduler instance visible. Must be inside cpus_read_lock().
5077	 * See handle_hotplug().
5078	 */
5079	rcu_assign_pointer(scx_root, sch);
5080
5081	scx_idle_enable(ops);
5082
5083	if (sch->ops.init) {
5084		ret = SCX_CALL_OP_RET(sch, SCX_KF_UNLOCKED, init, NULL);
5085		if (ret) {
5086			ret = ops_sanitize_err(sch, "init", ret);
5087			cpus_read_unlock();
5088			scx_error(sch, "ops.init() failed (%d)", ret);
5089			goto err_disable;
5090		}
5091		sch->exit_info->flags |= SCX_EFLAG_INITIALIZED;
5092	}
5093
5094	for (i = SCX_OPI_CPU_HOTPLUG_BEGIN; i < SCX_OPI_CPU_HOTPLUG_END; i++)
5095		if (((void (**)(void))ops)[i])
5096			set_bit(i, sch->has_op);
5097
5098	ret = check_hotplug_seq(sch, ops);
5099	if (ret) {
5100		cpus_read_unlock();
5101		goto err_disable;
5102	}
5103	scx_idle_update_selcpu_topology(ops);
5104
5105	cpus_read_unlock();
5106
5107	ret = validate_ops(sch, ops);
5108	if (ret)
5109		goto err_disable;
5110
5111	WARN_ON_ONCE(scx_dsp_ctx);
5112	scx_dsp_max_batch = ops->dispatch_max_batch ?: SCX_DSP_DFL_MAX_BATCH;
5113	scx_dsp_ctx = __alloc_percpu(struct_size_t(struct scx_dsp_ctx, buf,
5114						   scx_dsp_max_batch),
5115				     __alignof__(struct scx_dsp_ctx));
5116	if (!scx_dsp_ctx) {
5117		ret = -ENOMEM;
5118		goto err_disable;
5119	}
5120
5121	if (ops->timeout_ms)
5122		timeout = msecs_to_jiffies(ops->timeout_ms);
5123	else
5124		timeout = SCX_WATCHDOG_MAX_TIMEOUT;
5125
5126	WRITE_ONCE(scx_watchdog_timeout, timeout);
5127	WRITE_ONCE(scx_watchdog_timestamp, jiffies);
5128	queue_delayed_work(system_unbound_wq, &scx_watchdog_work,
5129			   scx_watchdog_timeout / 2);
5130
5131	/*
5132	 * Once __scx_enabled is set, %current can be switched to SCX anytime.
5133	 * This can lead to stalls as some BPF schedulers (e.g. userspace
5134	 * scheduling) may not function correctly before all tasks are switched.
5135	 * Init in bypass mode to guarantee forward progress.
5136	 */
5137	scx_bypass(true);
5138	scx_bypassed_for_enable = true;
5139
5140	for (i = SCX_OPI_NORMAL_BEGIN; i < SCX_OPI_NORMAL_END; i++)
5141		if (((void (**)(void))ops)[i])
5142			set_bit(i, sch->has_op);
5143
5144	if (sch->ops.cpu_acquire || sch->ops.cpu_release)
5145		sch->ops.flags |= SCX_OPS_HAS_CPU_PREEMPT;
5146
5147	/*
5148	 * Lock out forks, cgroup on/offlining and moves before opening the
5149	 * floodgate so that they don't wander into the operations prematurely.
5150	 */
5151	percpu_down_write(&scx_fork_rwsem);
5152
5153	WARN_ON_ONCE(scx_init_task_enabled);
5154	scx_init_task_enabled = true;
5155
5156	/*
5157	 * Enable ops for every task. Fork is excluded by scx_fork_rwsem
5158	 * preventing new tasks from being added. No need to exclude tasks
5159	 * leaving as sched_ext_free() can handle both prepped and enabled
5160	 * tasks. Prep all tasks first and then enable them with preemption
5161	 * disabled.
5162	 *
5163	 * All cgroups should be initialized before scx_init_task() so that the
5164	 * BPF scheduler can reliably track each task's cgroup membership from
5165	 * scx_init_task(). Lock out cgroup on/offlining and task migrations
5166	 * while tasks are being initialized so that scx_cgroup_can_attach()
5167	 * never sees uninitialized tasks.
5168	 */
5169	scx_cgroup_lock();
5170	ret = scx_cgroup_init(sch);
5171	if (ret)
5172		goto err_disable_unlock_all;
5173
5174	scx_task_iter_start(&sti);
5175	while ((p = scx_task_iter_next_locked(&sti))) {
5176		/*
5177		 * @p may already be dead, have lost all its usages counts and
5178		 * be waiting for RCU grace period before being freed. @p can't
5179		 * be initialized for SCX in such cases and should be ignored.
5180		 */
5181		if (!tryget_task_struct(p))
5182			continue;
5183
5184		scx_task_iter_unlock(&sti);
5185
5186		ret = scx_init_task(p, task_group(p), false);
5187		if (ret) {
5188			put_task_struct(p);
5189			scx_task_iter_stop(&sti);
5190			scx_error(sch, "ops.init_task() failed (%d) for %s[%d]",
5191				  ret, p->comm, p->pid);
5192			goto err_disable_unlock_all;
5193		}
5194
5195		scx_set_task_state(p, SCX_TASK_READY);
5196
5197		put_task_struct(p);
5198	}
5199	scx_task_iter_stop(&sti);
5200	scx_cgroup_unlock();
5201	percpu_up_write(&scx_fork_rwsem);
5202
5203	/*
5204	 * All tasks are READY. It's safe to turn on scx_enabled() and switch
5205	 * all eligible tasks.
5206	 */
5207	WRITE_ONCE(scx_switching_all, !(ops->flags & SCX_OPS_SWITCH_PARTIAL));
5208	static_branch_enable(&__scx_enabled);
5209
5210	/*
5211	 * We're fully committed and can't fail. The task READY -> ENABLED
5212	 * transitions here are synchronized against sched_ext_free() through
5213	 * scx_tasks_lock.
5214	 */
5215	percpu_down_write(&scx_fork_rwsem);
5216	scx_task_iter_start(&sti);
5217	while ((p = scx_task_iter_next_locked(&sti))) {
5218		unsigned int queue_flags = DEQUEUE_SAVE | DEQUEUE_MOVE;
5219		const struct sched_class *old_class = p->sched_class;
5220		const struct sched_class *new_class = scx_setscheduler_class(p);
5221
5222		if (scx_get_task_state(p) != SCX_TASK_READY)
5223			continue;
5224
5225		if (old_class != new_class)
5226			queue_flags |= DEQUEUE_CLASS;
5227
5228		scoped_guard (sched_change, p, queue_flags) {
5229			p->scx.slice = READ_ONCE(scx_slice_dfl);
5230			p->sched_class = new_class;
5231		}
5232	}
5233	scx_task_iter_stop(&sti);
5234	percpu_up_write(&scx_fork_rwsem);
5235
5236	scx_bypassed_for_enable = false;
5237	scx_bypass(false);
5238
5239	if (!scx_tryset_enable_state(SCX_ENABLED, SCX_ENABLING)) {
5240		WARN_ON_ONCE(atomic_read(&sch->exit_kind) == SCX_EXIT_NONE);
5241		goto err_disable;
5242	}
5243
5244	if (!(ops->flags & SCX_OPS_SWITCH_PARTIAL))
5245		static_branch_enable(&__scx_switched_all);
5246
5247	pr_info("sched_ext: BPF scheduler \"%s\" enabled%s\n",
5248		sch->ops.name, scx_switched_all() ? "" : " (partial)");
5249	kobject_uevent(&sch->kobj, KOBJ_ADD);
5250	mutex_unlock(&scx_enable_mutex);
5251
5252	atomic_long_inc(&scx_enable_seq);
5253
5254	cmd->ret = 0;
5255	return;
5256
5257err_free_ksyncs:
5258	free_kick_syncs();
5259err_unlock:
5260	mutex_unlock(&scx_enable_mutex);
5261	cmd->ret = ret;
5262	return;
5263
5264err_disable_unlock_all:
5265	scx_cgroup_unlock();
5266	percpu_up_write(&scx_fork_rwsem);
5267	/* we'll soon enter disable path, keep bypass on */
5268err_disable:
5269	mutex_unlock(&scx_enable_mutex);
5270	/*
5271	 * Returning an error code here would not pass all the error information
5272	 * to userspace. Record errno using scx_error() for cases scx_error()
5273	 * wasn't already invoked and exit indicating success so that the error
5274	 * is notified through ops.exit() with all the details.
5275	 *
5276	 * Flush scx_disable_work to ensure that error is reported before init
5277	 * completion. sch's base reference will be put by bpf_scx_unreg().
5278	 */
5279	scx_error(sch, "scx_enable() failed (%d)", ret);
5280	kthread_flush_work(&sch->disable_work);
5281	cmd->ret = 0;
5282}
5283
5284static int scx_enable(struct sched_ext_ops *ops, struct bpf_link *link)
5285{
5286	static struct kthread_worker *helper;
5287	static DEFINE_MUTEX(helper_mutex);
5288	struct scx_enable_cmd cmd;
5289
5290	if (!cpumask_equal(housekeeping_cpumask(HK_TYPE_DOMAIN),
5291			   cpu_possible_mask)) {
5292		pr_err("sched_ext: Not compatible with \"isolcpus=\" domain isolation\n");
5293		return -EINVAL;
5294	}
5295
5296	if (!READ_ONCE(helper)) {
5297		mutex_lock(&helper_mutex);
5298		if (!helper) {
5299			struct kthread_worker *w =
5300				kthread_run_worker(0, "scx_enable_helper");
5301			if (IS_ERR_OR_NULL(w)) {
5302				mutex_unlock(&helper_mutex);
5303				return -ENOMEM;
5304			}
5305			sched_set_fifo(w->task);
5306			WRITE_ONCE(helper, w);
5307		}
5308		mutex_unlock(&helper_mutex);
5309	}
5310
5311	kthread_init_work(&cmd.work, scx_enable_workfn);
5312	cmd.ops = ops;
5313
5314	kthread_queue_work(READ_ONCE(helper), &cmd.work);
5315	kthread_flush_work(&cmd.work);
5316	return cmd.ret;
5317}
5318
5319
5320/********************************************************************************
5321 * bpf_struct_ops plumbing.
5322 */
5323#include <linux/bpf_verifier.h>
5324#include <linux/bpf.h>
5325#include <linux/btf.h>
5326
5327static const struct btf_type *task_struct_type;
5328
5329static bool bpf_scx_is_valid_access(int off, int size,
5330				    enum bpf_access_type type,
5331				    const struct bpf_prog *prog,
5332				    struct bpf_insn_access_aux *info)
5333{
5334	if (type != BPF_READ)
5335		return false;
5336	if (off < 0 || off >= sizeof(__u64) * MAX_BPF_FUNC_ARGS)
5337		return false;
5338	if (off % size != 0)
5339		return false;
5340
5341	return btf_ctx_access(off, size, type, prog, info);
5342}
5343
5344static int bpf_scx_btf_struct_access(struct bpf_verifier_log *log,
5345				     const struct bpf_reg_state *reg, int off,
5346				     int size)
5347{
5348	const struct btf_type *t;
5349
5350	t = btf_type_by_id(reg->btf, reg->btf_id);
5351	if (t == task_struct_type) {
5352		if (off >= offsetof(struct task_struct, scx.slice) &&
5353		    off + size <= offsetofend(struct task_struct, scx.slice))
5354			return SCALAR_VALUE;
5355		if (off >= offsetof(struct task_struct, scx.dsq_vtime) &&
5356		    off + size <= offsetofend(struct task_struct, scx.dsq_vtime))
5357			return SCALAR_VALUE;
5358		if (off >= offsetof(struct task_struct, scx.disallow) &&
5359		    off + size <= offsetofend(struct task_struct, scx.disallow))
5360			return SCALAR_VALUE;
5361	}
5362
5363	return -EACCES;
5364}
5365
5366static const struct bpf_verifier_ops bpf_scx_verifier_ops = {
5367	.get_func_proto = bpf_base_func_proto,
5368	.is_valid_access = bpf_scx_is_valid_access,
5369	.btf_struct_access = bpf_scx_btf_struct_access,
5370};
5371
5372static int bpf_scx_init_member(const struct btf_type *t,
5373			       const struct btf_member *member,
5374			       void *kdata, const void *udata)
5375{
5376	const struct sched_ext_ops *uops = udata;
5377	struct sched_ext_ops *ops = kdata;
5378	u32 moff = __btf_member_bit_offset(t, member) / 8;
5379	int ret;
5380
5381	switch (moff) {
5382	case offsetof(struct sched_ext_ops, dispatch_max_batch):
5383		if (*(u32 *)(udata + moff) > INT_MAX)
5384			return -E2BIG;
5385		ops->dispatch_max_batch = *(u32 *)(udata + moff);
5386		return 1;
5387	case offsetof(struct sched_ext_ops, flags):
5388		if (*(u64 *)(udata + moff) & ~SCX_OPS_ALL_FLAGS)
5389			return -EINVAL;
5390		ops->flags = *(u64 *)(udata + moff);
5391		return 1;
5392	case offsetof(struct sched_ext_ops, name):
5393		ret = bpf_obj_name_cpy(ops->name, uops->name,
5394				       sizeof(ops->name));
5395		if (ret < 0)
5396			return ret;
5397		if (ret == 0)
5398			return -EINVAL;
5399		return 1;
5400	case offsetof(struct sched_ext_ops, timeout_ms):
5401		if (msecs_to_jiffies(*(u32 *)(udata + moff)) >
5402		    SCX_WATCHDOG_MAX_TIMEOUT)
5403			return -E2BIG;
5404		ops->timeout_ms = *(u32 *)(udata + moff);
5405		return 1;
5406	case offsetof(struct sched_ext_ops, exit_dump_len):
5407		ops->exit_dump_len =
5408			*(u32 *)(udata + moff) ?: SCX_EXIT_DUMP_DFL_LEN;
5409		return 1;
5410	case offsetof(struct sched_ext_ops, hotplug_seq):
5411		ops->hotplug_seq = *(u64 *)(udata + moff);
5412		return 1;
5413	}
5414
5415	return 0;
5416}
5417
5418static int bpf_scx_check_member(const struct btf_type *t,
5419				const struct btf_member *member,
5420				const struct bpf_prog *prog)
5421{
5422	u32 moff = __btf_member_bit_offset(t, member) / 8;
5423
5424	switch (moff) {
5425	case offsetof(struct sched_ext_ops, init_task):
5426#ifdef CONFIG_EXT_GROUP_SCHED
5427	case offsetof(struct sched_ext_ops, cgroup_init):
5428	case offsetof(struct sched_ext_ops, cgroup_exit):
5429	case offsetof(struct sched_ext_ops, cgroup_prep_move):
5430#endif
5431	case offsetof(struct sched_ext_ops, cpu_online):
5432	case offsetof(struct sched_ext_ops, cpu_offline):
5433	case offsetof(struct sched_ext_ops, init):
5434	case offsetof(struct sched_ext_ops, exit):
5435		break;
5436	default:
5437		if (prog->sleepable)
5438			return -EINVAL;
5439	}
5440
5441	return 0;
5442}
5443
5444static int bpf_scx_reg(void *kdata, struct bpf_link *link)
5445{
5446	return scx_enable(kdata, link);
5447}
5448
5449static void bpf_scx_unreg(void *kdata, struct bpf_link *link)
5450{
5451	struct sched_ext_ops *ops = kdata;
5452	struct scx_sched *sch = ops->priv;
5453
5454	scx_disable(SCX_EXIT_UNREG);
5455	kthread_flush_work(&sch->disable_work);
5456	kobject_put(&sch->kobj);
5457}
5458
5459static int bpf_scx_init(struct btf *btf)
5460{
5461	task_struct_type = btf_type_by_id(btf, btf_tracing_ids[BTF_TRACING_TYPE_TASK]);
5462
5463	return 0;
5464}
5465
5466static int bpf_scx_update(void *kdata, void *old_kdata, struct bpf_link *link)
5467{
5468	/*
5469	 * sched_ext does not support updating the actively-loaded BPF
5470	 * scheduler, as registering a BPF scheduler can always fail if the
5471	 * scheduler returns an error code for e.g. ops.init(), ops.init_task(),
5472	 * etc. Similarly, we can always race with unregistration happening
5473	 * elsewhere, such as with sysrq.
5474	 */
5475	return -EOPNOTSUPP;
5476}
5477
5478static int bpf_scx_validate(void *kdata)
5479{
5480	return 0;
5481}
5482
5483static s32 sched_ext_ops__select_cpu(struct task_struct *p, s32 prev_cpu, u64 wake_flags) { return -EINVAL; }
5484static void sched_ext_ops__enqueue(struct task_struct *p, u64 enq_flags) {}
5485static void sched_ext_ops__dequeue(struct task_struct *p, u64 enq_flags) {}
5486static void sched_ext_ops__dispatch(s32 prev_cpu, struct task_struct *prev__nullable) {}
5487static void sched_ext_ops__tick(struct task_struct *p) {}
5488static void sched_ext_ops__runnable(struct task_struct *p, u64 enq_flags) {}
5489static void sched_ext_ops__running(struct task_struct *p) {}
5490static void sched_ext_ops__stopping(struct task_struct *p, bool runnable) {}
5491static void sched_ext_ops__quiescent(struct task_struct *p, u64 deq_flags) {}
5492static bool sched_ext_ops__yield(struct task_struct *from, struct task_struct *to__nullable) { return false; }
5493static bool sched_ext_ops__core_sched_before(struct task_struct *a, struct task_struct *b) { return false; }
5494static void sched_ext_ops__set_weight(struct task_struct *p, u32 weight) {}
5495static void sched_ext_ops__set_cpumask(struct task_struct *p, const struct cpumask *mask) {}
5496static void sched_ext_ops__update_idle(s32 cpu, bool idle) {}
5497static void sched_ext_ops__cpu_acquire(s32 cpu, struct scx_cpu_acquire_args *args) {}
5498static void sched_ext_ops__cpu_release(s32 cpu, struct scx_cpu_release_args *args) {}
5499static s32 sched_ext_ops__init_task(struct task_struct *p, struct scx_init_task_args *args) { return -EINVAL; }
5500static void sched_ext_ops__exit_task(struct task_struct *p, struct scx_exit_task_args *args) {}
5501static void sched_ext_ops__enable(struct task_struct *p) {}
5502static void sched_ext_ops__disable(struct task_struct *p) {}
5503#ifdef CONFIG_EXT_GROUP_SCHED
5504static s32 sched_ext_ops__cgroup_init(struct cgroup *cgrp, struct scx_cgroup_init_args *args) { return -EINVAL; }
5505static void sched_ext_ops__cgroup_exit(struct cgroup *cgrp) {}
5506static s32 sched_ext_ops__cgroup_prep_move(struct task_struct *p, struct cgroup *from, struct cgroup *to) { return -EINVAL; }
5507static void sched_ext_ops__cgroup_move(struct task_struct *p, struct cgroup *from, struct cgroup *to) {}
5508static void sched_ext_ops__cgroup_cancel_move(struct task_struct *p, struct cgroup *from, struct cgroup *to) {}
5509static void sched_ext_ops__cgroup_set_weight(struct cgroup *cgrp, u32 weight) {}
5510static void sched_ext_ops__cgroup_set_bandwidth(struct cgroup *cgrp, u64 period_us, u64 quota_us, u64 burst_us) {}
5511static void sched_ext_ops__cgroup_set_idle(struct cgroup *cgrp, bool idle) {}
5512#endif
5513static void sched_ext_ops__cpu_online(s32 cpu) {}
5514static void sched_ext_ops__cpu_offline(s32 cpu) {}
5515static s32 sched_ext_ops__init(void) { return -EINVAL; }
5516static void sched_ext_ops__exit(struct scx_exit_info *info) {}
5517static void sched_ext_ops__dump(struct scx_dump_ctx *ctx) {}
5518static void sched_ext_ops__dump_cpu(struct scx_dump_ctx *ctx, s32 cpu, bool idle) {}
5519static void sched_ext_ops__dump_task(struct scx_dump_ctx *ctx, struct task_struct *p) {}
5520
5521static struct sched_ext_ops __bpf_ops_sched_ext_ops = {
5522	.select_cpu		= sched_ext_ops__select_cpu,
5523	.enqueue		= sched_ext_ops__enqueue,
5524	.dequeue		= sched_ext_ops__dequeue,
5525	.dispatch		= sched_ext_ops__dispatch,
5526	.tick			= sched_ext_ops__tick,
5527	.runnable		= sched_ext_ops__runnable,
5528	.running		= sched_ext_ops__running,
5529	.stopping		= sched_ext_ops__stopping,
5530	.quiescent		= sched_ext_ops__quiescent,
5531	.yield			= sched_ext_ops__yield,
5532	.core_sched_before	= sched_ext_ops__core_sched_before,
5533	.set_weight		= sched_ext_ops__set_weight,
5534	.set_cpumask		= sched_ext_ops__set_cpumask,
5535	.update_idle		= sched_ext_ops__update_idle,
5536	.cpu_acquire		= sched_ext_ops__cpu_acquire,
5537	.cpu_release		= sched_ext_ops__cpu_release,
5538	.init_task		= sched_ext_ops__init_task,
5539	.exit_task		= sched_ext_ops__exit_task,
5540	.enable			= sched_ext_ops__enable,
5541	.disable		= sched_ext_ops__disable,
5542#ifdef CONFIG_EXT_GROUP_SCHED
5543	.cgroup_init		= sched_ext_ops__cgroup_init,
5544	.cgroup_exit		= sched_ext_ops__cgroup_exit,
5545	.cgroup_prep_move	= sched_ext_ops__cgroup_prep_move,
5546	.cgroup_move		= sched_ext_ops__cgroup_move,
5547	.cgroup_cancel_move	= sched_ext_ops__cgroup_cancel_move,
5548	.cgroup_set_weight	= sched_ext_ops__cgroup_set_weight,
5549	.cgroup_set_bandwidth	= sched_ext_ops__cgroup_set_bandwidth,
5550	.cgroup_set_idle	= sched_ext_ops__cgroup_set_idle,
5551#endif
5552	.cpu_online		= sched_ext_ops__cpu_online,
5553	.cpu_offline		= sched_ext_ops__cpu_offline,
5554	.init			= sched_ext_ops__init,
5555	.exit			= sched_ext_ops__exit,
5556	.dump			= sched_ext_ops__dump,
5557	.dump_cpu		= sched_ext_ops__dump_cpu,
5558	.dump_task		= sched_ext_ops__dump_task,
5559};
5560
5561static struct bpf_struct_ops bpf_sched_ext_ops = {
5562	.verifier_ops = &bpf_scx_verifier_ops,
5563	.reg = bpf_scx_reg,
5564	.unreg = bpf_scx_unreg,
5565	.check_member = bpf_scx_check_member,
5566	.init_member = bpf_scx_init_member,
5567	.init = bpf_scx_init,
5568	.update = bpf_scx_update,
5569	.validate = bpf_scx_validate,
5570	.name = "sched_ext_ops",
5571	.owner = THIS_MODULE,
5572	.cfi_stubs = &__bpf_ops_sched_ext_ops
5573};
5574
5575
5576/********************************************************************************
5577 * System integration and init.
5578 */
5579
5580static void sysrq_handle_sched_ext_reset(u8 key)
5581{
5582	scx_disable(SCX_EXIT_SYSRQ);
5583}
5584
5585static const struct sysrq_key_op sysrq_sched_ext_reset_op = {
5586	.handler	= sysrq_handle_sched_ext_reset,
5587	.help_msg	= "reset-sched-ext(S)",
5588	.action_msg	= "Disable sched_ext and revert all tasks to CFS",
5589	.enable_mask	= SYSRQ_ENABLE_RTNICE,
5590};
5591
5592static void sysrq_handle_sched_ext_dump(u8 key)
5593{
5594	struct scx_exit_info ei = { .kind = SCX_EXIT_NONE, .reason = "SysRq-D" };
5595
5596	if (scx_enabled())
5597		scx_dump_state(&ei, 0);
5598}
5599
5600static const struct sysrq_key_op sysrq_sched_ext_dump_op = {
5601	.handler	= sysrq_handle_sched_ext_dump,
5602	.help_msg	= "dump-sched-ext(D)",
5603	.action_msg	= "Trigger sched_ext debug dump",
5604	.enable_mask	= SYSRQ_ENABLE_RTNICE,
5605};
5606
5607static bool can_skip_idle_kick(struct rq *rq)
5608{
5609	lockdep_assert_rq_held(rq);
5610
5611	/*
5612	 * We can skip idle kicking if @rq is going to go through at least one
5613	 * full SCX scheduling cycle before going idle. Just checking whether
5614	 * curr is not idle is insufficient because we could be racing
5615	 * balance_one() trying to pull the next task from a remote rq, which
5616	 * may fail, and @rq may become idle afterwards.
5617	 *
5618	 * The race window is small and we don't and can't guarantee that @rq is
5619	 * only kicked while idle anyway. Skip only when sure.
5620	 */
5621	return !is_idle_task(rq->curr) && !(rq->scx.flags & SCX_RQ_IN_BALANCE);
5622}
5623
5624static bool kick_one_cpu(s32 cpu, struct rq *this_rq, unsigned long *ksyncs)
5625{
5626	struct rq *rq = cpu_rq(cpu);
5627	struct scx_rq *this_scx = &this_rq->scx;
5628	const struct sched_class *cur_class;
5629	bool should_wait = false;
5630	unsigned long flags;
5631
5632	raw_spin_rq_lock_irqsave(rq, flags);
5633	cur_class = rq->curr->sched_class;
5634
5635	/*
5636	 * During CPU hotplug, a CPU may depend on kicking itself to make
5637	 * forward progress. Allow kicking self regardless of online state. If
5638	 * @cpu is running a higher class task, we have no control over @cpu.
5639	 * Skip kicking.
5640	 */
5641	if ((cpu_online(cpu) || cpu == cpu_of(this_rq)) &&
5642	    !sched_class_above(cur_class, &ext_sched_class)) {
5643		if (cpumask_test_cpu(cpu, this_scx->cpus_to_preempt)) {
5644			if (cur_class == &ext_sched_class)
5645				rq->curr->scx.slice = 0;
5646			cpumask_clear_cpu(cpu, this_scx->cpus_to_preempt);
5647		}
5648
5649		if (cpumask_test_cpu(cpu, this_scx->cpus_to_wait)) {
5650			if (cur_class == &ext_sched_class) {
5651				cpumask_set_cpu(cpu, this_scx->cpus_to_sync);
5652				ksyncs[cpu] = rq->scx.kick_sync;
5653				should_wait = true;
5654			}
5655			cpumask_clear_cpu(cpu, this_scx->cpus_to_wait);
5656		}
5657
5658		resched_curr(rq);
5659	} else {
5660		cpumask_clear_cpu(cpu, this_scx->cpus_to_preempt);
5661		cpumask_clear_cpu(cpu, this_scx->cpus_to_wait);
5662	}
5663
5664	raw_spin_rq_unlock_irqrestore(rq, flags);
5665
5666	return should_wait;
5667}
5668
5669static void kick_one_cpu_if_idle(s32 cpu, struct rq *this_rq)
5670{
5671	struct rq *rq = cpu_rq(cpu);
5672	unsigned long flags;
5673
5674	raw_spin_rq_lock_irqsave(rq, flags);
5675
5676	if (!can_skip_idle_kick(rq) &&
5677	    (cpu_online(cpu) || cpu == cpu_of(this_rq)))
5678		resched_curr(rq);
5679
5680	raw_spin_rq_unlock_irqrestore(rq, flags);
5681}
5682
5683static void kick_cpus_irq_workfn(struct irq_work *irq_work)
5684{
5685	struct rq *this_rq = this_rq();
5686	struct scx_rq *this_scx = &this_rq->scx;
5687	struct scx_kick_syncs __rcu *ksyncs_pcpu = __this_cpu_read(scx_kick_syncs);
5688	bool should_wait = false;
5689	unsigned long *ksyncs;
5690	s32 cpu;
5691
5692	if (unlikely(!ksyncs_pcpu)) {
5693		pr_warn_once("kick_cpus_irq_workfn() called with NULL scx_kick_syncs");
5694		return;
5695	}
5696
5697	ksyncs = rcu_dereference_bh(ksyncs_pcpu)->syncs;
5698
5699	for_each_cpu(cpu, this_scx->cpus_to_kick) {
5700		should_wait |= kick_one_cpu(cpu, this_rq, ksyncs);
5701		cpumask_clear_cpu(cpu, this_scx->cpus_to_kick);
5702		cpumask_clear_cpu(cpu, this_scx->cpus_to_kick_if_idle);
5703	}
5704
5705	for_each_cpu(cpu, this_scx->cpus_to_kick_if_idle) {
5706		kick_one_cpu_if_idle(cpu, this_rq);
5707		cpumask_clear_cpu(cpu, this_scx->cpus_to_kick_if_idle);
5708	}
5709
5710	/*
5711	 * Can't wait in hardirq — kick_sync can't advance, deadlocking if
5712	 * CPUs wait for each other. Defer to kick_sync_wait_bal_cb().
5713	 */
5714	if (should_wait) {
5715		raw_spin_rq_lock(this_rq);
5716		this_scx->kick_sync_pending = true;
5717		resched_curr(this_rq);
5718		raw_spin_rq_unlock(this_rq);
5719	}
5720}
5721
5722/**
5723 * print_scx_info - print out sched_ext scheduler state
5724 * @log_lvl: the log level to use when printing
5725 * @p: target task
5726 *
5727 * If a sched_ext scheduler is enabled, print the name and state of the
5728 * scheduler. If @p is on sched_ext, print further information about the task.
5729 *
5730 * This function can be safely called on any task as long as the task_struct
5731 * itself is accessible. While safe, this function isn't synchronized and may
5732 * print out mixups or garbages of limited length.
5733 */
5734void print_scx_info(const char *log_lvl, struct task_struct *p)
5735{
5736	struct scx_sched *sch = scx_root;
5737	enum scx_enable_state state = scx_enable_state();
5738	const char *all = READ_ONCE(scx_switching_all) ? "+all" : "";
5739	char runnable_at_buf[22] = "?";
5740	struct sched_class *class;
5741	unsigned long runnable_at;
5742
5743	if (state == SCX_DISABLED)
5744		return;
5745
5746	/*
5747	 * Carefully check if the task was running on sched_ext, and then
5748	 * carefully copy the time it's been runnable, and its state.
5749	 */
5750	if (copy_from_kernel_nofault(&class, &p->sched_class, sizeof(class)) ||
5751	    class != &ext_sched_class) {
5752		printk("%sSched_ext: %s (%s%s)", log_lvl, sch->ops.name,
5753		       scx_enable_state_str[state], all);
5754		return;
5755	}
5756
5757	if (!copy_from_kernel_nofault(&runnable_at, &p->scx.runnable_at,
5758				      sizeof(runnable_at)))
5759		scnprintf(runnable_at_buf, sizeof(runnable_at_buf), "%+ldms",
5760			  jiffies_delta_msecs(runnable_at, jiffies));
5761
5762	/* print everything onto one line to conserve console space */
5763	printk("%sSched_ext: %s (%s%s), task: runnable_at=%s",
5764	       log_lvl, sch->ops.name, scx_enable_state_str[state], all,
5765	       runnable_at_buf);
5766}
5767
5768static int scx_pm_handler(struct notifier_block *nb, unsigned long event, void *ptr)
5769{
5770	/*
5771	 * SCX schedulers often have userspace components which are sometimes
5772	 * involved in critial scheduling paths. PM operations involve freezing
5773	 * userspace which can lead to scheduling misbehaviors including stalls.
5774	 * Let's bypass while PM operations are in progress.
5775	 */
5776	switch (event) {
5777	case PM_HIBERNATION_PREPARE:
5778	case PM_SUSPEND_PREPARE:
5779	case PM_RESTORE_PREPARE:
5780		scx_bypass(true);
5781		break;
5782	case PM_POST_HIBERNATION:
5783	case PM_POST_SUSPEND:
5784	case PM_POST_RESTORE:
5785		scx_bypass(false);
5786		break;
5787	}
5788
5789	return NOTIFY_OK;
5790}
5791
5792static struct notifier_block scx_pm_notifier = {
5793	.notifier_call = scx_pm_handler,
5794};
5795
5796void __init init_sched_ext_class(void)
5797{
5798	s32 cpu, v;
5799
5800	/*
5801	 * The following is to prevent the compiler from optimizing out the enum
5802	 * definitions so that BPF scheduler implementations can use them
5803	 * through the generated vmlinux.h.
5804	 */
5805	WRITE_ONCE(v, SCX_ENQ_WAKEUP | SCX_DEQ_SLEEP | SCX_KICK_PREEMPT |
5806		   SCX_TG_ONLINE);
5807
5808	scx_idle_init_masks();
5809
5810	for_each_possible_cpu(cpu) {
5811		struct rq *rq = cpu_rq(cpu);
5812		int  n = cpu_to_node(cpu);
5813
5814		init_dsq(&rq->scx.local_dsq, SCX_DSQ_LOCAL);
5815		init_dsq(&rq->scx.bypass_dsq, SCX_DSQ_BYPASS);
5816		INIT_LIST_HEAD(&rq->scx.runnable_list);
5817		INIT_LIST_HEAD(&rq->scx.ddsp_deferred_locals);
5818
5819		BUG_ON(!zalloc_cpumask_var_node(&rq->scx.cpus_to_kick, GFP_KERNEL, n));
5820		BUG_ON(!zalloc_cpumask_var_node(&rq->scx.cpus_to_kick_if_idle, GFP_KERNEL, n));
5821		BUG_ON(!zalloc_cpumask_var_node(&rq->scx.cpus_to_preempt, GFP_KERNEL, n));
5822		BUG_ON(!zalloc_cpumask_var_node(&rq->scx.cpus_to_wait, GFP_KERNEL, n));
5823		BUG_ON(!zalloc_cpumask_var_node(&rq->scx.cpus_to_sync, GFP_KERNEL, n));
5824		rq->scx.deferred_irq_work = IRQ_WORK_INIT_HARD(deferred_irq_workfn);
5825		rq->scx.kick_cpus_irq_work = IRQ_WORK_INIT_HARD(kick_cpus_irq_workfn);
5826
5827		if (cpu_online(cpu))
5828			cpu_rq(cpu)->scx.flags |= SCX_RQ_ONLINE;
5829	}
5830
5831	register_sysrq_key('S', &sysrq_sched_ext_reset_op);
5832	register_sysrq_key('D', &sysrq_sched_ext_dump_op);
5833	INIT_DELAYED_WORK(&scx_watchdog_work, scx_watchdog_workfn);
5834}
5835
5836
5837/********************************************************************************
5838 * Helpers that can be called from the BPF scheduler.
5839 */
5840static bool scx_dsq_insert_preamble(struct scx_sched *sch, struct task_struct *p,
5841				    u64 enq_flags)
5842{
5843	if (!scx_kf_allowed(sch, SCX_KF_ENQUEUE | SCX_KF_DISPATCH))
5844		return false;
5845
5846	lockdep_assert_irqs_disabled();
5847
5848	if (unlikely(!p)) {
5849		scx_error(sch, "called with NULL task");
5850		return false;
5851	}
5852
5853	if (unlikely(enq_flags & __SCX_ENQ_INTERNAL_MASK)) {
5854		scx_error(sch, "invalid enq_flags 0x%llx", enq_flags);
5855		return false;
5856	}
5857
5858	return true;
5859}
5860
5861static void scx_dsq_insert_commit(struct scx_sched *sch, struct task_struct *p,
5862				  u64 dsq_id, u64 enq_flags)
5863{
5864	struct scx_dsp_ctx *dspc = this_cpu_ptr(scx_dsp_ctx);
5865	struct task_struct *ddsp_task;
5866
5867	ddsp_task = __this_cpu_read(direct_dispatch_task);
5868	if (ddsp_task) {
5869		mark_direct_dispatch(sch, ddsp_task, p, dsq_id, enq_flags);
5870		return;
5871	}
5872
5873	if (unlikely(dspc->cursor >= scx_dsp_max_batch)) {
5874		scx_error(sch, "dispatch buffer overflow");
5875		return;
5876	}
5877
5878	dspc->buf[dspc->cursor++] = (struct scx_dsp_buf_ent){
5879		.task = p,
5880		.qseq = atomic_long_read(&p->scx.ops_state) & SCX_OPSS_QSEQ_MASK,
5881		.dsq_id = dsq_id,
5882		.enq_flags = enq_flags,
5883	};
5884}
5885
5886__bpf_kfunc_start_defs();
5887
5888/**
5889 * scx_bpf_dsq_insert - Insert a task into the FIFO queue of a DSQ
5890 * @p: task_struct to insert
5891 * @dsq_id: DSQ to insert into
5892 * @slice: duration @p can run for in nsecs, 0 to keep the current value
5893 * @enq_flags: SCX_ENQ_*
5894 *
5895 * Insert @p into the FIFO queue of the DSQ identified by @dsq_id. It is safe to
5896 * call this function spuriously. Can be called from ops.enqueue(),
5897 * ops.select_cpu(), and ops.dispatch().
5898 *
5899 * When called from ops.select_cpu() or ops.enqueue(), it's for direct dispatch
5900 * and @p must match the task being enqueued.
5901 *
5902 * When called from ops.select_cpu(), @enq_flags and @dsp_id are stored, and @p
5903 * will be directly inserted into the corresponding dispatch queue after
5904 * ops.select_cpu() returns. If @p is inserted into SCX_DSQ_LOCAL, it will be
5905 * inserted into the local DSQ of the CPU returned by ops.select_cpu().
5906 * @enq_flags are OR'd with the enqueue flags on the enqueue path before the
5907 * task is inserted.
5908 *
5909 * When called from ops.dispatch(), there are no restrictions on @p or @dsq_id
5910 * and this function can be called upto ops.dispatch_max_batch times to insert
5911 * multiple tasks. scx_bpf_dispatch_nr_slots() returns the number of the
5912 * remaining slots. scx_bpf_dsq_move_to_local() flushes the batch and resets the
5913 * counter.
5914 *
5915 * This function doesn't have any locking restrictions and may be called under
5916 * BPF locks (in the future when BPF introduces more flexible locking).
5917 *
5918 * @p is allowed to run for @slice. The scheduling path is triggered on slice
5919 * exhaustion. If zero, the current residual slice is maintained. If
5920 * %SCX_SLICE_INF, @p never expires and the BPF scheduler must kick the CPU with
5921 * scx_bpf_kick_cpu() to trigger scheduling.
5922 *
5923 * Returns %true on successful insertion, %false on failure. On the root
5924 * scheduler, %false return triggers scheduler abort and the caller doesn't need
5925 * to check the return value.
5926 */
5927__bpf_kfunc bool scx_bpf_dsq_insert___v2(struct task_struct *p, u64 dsq_id,
5928					 u64 slice, u64 enq_flags)
5929{
5930	struct scx_sched *sch;
5931
5932	guard(rcu)();
5933	sch = rcu_dereference(scx_root);
5934	if (unlikely(!sch))
5935		return false;
5936
5937	if (!scx_dsq_insert_preamble(sch, p, enq_flags))
5938		return false;
5939
5940	if (slice)
5941		p->scx.slice = slice;
5942	else
5943		p->scx.slice = p->scx.slice ?: 1;
5944
5945	scx_dsq_insert_commit(sch, p, dsq_id, enq_flags);
5946
5947	return true;
5948}
5949
5950/*
5951 * COMPAT: Will be removed in v6.23 along with the ___v2 suffix.
5952 */
5953__bpf_kfunc void scx_bpf_dsq_insert(struct task_struct *p, u64 dsq_id,
5954					     u64 slice, u64 enq_flags)
5955{
5956	scx_bpf_dsq_insert___v2(p, dsq_id, slice, enq_flags);
5957}
5958
5959static bool scx_dsq_insert_vtime(struct scx_sched *sch, struct task_struct *p,
5960				 u64 dsq_id, u64 slice, u64 vtime, u64 enq_flags)
5961{
5962	if (!scx_dsq_insert_preamble(sch, p, enq_flags))
5963		return false;
5964
5965	if (slice)
5966		p->scx.slice = slice;
5967	else
5968		p->scx.slice = p->scx.slice ?: 1;
5969
5970	p->scx.dsq_vtime = vtime;
5971
5972	scx_dsq_insert_commit(sch, p, dsq_id, enq_flags | SCX_ENQ_DSQ_PRIQ);
5973
5974	return true;
5975}
5976
5977struct scx_bpf_dsq_insert_vtime_args {
5978	/* @p can't be packed together as KF_RCU is not transitive */
5979	u64			dsq_id;
5980	u64			slice;
5981	u64			vtime;
5982	u64			enq_flags;
5983};
5984
5985/**
5986 * __scx_bpf_dsq_insert_vtime - Arg-wrapped vtime DSQ insertion
5987 * @p: task_struct to insert
5988 * @args: struct containing the rest of the arguments
5989 *       @args->dsq_id: DSQ to insert into
5990 *       @args->slice: duration @p can run for in nsecs, 0 to keep the current value
5991 *       @args->vtime: @p's ordering inside the vtime-sorted queue of the target DSQ
5992 *       @args->enq_flags: SCX_ENQ_*
5993 *
5994 * Wrapper kfunc that takes arguments via struct to work around BPF's 5 argument
5995 * limit. BPF programs should use scx_bpf_dsq_insert_vtime() which is provided
5996 * as an inline wrapper in common.bpf.h.
5997 *
5998 * Insert @p into the vtime priority queue of the DSQ identified by
5999 * @args->dsq_id. Tasks queued into the priority queue are ordered by
6000 * @args->vtime. All other aspects are identical to scx_bpf_dsq_insert().
6001 *
6002 * @args->vtime ordering is according to time_before64() which considers
6003 * wrapping. A numerically larger vtime may indicate an earlier position in the
6004 * ordering and vice-versa.
6005 *
6006 * A DSQ can only be used as a FIFO or priority queue at any given time and this
6007 * function must not be called on a DSQ which already has one or more FIFO tasks
6008 * queued and vice-versa. Also, the built-in DSQs (SCX_DSQ_LOCAL and
6009 * SCX_DSQ_GLOBAL) cannot be used as priority queues.
6010 *
6011 * Returns %true on successful insertion, %false on failure. On the root
6012 * scheduler, %false return triggers scheduler abort and the caller doesn't need
6013 * to check the return value.
6014 */
6015__bpf_kfunc bool
6016__scx_bpf_dsq_insert_vtime(struct task_struct *p,
6017			   struct scx_bpf_dsq_insert_vtime_args *args)
6018{
6019	struct scx_sched *sch;
6020
6021	guard(rcu)();
6022
6023	sch = rcu_dereference(scx_root);
6024	if (unlikely(!sch))
6025		return false;
6026
6027	return scx_dsq_insert_vtime(sch, p, args->dsq_id, args->slice,
6028				    args->vtime, args->enq_flags);
6029}
6030
6031/*
6032 * COMPAT: Will be removed in v6.23.
6033 */
6034__bpf_kfunc void scx_bpf_dsq_insert_vtime(struct task_struct *p, u64 dsq_id,
6035					  u64 slice, u64 vtime, u64 enq_flags)
6036{
6037	struct scx_sched *sch;
6038
6039	guard(rcu)();
6040
6041	sch = rcu_dereference(scx_root);
6042	if (unlikely(!sch))
6043		return;
6044
6045	scx_dsq_insert_vtime(sch, p, dsq_id, slice, vtime, enq_flags);
6046}
6047
6048__bpf_kfunc_end_defs();
6049
6050BTF_KFUNCS_START(scx_kfunc_ids_enqueue_dispatch)
6051BTF_ID_FLAGS(func, scx_bpf_dsq_insert, KF_RCU)
6052BTF_ID_FLAGS(func, scx_bpf_dsq_insert___v2, KF_RCU)
6053BTF_ID_FLAGS(func, __scx_bpf_dsq_insert_vtime, KF_RCU)
6054BTF_ID_FLAGS(func, scx_bpf_dsq_insert_vtime, KF_RCU)
6055BTF_KFUNCS_END(scx_kfunc_ids_enqueue_dispatch)
6056
6057static const struct btf_kfunc_id_set scx_kfunc_set_enqueue_dispatch = {
6058	.owner			= THIS_MODULE,
6059	.set			= &scx_kfunc_ids_enqueue_dispatch,
6060};
6061
6062static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
6063			 struct task_struct *p, u64 dsq_id, u64 enq_flags)
6064{
6065	struct scx_sched *sch = scx_root;
6066	struct scx_dispatch_q *src_dsq = kit->dsq, *dst_dsq;
6067	struct rq *this_rq, *src_rq, *locked_rq;
6068	bool dispatched = false;
6069	bool in_balance;
6070	unsigned long flags;
6071
6072	if (!scx_kf_allowed_if_unlocked() &&
6073	    !scx_kf_allowed(sch, SCX_KF_DISPATCH))
6074		return false;
6075
6076	/*
6077	 * If the BPF scheduler keeps calling this function repeatedly, it can
6078	 * cause similar live-lock conditions as consume_dispatch_q().
6079	 */
6080	if (unlikely(READ_ONCE(scx_aborting)))
6081		return false;
6082
6083	/*
6084	 * Can be called from either ops.dispatch() locking this_rq() or any
6085	 * context where no rq lock is held. If latter, lock @p's task_rq which
6086	 * we'll likely need anyway.
6087	 */
6088	src_rq = task_rq(p);
6089
6090	local_irq_save(flags);
6091	this_rq = this_rq();
6092	in_balance = this_rq->scx.flags & SCX_RQ_IN_BALANCE;
6093
6094	if (in_balance) {
6095		if (this_rq != src_rq) {
6096			raw_spin_rq_unlock(this_rq);
6097			raw_spin_rq_lock(src_rq);
6098		}
6099	} else {
6100		raw_spin_rq_lock(src_rq);
6101	}
6102
6103	locked_rq = src_rq;
6104	raw_spin_lock(&src_dsq->lock);
6105
6106	/*
6107	 * Did someone else get to it? @p could have already left $src_dsq, got
6108	 * re-enqueud, or be in the process of being consumed by someone else.
6109	 */
6110	if (unlikely(p->scx.dsq != src_dsq ||
6111		     u32_before(kit->cursor.priv, p->scx.dsq_seq) ||
6112		     p->scx.holding_cpu >= 0) ||
6113	    WARN_ON_ONCE(src_rq != task_rq(p))) {
6114		raw_spin_unlock(&src_dsq->lock);
6115		goto out;
6116	}
6117
6118	/* @p is still on $src_dsq and stable, determine the destination */
6119	dst_dsq = find_dsq_for_dispatch(sch, this_rq, dsq_id, p);
6120
6121	/*
6122	 * Apply vtime and slice updates before moving so that the new time is
6123	 * visible before inserting into $dst_dsq. @p is still on $src_dsq but
6124	 * this is safe as we're locking it.
6125	 */
6126	if (kit->cursor.flags & __SCX_DSQ_ITER_HAS_VTIME)
6127		p->scx.dsq_vtime = kit->vtime;
6128	if (kit->cursor.flags & __SCX_DSQ_ITER_HAS_SLICE)
6129		p->scx.slice = kit->slice;
6130
6131	/* execute move */
6132	locked_rq = move_task_between_dsqs(sch, p, enq_flags, src_dsq, dst_dsq);
6133	dispatched = true;
6134out:
6135	if (in_balance) {
6136		if (this_rq != locked_rq) {
6137			raw_spin_rq_unlock(locked_rq);
6138			raw_spin_rq_lock(this_rq);
6139		}
6140	} else {
6141		raw_spin_rq_unlock_irqrestore(locked_rq, flags);
6142	}
6143
6144	kit->cursor.flags &= ~(__SCX_DSQ_ITER_HAS_SLICE |
6145			       __SCX_DSQ_ITER_HAS_VTIME);
6146	return dispatched;
6147}
6148
6149__bpf_kfunc_start_defs();
6150
6151/**
6152 * scx_bpf_dispatch_nr_slots - Return the number of remaining dispatch slots
6153 *
6154 * Can only be called from ops.dispatch().
6155 */
6156__bpf_kfunc u32 scx_bpf_dispatch_nr_slots(void)
6157{
6158	struct scx_sched *sch;
6159
6160	guard(rcu)();
6161
6162	sch = rcu_dereference(scx_root);
6163	if (unlikely(!sch))
6164		return 0;
6165
6166	if (!scx_kf_allowed(sch, SCX_KF_DISPATCH))
6167		return 0;
6168
6169	return scx_dsp_max_batch - __this_cpu_read(scx_dsp_ctx->cursor);
6170}
6171
6172/**
6173 * scx_bpf_dispatch_cancel - Cancel the latest dispatch
6174 *
6175 * Cancel the latest dispatch. Can be called multiple times to cancel further
6176 * dispatches. Can only be called from ops.dispatch().
6177 */
6178__bpf_kfunc void scx_bpf_dispatch_cancel(void)
6179{
6180	struct scx_dsp_ctx *dspc = this_cpu_ptr(scx_dsp_ctx);
6181	struct scx_sched *sch;
6182
6183	guard(rcu)();
6184
6185	sch = rcu_dereference(scx_root);
6186	if (unlikely(!sch))
6187		return;
6188
6189	if (!scx_kf_allowed(sch, SCX_KF_DISPATCH))
6190		return;
6191
6192	if (dspc->cursor > 0)
6193		dspc->cursor--;
6194	else
6195		scx_error(sch, "dispatch buffer underflow");
6196}
6197
6198/**
6199 * scx_bpf_dsq_move_to_local - move a task from a DSQ to the current CPU's local DSQ
6200 * @dsq_id: DSQ to move task from
6201 *
6202 * Move a task from the non-local DSQ identified by @dsq_id to the current CPU's
6203 * local DSQ for execution. Can only be called from ops.dispatch().
6204 *
6205 * This function flushes the in-flight dispatches from scx_bpf_dsq_insert()
6206 * before trying to move from the specified DSQ. It may also grab rq locks and
6207 * thus can't be called under any BPF locks.
6208 *
6209 * Returns %true if a task has been moved, %false if there isn't any task to
6210 * move.
6211 */
6212__bpf_kfunc bool scx_bpf_dsq_move_to_local(u64 dsq_id)
6213{
6214	struct scx_dsp_ctx *dspc = this_cpu_ptr(scx_dsp_ctx);
6215	struct scx_dispatch_q *dsq;
6216	struct scx_sched *sch;
6217
6218	guard(rcu)();
6219
6220	sch = rcu_dereference(scx_root);
6221	if (unlikely(!sch))
6222		return false;
6223
6224	if (!scx_kf_allowed(sch, SCX_KF_DISPATCH))
6225		return false;
6226
6227	flush_dispatch_buf(sch, dspc->rq);
6228
6229	dsq = find_user_dsq(sch, dsq_id);
6230	if (unlikely(!dsq)) {
6231		scx_error(sch, "invalid DSQ ID 0x%016llx", dsq_id);
6232		return false;
6233	}
6234
6235	if (consume_dispatch_q(sch, dspc->rq, dsq)) {
6236		/*
6237		 * A successfully consumed task can be dequeued before it starts
6238		 * running while the CPU is trying to migrate other dispatched
6239		 * tasks. Bump nr_tasks to tell balance_one() to retry on empty
6240		 * local DSQ.
6241		 */
6242		dspc->nr_tasks++;
6243		return true;
6244	} else {
6245		return false;
6246	}
6247}
6248
6249/**
6250 * scx_bpf_dsq_move_set_slice - Override slice when moving between DSQs
6251 * @it__iter: DSQ iterator in progress
6252 * @slice: duration the moved task can run for in nsecs
6253 *
6254 * Override the slice of the next task that will be moved from @it__iter using
6255 * scx_bpf_dsq_move[_vtime](). If this function is not called, the previous
6256 * slice duration is kept.
6257 */
6258__bpf_kfunc void scx_bpf_dsq_move_set_slice(struct bpf_iter_scx_dsq *it__iter,
6259					    u64 slice)
6260{
6261	struct bpf_iter_scx_dsq_kern *kit = (void *)it__iter;
6262
6263	kit->slice = slice;
6264	kit->cursor.flags |= __SCX_DSQ_ITER_HAS_SLICE;
6265}
6266
6267/**
6268 * scx_bpf_dsq_move_set_vtime - Override vtime when moving between DSQs
6269 * @it__iter: DSQ iterator in progress
6270 * @vtime: task's ordering inside the vtime-sorted queue of the target DSQ
6271 *
6272 * Override the vtime of the next task that will be moved from @it__iter using
6273 * scx_bpf_dsq_move_vtime(). If this function is not called, the previous slice
6274 * vtime is kept. If scx_bpf_dsq_move() is used to dispatch the next task, the
6275 * override is ignored and cleared.
6276 */
6277__bpf_kfunc void scx_bpf_dsq_move_set_vtime(struct bpf_iter_scx_dsq *it__iter,
6278					    u64 vtime)
6279{
6280	struct bpf_iter_scx_dsq_kern *kit = (void *)it__iter;
6281
6282	kit->vtime = vtime;
6283	kit->cursor.flags |= __SCX_DSQ_ITER_HAS_VTIME;
6284}
6285
6286/**
6287 * scx_bpf_dsq_move - Move a task from DSQ iteration to a DSQ
6288 * @it__iter: DSQ iterator in progress
6289 * @p: task to transfer
6290 * @dsq_id: DSQ to move @p to
6291 * @enq_flags: SCX_ENQ_*
6292 *
6293 * Transfer @p which is on the DSQ currently iterated by @it__iter to the DSQ
6294 * specified by @dsq_id. All DSQs - local DSQs, global DSQ and user DSQs - can
6295 * be the destination.
6296 *
6297 * For the transfer to be successful, @p must still be on the DSQ and have been
6298 * queued before the DSQ iteration started. This function doesn't care whether
6299 * @p was obtained from the DSQ iteration. @p just has to be on the DSQ and have
6300 * been queued before the iteration started.
6301 *
6302 * @p's slice is kept by default. Use scx_bpf_dsq_move_set_slice() to update.
6303 *
6304 * Can be called from ops.dispatch() or any BPF context which doesn't hold a rq
6305 * lock (e.g. BPF timers or SYSCALL programs).
6306 *
6307 * Returns %true if @p has been consumed, %false if @p had already been
6308 * consumed, dequeued, or, for sub-scheds, @dsq_id points to a disallowed local
6309 * DSQ.
6310 */
6311__bpf_kfunc bool scx_bpf_dsq_move(struct bpf_iter_scx_dsq *it__iter,
6312				  struct task_struct *p, u64 dsq_id,
6313				  u64 enq_flags)
6314{
6315	return scx_dsq_move((struct bpf_iter_scx_dsq_kern *)it__iter,
6316			    p, dsq_id, enq_flags);
6317}
6318
6319/**
6320 * scx_bpf_dsq_move_vtime - Move a task from DSQ iteration to a PRIQ DSQ
6321 * @it__iter: DSQ iterator in progress
6322 * @p: task to transfer
6323 * @dsq_id: DSQ to move @p to
6324 * @enq_flags: SCX_ENQ_*
6325 *
6326 * Transfer @p which is on the DSQ currently iterated by @it__iter to the
6327 * priority queue of the DSQ specified by @dsq_id. The destination must be a
6328 * user DSQ as only user DSQs support priority queue.
6329 *
6330 * @p's slice and vtime are kept by default. Use scx_bpf_dsq_move_set_slice()
6331 * and scx_bpf_dsq_move_set_vtime() to update.
6332 *
6333 * All other aspects are identical to scx_bpf_dsq_move(). See
6334 * scx_bpf_dsq_insert_vtime() for more information on @vtime.
6335 */
6336__bpf_kfunc bool scx_bpf_dsq_move_vtime(struct bpf_iter_scx_dsq *it__iter,
6337					struct task_struct *p, u64 dsq_id,
6338					u64 enq_flags)
6339{
6340	return scx_dsq_move((struct bpf_iter_scx_dsq_kern *)it__iter,
6341			    p, dsq_id, enq_flags | SCX_ENQ_DSQ_PRIQ);
6342}
6343
6344__bpf_kfunc_end_defs();
6345
6346BTF_KFUNCS_START(scx_kfunc_ids_dispatch)
6347BTF_ID_FLAGS(func, scx_bpf_dispatch_nr_slots)
6348BTF_ID_FLAGS(func, scx_bpf_dispatch_cancel)
6349BTF_ID_FLAGS(func, scx_bpf_dsq_move_to_local)
6350BTF_ID_FLAGS(func, scx_bpf_dsq_move_set_slice, KF_RCU)
6351BTF_ID_FLAGS(func, scx_bpf_dsq_move_set_vtime, KF_RCU)
6352BTF_ID_FLAGS(func, scx_bpf_dsq_move, KF_RCU)
6353BTF_ID_FLAGS(func, scx_bpf_dsq_move_vtime, KF_RCU)
6354BTF_KFUNCS_END(scx_kfunc_ids_dispatch)
6355
6356static const struct btf_kfunc_id_set scx_kfunc_set_dispatch = {
6357	.owner			= THIS_MODULE,
6358	.set			= &scx_kfunc_ids_dispatch,
6359};
6360
6361static u32 reenq_local(struct rq *rq)
6362{
6363	LIST_HEAD(tasks);
6364	u32 nr_enqueued = 0;
6365	struct task_struct *p, *n;
6366
6367	lockdep_assert_rq_held(rq);
6368
6369	/*
6370	 * The BPF scheduler may choose to dispatch tasks back to
6371	 * @rq->scx.local_dsq. Move all candidate tasks off to a private list
6372	 * first to avoid processing the same tasks repeatedly.
6373	 */
6374	list_for_each_entry_safe(p, n, &rq->scx.local_dsq.list,
6375				 scx.dsq_list.node) {
6376		/*
6377		 * If @p is being migrated, @p's current CPU may not agree with
6378		 * its allowed CPUs and the migration_cpu_stop is about to
6379		 * deactivate and re-activate @p anyway. Skip re-enqueueing.
6380		 *
6381		 * While racing sched property changes may also dequeue and
6382		 * re-enqueue a migrating task while its current CPU and allowed
6383		 * CPUs disagree, they use %ENQUEUE_RESTORE which is bypassed to
6384		 * the current local DSQ for running tasks and thus are not
6385		 * visible to the BPF scheduler.
6386		 */
6387		if (p->migration_pending)
6388			continue;
6389
6390		dispatch_dequeue(rq, p);
6391		list_add_tail(&p->scx.dsq_list.node, &tasks);
6392	}
6393
6394	list_for_each_entry_safe(p, n, &tasks, scx.dsq_list.node) {
6395		list_del_init(&p->scx.dsq_list.node);
6396		do_enqueue_task(rq, p, SCX_ENQ_REENQ, -1);
6397		nr_enqueued++;
6398	}
6399
6400	return nr_enqueued;
6401}
6402
6403__bpf_kfunc_start_defs();
6404
6405/**
6406 * scx_bpf_reenqueue_local - Re-enqueue tasks on a local DSQ
6407 *
6408 * Iterate over all of the tasks currently enqueued on the local DSQ of the
6409 * caller's CPU, and re-enqueue them in the BPF scheduler. Returns the number of
6410 * processed tasks. Can only be called from ops.cpu_release().
6411 *
6412 * COMPAT: Will be removed in v6.23 along with the ___v2 suffix on the void
6413 * returning variant that can be called from anywhere.
6414 */
6415__bpf_kfunc u32 scx_bpf_reenqueue_local(void)
6416{
6417	struct scx_sched *sch;
6418	struct rq *rq;
6419
6420	guard(rcu)();
6421	sch = rcu_dereference(scx_root);
6422	if (unlikely(!sch))
6423		return 0;
6424
6425	if (!scx_kf_allowed(sch, SCX_KF_CPU_RELEASE))
6426		return 0;
6427
6428	rq = cpu_rq(smp_processor_id());
6429	lockdep_assert_rq_held(rq);
6430
6431	return reenq_local(rq);
6432}
6433
6434__bpf_kfunc_end_defs();
6435
6436BTF_KFUNCS_START(scx_kfunc_ids_cpu_release)
6437BTF_ID_FLAGS(func, scx_bpf_reenqueue_local)
6438BTF_KFUNCS_END(scx_kfunc_ids_cpu_release)
6439
6440static const struct btf_kfunc_id_set scx_kfunc_set_cpu_release = {
6441	.owner			= THIS_MODULE,
6442	.set			= &scx_kfunc_ids_cpu_release,
6443};
6444
6445__bpf_kfunc_start_defs();
6446
6447/**
6448 * scx_bpf_create_dsq - Create a custom DSQ
6449 * @dsq_id: DSQ to create
6450 * @node: NUMA node to allocate from
6451 *
6452 * Create a custom DSQ identified by @dsq_id. Can be called from any sleepable
6453 * scx callback, and any BPF_PROG_TYPE_SYSCALL prog.
6454 */
6455__bpf_kfunc s32 scx_bpf_create_dsq(u64 dsq_id, s32 node)
6456{
6457	struct scx_dispatch_q *dsq;
6458	struct scx_sched *sch;
6459	s32 ret;
6460
6461	if (unlikely(node >= (int)nr_node_ids ||
6462		     (node < 0 && node != NUMA_NO_NODE)))
6463		return -EINVAL;
6464
6465	if (unlikely(dsq_id & SCX_DSQ_FLAG_BUILTIN))
6466		return -EINVAL;
6467
6468	dsq = kmalloc_node(sizeof(*dsq), GFP_KERNEL, node);
6469	if (!dsq)
6470		return -ENOMEM;
6471
6472	init_dsq(dsq, dsq_id);
6473
6474	rcu_read_lock();
6475
6476	sch = rcu_dereference(scx_root);
6477	if (sch)
6478		ret = rhashtable_lookup_insert_fast(&sch->dsq_hash, &dsq->hash_node,
6479						    dsq_hash_params);
6480	else
6481		ret = -ENODEV;
6482
6483	rcu_read_unlock();
6484	if (ret)
6485		kfree(dsq);
6486	return ret;
6487}
6488
6489__bpf_kfunc_end_defs();
6490
6491BTF_KFUNCS_START(scx_kfunc_ids_unlocked)
6492BTF_ID_FLAGS(func, scx_bpf_create_dsq, KF_SLEEPABLE)
6493BTF_ID_FLAGS(func, scx_bpf_dsq_move_set_slice, KF_RCU)
6494BTF_ID_FLAGS(func, scx_bpf_dsq_move_set_vtime, KF_RCU)
6495BTF_ID_FLAGS(func, scx_bpf_dsq_move, KF_RCU)
6496BTF_ID_FLAGS(func, scx_bpf_dsq_move_vtime, KF_RCU)
6497BTF_KFUNCS_END(scx_kfunc_ids_unlocked)
6498
6499static const struct btf_kfunc_id_set scx_kfunc_set_unlocked = {
6500	.owner			= THIS_MODULE,
6501	.set			= &scx_kfunc_ids_unlocked,
6502};
6503
6504__bpf_kfunc_start_defs();
6505
6506/**
6507 * scx_bpf_task_set_slice - Set task's time slice
6508 * @p: task of interest
6509 * @slice: time slice to set in nsecs
6510 *
6511 * Set @p's time slice to @slice. Returns %true on success, %false if the
6512 * calling scheduler doesn't have authority over @p.
6513 */
6514__bpf_kfunc bool scx_bpf_task_set_slice(struct task_struct *p, u64 slice)
6515{
6516	p->scx.slice = slice;
6517	return true;
6518}
6519
6520/**
6521 * scx_bpf_task_set_dsq_vtime - Set task's virtual time for DSQ ordering
6522 * @p: task of interest
6523 * @vtime: virtual time to set
6524 *
6525 * Set @p's virtual time to @vtime. Returns %true on success, %false if the
6526 * calling scheduler doesn't have authority over @p.
6527 */
6528__bpf_kfunc bool scx_bpf_task_set_dsq_vtime(struct task_struct *p, u64 vtime)
6529{
6530	p->scx.dsq_vtime = vtime;
6531	return true;
6532}
6533
6534static void scx_kick_cpu(struct scx_sched *sch, s32 cpu, u64 flags)
6535{
6536	struct rq *this_rq;
6537	unsigned long irq_flags;
6538
6539	if (!ops_cpu_valid(sch, cpu, NULL))
6540		return;
6541
6542	local_irq_save(irq_flags);
6543
6544	this_rq = this_rq();
6545
6546	/*
6547	 * While bypassing for PM ops, IRQ handling may not be online which can
6548	 * lead to irq_work_queue() malfunction such as infinite busy wait for
6549	 * IRQ status update. Suppress kicking.
6550	 */
6551	if (scx_rq_bypassing(this_rq))
6552		goto out;
6553
6554	/*
6555	 * Actual kicking is bounced to kick_cpus_irq_workfn() to avoid nesting
6556	 * rq locks. We can probably be smarter and avoid bouncing if called
6557	 * from ops which don't hold a rq lock.
6558	 */
6559	if (flags & SCX_KICK_IDLE) {
6560		struct rq *target_rq = cpu_rq(cpu);
6561
6562		if (unlikely(flags & (SCX_KICK_PREEMPT | SCX_KICK_WAIT)))
6563			scx_error(sch, "PREEMPT/WAIT cannot be used with SCX_KICK_IDLE");
6564
6565		if (raw_spin_rq_trylock(target_rq)) {
6566			if (can_skip_idle_kick(target_rq)) {
6567				raw_spin_rq_unlock(target_rq);
6568				goto out;
6569			}
6570			raw_spin_rq_unlock(target_rq);
6571		}
6572		cpumask_set_cpu(cpu, this_rq->scx.cpus_to_kick_if_idle);
6573	} else {
6574		cpumask_set_cpu(cpu, this_rq->scx.cpus_to_kick);
6575
6576		if (flags & SCX_KICK_PREEMPT)
6577			cpumask_set_cpu(cpu, this_rq->scx.cpus_to_preempt);
6578		if (flags & SCX_KICK_WAIT)
6579			cpumask_set_cpu(cpu, this_rq->scx.cpus_to_wait);
6580	}
6581
6582	irq_work_queue(&this_rq->scx.kick_cpus_irq_work);
6583out:
6584	local_irq_restore(irq_flags);
6585}
6586
6587/**
6588 * scx_bpf_kick_cpu - Trigger reschedule on a CPU
6589 * @cpu: cpu to kick
6590 * @flags: %SCX_KICK_* flags
6591 *
6592 * Kick @cpu into rescheduling. This can be used to wake up an idle CPU or
6593 * trigger rescheduling on a busy CPU. This can be called from any online
6594 * scx_ops operation and the actual kicking is performed asynchronously through
6595 * an irq work.
6596 */
6597__bpf_kfunc void scx_bpf_kick_cpu(s32 cpu, u64 flags)
6598{
6599	struct scx_sched *sch;
6600
6601	guard(rcu)();
6602	sch = rcu_dereference(scx_root);
6603	if (likely(sch))
6604		scx_kick_cpu(sch, cpu, flags);
6605}
6606
6607/**
6608 * scx_bpf_dsq_nr_queued - Return the number of queued tasks
6609 * @dsq_id: id of the DSQ
6610 *
6611 * Return the number of tasks in the DSQ matching @dsq_id. If not found,
6612 * -%ENOENT is returned.
6613 */
6614__bpf_kfunc s32 scx_bpf_dsq_nr_queued(u64 dsq_id)
6615{
6616	struct scx_sched *sch;
6617	struct scx_dispatch_q *dsq;
6618	s32 ret;
6619
6620	preempt_disable();
6621
6622	sch = rcu_dereference_sched(scx_root);
6623	if (unlikely(!sch)) {
6624		ret = -ENODEV;
6625		goto out;
6626	}
6627
6628	if (dsq_id == SCX_DSQ_LOCAL) {
6629		ret = READ_ONCE(this_rq()->scx.local_dsq.nr);
6630		goto out;
6631	} else if ((dsq_id & SCX_DSQ_LOCAL_ON) == SCX_DSQ_LOCAL_ON) {
6632		s32 cpu = dsq_id & SCX_DSQ_LOCAL_CPU_MASK;
6633
6634		if (ops_cpu_valid(sch, cpu, NULL)) {
6635			ret = READ_ONCE(cpu_rq(cpu)->scx.local_dsq.nr);
6636			goto out;
6637		}
6638	} else {
6639		dsq = find_user_dsq(sch, dsq_id);
6640		if (dsq) {
6641			ret = READ_ONCE(dsq->nr);
6642			goto out;
6643		}
6644	}
6645	ret = -ENOENT;
6646out:
6647	preempt_enable();
6648	return ret;
6649}
6650
6651/**
6652 * scx_bpf_destroy_dsq - Destroy a custom DSQ
6653 * @dsq_id: DSQ to destroy
6654 *
6655 * Destroy the custom DSQ identified by @dsq_id. Only DSQs created with
6656 * scx_bpf_create_dsq() can be destroyed. The caller must ensure that the DSQ is
6657 * empty and no further tasks are dispatched to it. Ignored if called on a DSQ
6658 * which doesn't exist. Can be called from any online scx_ops operations.
6659 */
6660__bpf_kfunc void scx_bpf_destroy_dsq(u64 dsq_id)
6661{
6662	struct scx_sched *sch;
6663
6664	rcu_read_lock();
6665	sch = rcu_dereference(scx_root);
6666	if (sch)
6667		destroy_dsq(sch, dsq_id);
6668	rcu_read_unlock();
6669}
6670
6671/**
6672 * bpf_iter_scx_dsq_new - Create a DSQ iterator
6673 * @it: iterator to initialize
6674 * @dsq_id: DSQ to iterate
6675 * @flags: %SCX_DSQ_ITER_*
6676 *
6677 * Initialize BPF iterator @it which can be used with bpf_for_each() to walk
6678 * tasks in the DSQ specified by @dsq_id. Iteration using @it only includes
6679 * tasks which are already queued when this function is invoked.
6680 */
6681__bpf_kfunc int bpf_iter_scx_dsq_new(struct bpf_iter_scx_dsq *it, u64 dsq_id,
6682				     u64 flags)
6683{
6684	struct bpf_iter_scx_dsq_kern *kit = (void *)it;
6685	struct scx_sched *sch;
6686
6687	BUILD_BUG_ON(sizeof(struct bpf_iter_scx_dsq_kern) >
6688		     sizeof(struct bpf_iter_scx_dsq));
6689	BUILD_BUG_ON(__alignof__(struct bpf_iter_scx_dsq_kern) !=
6690		     __alignof__(struct bpf_iter_scx_dsq));
6691	BUILD_BUG_ON(__SCX_DSQ_ITER_ALL_FLAGS &
6692		     ((1U << __SCX_DSQ_LNODE_PRIV_SHIFT) - 1));
6693
6694	/*
6695	 * next() and destroy() will be called regardless of the return value.
6696	 * Always clear $kit->dsq.
6697	 */
6698	kit->dsq = NULL;
6699
6700	sch = rcu_dereference_check(scx_root, rcu_read_lock_bh_held());
6701	if (unlikely(!sch))
6702		return -ENODEV;
6703
6704	if (flags & ~__SCX_DSQ_ITER_USER_FLAGS)
6705		return -EINVAL;
6706
6707	kit->dsq = find_user_dsq(sch, dsq_id);
6708	if (!kit->dsq)
6709		return -ENOENT;
6710
6711	kit->cursor = INIT_DSQ_LIST_CURSOR(kit->cursor, flags,
6712					   READ_ONCE(kit->dsq->seq));
6713
6714	return 0;
6715}
6716
6717/**
6718 * bpf_iter_scx_dsq_next - Progress a DSQ iterator
6719 * @it: iterator to progress
6720 *
6721 * Return the next task. See bpf_iter_scx_dsq_new().
6722 */
6723__bpf_kfunc struct task_struct *bpf_iter_scx_dsq_next(struct bpf_iter_scx_dsq *it)
6724{
6725	struct bpf_iter_scx_dsq_kern *kit = (void *)it;
6726	bool rev = kit->cursor.flags & SCX_DSQ_ITER_REV;
6727	struct task_struct *p;
6728	unsigned long flags;
6729
6730	if (!kit->dsq)
6731		return NULL;
6732
6733	raw_spin_lock_irqsave(&kit->dsq->lock, flags);
6734
6735	if (list_empty(&kit->cursor.node))
6736		p = NULL;
6737	else
6738		p = container_of(&kit->cursor, struct task_struct, scx.dsq_list);
6739
6740	/*
6741	 * Only tasks which were queued before the iteration started are
6742	 * visible. This bounds BPF iterations and guarantees that vtime never
6743	 * jumps in the other direction while iterating.
6744	 */
6745	do {
6746		p = nldsq_next_task(kit->dsq, p, rev);
6747	} while (p && unlikely(u32_before(kit->cursor.priv, p->scx.dsq_seq)));
6748
6749	if (p) {
6750		if (rev)
6751			list_move_tail(&kit->cursor.node, &p->scx.dsq_list.node);
6752		else
6753			list_move(&kit->cursor.node, &p->scx.dsq_list.node);
6754	} else {
6755		list_del_init(&kit->cursor.node);
6756	}
6757
6758	raw_spin_unlock_irqrestore(&kit->dsq->lock, flags);
6759
6760	return p;
6761}
6762
6763/**
6764 * bpf_iter_scx_dsq_destroy - Destroy a DSQ iterator
6765 * @it: iterator to destroy
6766 *
6767 * Undo scx_iter_scx_dsq_new().
6768 */
6769__bpf_kfunc void bpf_iter_scx_dsq_destroy(struct bpf_iter_scx_dsq *it)
6770{
6771	struct bpf_iter_scx_dsq_kern *kit = (void *)it;
6772
6773	if (!kit->dsq)
6774		return;
6775
6776	if (!list_empty(&kit->cursor.node)) {
6777		unsigned long flags;
6778
6779		raw_spin_lock_irqsave(&kit->dsq->lock, flags);
6780		list_del_init(&kit->cursor.node);
6781		raw_spin_unlock_irqrestore(&kit->dsq->lock, flags);
6782	}
6783	kit->dsq = NULL;
6784}
6785
6786/**
6787 * scx_bpf_dsq_peek - Lockless peek at the first element.
6788 * @dsq_id: DSQ to examine.
6789 *
6790 * Read the first element in the DSQ. This is semantically equivalent to using
6791 * the DSQ iterator, but is lockfree. Of course, like any lockless operation,
6792 * this provides only a point-in-time snapshot, and the contents may change
6793 * by the time any subsequent locking operation reads the queue.
6794 *
6795 * Returns the pointer, or NULL indicates an empty queue OR internal error.
6796 */
6797__bpf_kfunc struct task_struct *scx_bpf_dsq_peek(u64 dsq_id)
6798{
6799	struct scx_sched *sch;
6800	struct scx_dispatch_q *dsq;
6801
6802	sch = rcu_dereference(scx_root);
6803	if (unlikely(!sch))
6804		return NULL;
6805
6806	if (unlikely(dsq_id & SCX_DSQ_FLAG_BUILTIN)) {
6807		scx_error(sch, "peek disallowed on builtin DSQ 0x%llx", dsq_id);
6808		return NULL;
6809	}
6810
6811	dsq = find_user_dsq(sch, dsq_id);
6812	if (unlikely(!dsq)) {
6813		scx_error(sch, "peek on non-existent DSQ 0x%llx", dsq_id);
6814		return NULL;
6815	}
6816
6817	return rcu_dereference(dsq->first_task);
6818}
6819
6820__bpf_kfunc_end_defs();
6821
6822static s32 __bstr_format(struct scx_sched *sch, u64 *data_buf, char *line_buf,
6823			 size_t line_size, char *fmt, unsigned long long *data,
6824			 u32 data__sz)
6825{
6826	struct bpf_bprintf_data bprintf_data = { .get_bin_args = true };
6827	s32 ret;
6828
6829	if (data__sz % 8 || data__sz > MAX_BPRINTF_VARARGS * 8 ||
6830	    (data__sz && !data)) {
6831		scx_error(sch, "invalid data=%p and data__sz=%u", (void *)data, data__sz);
6832		return -EINVAL;
6833	}
6834
6835	ret = copy_from_kernel_nofault(data_buf, data, data__sz);
6836	if (ret < 0) {
6837		scx_error(sch, "failed to read data fields (%d)", ret);
6838		return ret;
6839	}
6840
6841	ret = bpf_bprintf_prepare(fmt, UINT_MAX, data_buf, data__sz / 8,
6842				  &bprintf_data);
6843	if (ret < 0) {
6844		scx_error(sch, "format preparation failed (%d)", ret);
6845		return ret;
6846	}
6847
6848	ret = bstr_printf(line_buf, line_size, fmt,
6849			  bprintf_data.bin_args);
6850	bpf_bprintf_cleanup(&bprintf_data);
6851	if (ret < 0) {
6852		scx_error(sch, "(\"%s\", %p, %u) failed to format", fmt, data, data__sz);
6853		return ret;
6854	}
6855
6856	return ret;
6857}
6858
6859static s32 bstr_format(struct scx_sched *sch, struct scx_bstr_buf *buf,
6860		       char *fmt, unsigned long long *data, u32 data__sz)
6861{
6862	return __bstr_format(sch, buf->data, buf->line, sizeof(buf->line),
6863			     fmt, data, data__sz);
6864}
6865
6866__bpf_kfunc_start_defs();
6867
6868/**
6869 * scx_bpf_exit_bstr - Gracefully exit the BPF scheduler.
6870 * @exit_code: Exit value to pass to user space via struct scx_exit_info.
6871 * @fmt: error message format string
6872 * @data: format string parameters packaged using ___bpf_fill() macro
6873 * @data__sz: @data len, must end in '__sz' for the verifier
6874 *
6875 * Indicate that the BPF scheduler wants to exit gracefully, and initiate ops
6876 * disabling.
6877 */
6878__bpf_kfunc void scx_bpf_exit_bstr(s64 exit_code, char *fmt,
6879				   unsigned long long *data, u32 data__sz)
6880{
6881	struct scx_sched *sch;
6882	unsigned long flags;
6883
6884	raw_spin_lock_irqsave(&scx_exit_bstr_buf_lock, flags);
6885	sch = rcu_dereference_bh(scx_root);
6886	if (likely(sch) &&
6887	    bstr_format(sch, &scx_exit_bstr_buf, fmt, data, data__sz) >= 0)
6888		scx_exit(sch, SCX_EXIT_UNREG_BPF, exit_code, "%s", scx_exit_bstr_buf.line);
6889	raw_spin_unlock_irqrestore(&scx_exit_bstr_buf_lock, flags);
6890}
6891
6892/**
6893 * scx_bpf_error_bstr - Indicate fatal error
6894 * @fmt: error message format string
6895 * @data: format string parameters packaged using ___bpf_fill() macro
6896 * @data__sz: @data len, must end in '__sz' for the verifier
6897 *
6898 * Indicate that the BPF scheduler encountered a fatal error and initiate ops
6899 * disabling.
6900 */
6901__bpf_kfunc void scx_bpf_error_bstr(char *fmt, unsigned long long *data,
6902				    u32 data__sz)
6903{
6904	struct scx_sched *sch;
6905	unsigned long flags;
6906
6907	raw_spin_lock_irqsave(&scx_exit_bstr_buf_lock, flags);
6908	sch = rcu_dereference_bh(scx_root);
6909	if (likely(sch) &&
6910	    bstr_format(sch, &scx_exit_bstr_buf, fmt, data, data__sz) >= 0)
6911		scx_exit(sch, SCX_EXIT_ERROR_BPF, 0, "%s", scx_exit_bstr_buf.line);
6912	raw_spin_unlock_irqrestore(&scx_exit_bstr_buf_lock, flags);
6913}
6914
6915/**
6916 * scx_bpf_dump_bstr - Generate extra debug dump specific to the BPF scheduler
6917 * @fmt: format string
6918 * @data: format string parameters packaged using ___bpf_fill() macro
6919 * @data__sz: @data len, must end in '__sz' for the verifier
6920 *
6921 * To be called through scx_bpf_dump() helper from ops.dump(), dump_cpu() and
6922 * dump_task() to generate extra debug dump specific to the BPF scheduler.
6923 *
6924 * The extra dump may be multiple lines. A single line may be split over
6925 * multiple calls. The last line is automatically terminated.
6926 */
6927__bpf_kfunc void scx_bpf_dump_bstr(char *fmt, unsigned long long *data,
6928				   u32 data__sz)
6929{
6930	struct scx_sched *sch;
6931	struct scx_dump_data *dd = &scx_dump_data;
6932	struct scx_bstr_buf *buf = &dd->buf;
6933	s32 ret;
6934
6935	guard(rcu)();
6936
6937	sch = rcu_dereference(scx_root);
6938	if (unlikely(!sch))
6939		return;
6940
6941	if (raw_smp_processor_id() != dd->cpu) {
6942		scx_error(sch, "scx_bpf_dump() must only be called from ops.dump() and friends");
6943		return;
6944	}
6945
6946	/* append the formatted string to the line buf */
6947	ret = __bstr_format(sch, buf->data, buf->line + dd->cursor,
6948			    sizeof(buf->line) - dd->cursor, fmt, data, data__sz);
6949	if (ret < 0) {
6950		dump_line(dd->s, "%s[!] (\"%s\", %p, %u) failed to format (%d)",
6951			  dd->prefix, fmt, data, data__sz, ret);
6952		return;
6953	}
6954
6955	dd->cursor += ret;
6956	dd->cursor = min_t(s32, dd->cursor, sizeof(buf->line));
6957
6958	if (!dd->cursor)
6959		return;
6960
6961	/*
6962	 * If the line buf overflowed or ends in a newline, flush it into the
6963	 * dump. This is to allow the caller to generate a single line over
6964	 * multiple calls. As ops_dump_flush() can also handle multiple lines in
6965	 * the line buf, the only case which can lead to an unexpected
6966	 * truncation is when the caller keeps generating newlines in the middle
6967	 * instead of the end consecutively. Don't do that.
6968	 */
6969	if (dd->cursor >= sizeof(buf->line) || buf->line[dd->cursor - 1] == '\n')
6970		ops_dump_flush();
6971}
6972
6973/**
6974 * scx_bpf_reenqueue_local - Re-enqueue tasks on a local DSQ
6975 *
6976 * Iterate over all of the tasks currently enqueued on the local DSQ of the
6977 * caller's CPU, and re-enqueue them in the BPF scheduler. Can be called from
6978 * anywhere.
6979 */
6980__bpf_kfunc void scx_bpf_reenqueue_local___v2(void)
6981{
6982	struct rq *rq;
6983
6984	guard(preempt)();
6985
6986	rq = this_rq();
6987	local_set(&rq->scx.reenq_local_deferred, 1);
6988	schedule_deferred(rq);
6989}
6990
6991/**
6992 * scx_bpf_cpuperf_cap - Query the maximum relative capacity of a CPU
6993 * @cpu: CPU of interest
6994 *
6995 * Return the maximum relative capacity of @cpu in relation to the most
6996 * performant CPU in the system. The return value is in the range [1,
6997 * %SCX_CPUPERF_ONE]. See scx_bpf_cpuperf_cur().
6998 */
6999__bpf_kfunc u32 scx_bpf_cpuperf_cap(s32 cpu)
7000{
7001	struct scx_sched *sch;
7002
7003	guard(rcu)();
7004
7005	sch = rcu_dereference(scx_root);
7006	if (likely(sch) && ops_cpu_valid(sch, cpu, NULL))
7007		return arch_scale_cpu_capacity(cpu);
7008	else
7009		return SCX_CPUPERF_ONE;
7010}
7011
7012/**
7013 * scx_bpf_cpuperf_cur - Query the current relative performance of a CPU
7014 * @cpu: CPU of interest
7015 *
7016 * Return the current relative performance of @cpu in relation to its maximum.
7017 * The return value is in the range [1, %SCX_CPUPERF_ONE].
7018 *
7019 * The current performance level of a CPU in relation to the maximum performance
7020 * available in the system can be calculated as follows:
7021 *
7022 *   scx_bpf_cpuperf_cap() * scx_bpf_cpuperf_cur() / %SCX_CPUPERF_ONE
7023 *
7024 * The result is in the range [1, %SCX_CPUPERF_ONE].
7025 */
7026__bpf_kfunc u32 scx_bpf_cpuperf_cur(s32 cpu)
7027{
7028	struct scx_sched *sch;
7029
7030	guard(rcu)();
7031
7032	sch = rcu_dereference(scx_root);
7033	if (likely(sch) && ops_cpu_valid(sch, cpu, NULL))
7034		return arch_scale_freq_capacity(cpu);
7035	else
7036		return SCX_CPUPERF_ONE;
7037}
7038
7039/**
7040 * scx_bpf_cpuperf_set - Set the relative performance target of a CPU
7041 * @cpu: CPU of interest
7042 * @perf: target performance level [0, %SCX_CPUPERF_ONE]
7043 *
7044 * Set the target performance level of @cpu to @perf. @perf is in linear
7045 * relative scale between 0 and %SCX_CPUPERF_ONE. This determines how the
7046 * schedutil cpufreq governor chooses the target frequency.
7047 *
7048 * The actual performance level chosen, CPU grouping, and the overhead and
7049 * latency of the operations are dependent on the hardware and cpufreq driver in
7050 * use. Consult hardware and cpufreq documentation for more information. The
7051 * current performance level can be monitored using scx_bpf_cpuperf_cur().
7052 */
7053__bpf_kfunc void scx_bpf_cpuperf_set(s32 cpu, u32 perf)
7054{
7055	struct scx_sched *sch;
7056
7057	guard(rcu)();
7058
7059	sch = rcu_dereference(scx_root);
7060	if (unlikely(!sch))
7061		return;
7062
7063	if (unlikely(perf > SCX_CPUPERF_ONE)) {
7064		scx_error(sch, "Invalid cpuperf target %u for CPU %d", perf, cpu);
7065		return;
7066	}
7067
7068	if (ops_cpu_valid(sch, cpu, NULL)) {
7069		struct rq *rq = cpu_rq(cpu), *locked_rq = scx_locked_rq();
7070		struct rq_flags rf;
7071
7072		/*
7073		 * When called with an rq lock held, restrict the operation
7074		 * to the corresponding CPU to prevent ABBA deadlocks.
7075		 */
7076		if (locked_rq && rq != locked_rq) {
7077			scx_error(sch, "Invalid target CPU %d", cpu);
7078			return;
7079		}
7080
7081		/*
7082		 * If no rq lock is held, allow to operate on any CPU by
7083		 * acquiring the corresponding rq lock.
7084		 */
7085		if (!locked_rq) {
7086			rq_lock_irqsave(rq, &rf);
7087			update_rq_clock(rq);
7088		}
7089
7090		rq->scx.cpuperf_target = perf;
7091		cpufreq_update_util(rq, 0);
7092
7093		if (!locked_rq)
7094			rq_unlock_irqrestore(rq, &rf);
7095	}
7096}
7097
7098/**
7099 * scx_bpf_nr_node_ids - Return the number of possible node IDs
7100 *
7101 * All valid node IDs in the system are smaller than the returned value.
7102 */
7103__bpf_kfunc u32 scx_bpf_nr_node_ids(void)
7104{
7105	return nr_node_ids;
7106}
7107
7108/**
7109 * scx_bpf_nr_cpu_ids - Return the number of possible CPU IDs
7110 *
7111 * All valid CPU IDs in the system are smaller than the returned value.
7112 */
7113__bpf_kfunc u32 scx_bpf_nr_cpu_ids(void)
7114{
7115	return nr_cpu_ids;
7116}
7117
7118/**
7119 * scx_bpf_get_possible_cpumask - Get a referenced kptr to cpu_possible_mask
7120 */
7121__bpf_kfunc const struct cpumask *scx_bpf_get_possible_cpumask(void)
7122{
7123	return cpu_possible_mask;
7124}
7125
7126/**
7127 * scx_bpf_get_online_cpumask - Get a referenced kptr to cpu_online_mask
7128 */
7129__bpf_kfunc const struct cpumask *scx_bpf_get_online_cpumask(void)
7130{
7131	return cpu_online_mask;
7132}
7133
7134/**
7135 * scx_bpf_put_cpumask - Release a possible/online cpumask
7136 * @cpumask: cpumask to release
7137 */
7138__bpf_kfunc void scx_bpf_put_cpumask(const struct cpumask *cpumask)
7139{
7140	/*
7141	 * Empty function body because we aren't actually acquiring or releasing
7142	 * a reference to a global cpumask, which is read-only in the caller and
7143	 * is never released. The acquire / release semantics here are just used
7144	 * to make the cpumask is a trusted pointer in the caller.
7145	 */
7146}
7147
7148/**
7149 * scx_bpf_task_running - Is task currently running?
7150 * @p: task of interest
7151 */
7152__bpf_kfunc bool scx_bpf_task_running(const struct task_struct *p)
7153{
7154	return task_rq(p)->curr == p;
7155}
7156
7157/**
7158 * scx_bpf_task_cpu - CPU a task is currently associated with
7159 * @p: task of interest
7160 */
7161__bpf_kfunc s32 scx_bpf_task_cpu(const struct task_struct *p)
7162{
7163	return task_cpu(p);
7164}
7165
7166/**
7167 * scx_bpf_cpu_rq - Fetch the rq of a CPU
7168 * @cpu: CPU of the rq
7169 */
7170__bpf_kfunc struct rq *scx_bpf_cpu_rq(s32 cpu)
7171{
7172	struct scx_sched *sch;
7173
7174	guard(rcu)();
7175
7176	sch = rcu_dereference(scx_root);
7177	if (unlikely(!sch))
7178		return NULL;
7179
7180	if (!ops_cpu_valid(sch, cpu, NULL))
7181		return NULL;
7182
7183	if (!sch->warned_deprecated_rq) {
7184		printk_deferred(KERN_WARNING "sched_ext: %s() is deprecated; "
7185				"use scx_bpf_locked_rq() when holding rq lock "
7186				"or scx_bpf_cpu_curr() to read remote curr safely.\n", __func__);
7187		sch->warned_deprecated_rq = true;
7188	}
7189
7190	return cpu_rq(cpu);
7191}
7192
7193/**
7194 * scx_bpf_locked_rq - Return the rq currently locked by SCX
7195 *
7196 * Returns the rq if a rq lock is currently held by SCX.
7197 * Otherwise emits an error and returns NULL.
7198 */
7199__bpf_kfunc struct rq *scx_bpf_locked_rq(void)
7200{
7201	struct scx_sched *sch;
7202	struct rq *rq;
7203
7204	guard(preempt)();
7205
7206	sch = rcu_dereference_sched(scx_root);
7207	if (unlikely(!sch))
7208		return NULL;
7209
7210	rq = scx_locked_rq();
7211	if (!rq) {
7212		scx_error(sch, "accessing rq without holding rq lock");
7213		return NULL;
7214	}
7215
7216	return rq;
7217}
7218
7219/**
7220 * scx_bpf_cpu_curr - Return remote CPU's curr task
7221 * @cpu: CPU of interest
7222 *
7223 * Callers must hold RCU read lock (KF_RCU).
7224 */
7225__bpf_kfunc struct task_struct *scx_bpf_cpu_curr(s32 cpu)
7226{
7227	struct scx_sched *sch;
7228
7229	guard(rcu)();
7230
7231	sch = rcu_dereference(scx_root);
7232	if (unlikely(!sch))
7233		return NULL;
7234
7235	if (!ops_cpu_valid(sch, cpu, NULL))
7236		return NULL;
7237
7238	return rcu_dereference(cpu_rq(cpu)->curr);
7239}
7240
7241/**
7242 * scx_bpf_task_cgroup - Return the sched cgroup of a task
7243 * @p: task of interest
7244 *
7245 * @p->sched_task_group->css.cgroup represents the cgroup @p is associated with
7246 * from the scheduler's POV. SCX operations should use this function to
7247 * determine @p's current cgroup as, unlike following @p->cgroups,
7248 * @p->sched_task_group is protected by @p's rq lock and thus atomic w.r.t. all
7249 * rq-locked operations. Can be called on the parameter tasks of rq-locked
7250 * operations. The restriction guarantees that @p's rq is locked by the caller.
7251 */
7252#ifdef CONFIG_CGROUP_SCHED
7253__bpf_kfunc struct cgroup *scx_bpf_task_cgroup(struct task_struct *p)
7254{
7255	struct task_group *tg = p->sched_task_group;
7256	struct cgroup *cgrp = &cgrp_dfl_root.cgrp;
7257	struct scx_sched *sch;
7258
7259	guard(rcu)();
7260
7261	sch = rcu_dereference(scx_root);
7262	if (unlikely(!sch))
7263		goto out;
7264
7265	if (!scx_kf_allowed_on_arg_tasks(sch, __SCX_KF_RQ_LOCKED, p))
7266		goto out;
7267
7268	cgrp = tg_cgrp(tg);
7269
7270out:
7271	cgroup_get(cgrp);
7272	return cgrp;
7273}
7274#endif
7275
7276/**
7277 * scx_bpf_now - Returns a high-performance monotonically non-decreasing
7278 * clock for the current CPU. The clock returned is in nanoseconds.
7279 *
7280 * It provides the following properties:
7281 *
7282 * 1) High performance: Many BPF schedulers call bpf_ktime_get_ns() frequently
7283 *  to account for execution time and track tasks' runtime properties.
7284 *  Unfortunately, in some hardware platforms, bpf_ktime_get_ns() -- which
7285 *  eventually reads a hardware timestamp counter -- is neither performant nor
7286 *  scalable. scx_bpf_now() aims to provide a high-performance clock by
7287 *  using the rq clock in the scheduler core whenever possible.
7288 *
7289 * 2) High enough resolution for the BPF scheduler use cases: In most BPF
7290 *  scheduler use cases, the required clock resolution is lower than the most
7291 *  accurate hardware clock (e.g., rdtsc in x86). scx_bpf_now() basically
7292 *  uses the rq clock in the scheduler core whenever it is valid. It considers
7293 *  that the rq clock is valid from the time the rq clock is updated
7294 *  (update_rq_clock) until the rq is unlocked (rq_unpin_lock).
7295 *
7296 * 3) Monotonically non-decreasing clock for the same CPU: scx_bpf_now()
7297 *  guarantees the clock never goes backward when comparing them in the same
7298 *  CPU. On the other hand, when comparing clocks in different CPUs, there
7299 *  is no such guarantee -- the clock can go backward. It provides a
7300 *  monotonically *non-decreasing* clock so that it would provide the same
7301 *  clock values in two different scx_bpf_now() calls in the same CPU
7302 *  during the same period of when the rq clock is valid.
7303 */
7304__bpf_kfunc u64 scx_bpf_now(void)
7305{
7306	struct rq *rq;
7307	u64 clock;
7308
7309	preempt_disable();
7310
7311	rq = this_rq();
7312	if (smp_load_acquire(&rq->scx.flags) & SCX_RQ_CLK_VALID) {
7313		/*
7314		 * If the rq clock is valid, use the cached rq clock.
7315		 *
7316		 * Note that scx_bpf_now() is re-entrant between a process
7317		 * context and an interrupt context (e.g., timer interrupt).
7318		 * However, we don't need to consider the race between them
7319		 * because such race is not observable from a caller.
7320		 */
7321		clock = READ_ONCE(rq->scx.clock);
7322	} else {
7323		/*
7324		 * Otherwise, return a fresh rq clock.
7325		 *
7326		 * The rq clock is updated outside of the rq lock.
7327		 * In this case, keep the updated rq clock invalid so the next
7328		 * kfunc call outside the rq lock gets a fresh rq clock.
7329		 */
7330		clock = sched_clock_cpu(cpu_of(rq));
7331	}
7332
7333	preempt_enable();
7334
7335	return clock;
7336}
7337
7338static void scx_read_events(struct scx_sched *sch, struct scx_event_stats *events)
7339{
7340	struct scx_event_stats *e_cpu;
7341	int cpu;
7342
7343	/* Aggregate per-CPU event counters into @events. */
7344	memset(events, 0, sizeof(*events));
7345	for_each_possible_cpu(cpu) {
7346		e_cpu = &per_cpu_ptr(sch->pcpu, cpu)->event_stats;
7347		scx_agg_event(events, e_cpu, SCX_EV_SELECT_CPU_FALLBACK);
7348		scx_agg_event(events, e_cpu, SCX_EV_DISPATCH_LOCAL_DSQ_OFFLINE);
7349		scx_agg_event(events, e_cpu, SCX_EV_DISPATCH_KEEP_LAST);
7350		scx_agg_event(events, e_cpu, SCX_EV_ENQ_SKIP_EXITING);
7351		scx_agg_event(events, e_cpu, SCX_EV_ENQ_SKIP_MIGRATION_DISABLED);
7352		scx_agg_event(events, e_cpu, SCX_EV_REFILL_SLICE_DFL);
7353		scx_agg_event(events, e_cpu, SCX_EV_BYPASS_DURATION);
7354		scx_agg_event(events, e_cpu, SCX_EV_BYPASS_DISPATCH);
7355		scx_agg_event(events, e_cpu, SCX_EV_BYPASS_ACTIVATE);
7356	}
7357}
7358
7359/*
7360 * scx_bpf_events - Get a system-wide event counter to
7361 * @events: output buffer from a BPF program
7362 * @events__sz: @events len, must end in '__sz'' for the verifier
7363 */
7364__bpf_kfunc void scx_bpf_events(struct scx_event_stats *events,
7365				size_t events__sz)
7366{
7367	struct scx_sched *sch;
7368	struct scx_event_stats e_sys;
7369
7370	rcu_read_lock();
7371	sch = rcu_dereference(scx_root);
7372	if (sch)
7373		scx_read_events(sch, &e_sys);
7374	else
7375		memset(&e_sys, 0, sizeof(e_sys));
7376	rcu_read_unlock();
7377
7378	/*
7379	 * We cannot entirely trust a BPF-provided size since a BPF program
7380	 * might be compiled against a different vmlinux.h, of which
7381	 * scx_event_stats would be larger (a newer vmlinux.h) or smaller
7382	 * (an older vmlinux.h). Hence, we use the smaller size to avoid
7383	 * memory corruption.
7384	 */
7385	events__sz = min(events__sz, sizeof(*events));
7386	memcpy(events, &e_sys, events__sz);
7387}
7388
7389__bpf_kfunc_end_defs();
7390
7391BTF_KFUNCS_START(scx_kfunc_ids_any)
7392BTF_ID_FLAGS(func, scx_bpf_task_set_slice, KF_RCU);
7393BTF_ID_FLAGS(func, scx_bpf_task_set_dsq_vtime, KF_RCU);
7394BTF_ID_FLAGS(func, scx_bpf_kick_cpu)
7395BTF_ID_FLAGS(func, scx_bpf_dsq_nr_queued)
7396BTF_ID_FLAGS(func, scx_bpf_destroy_dsq)
7397BTF_ID_FLAGS(func, scx_bpf_dsq_peek, KF_RCU_PROTECTED | KF_RET_NULL)
7398BTF_ID_FLAGS(func, bpf_iter_scx_dsq_new, KF_ITER_NEW | KF_RCU_PROTECTED)
7399BTF_ID_FLAGS(func, bpf_iter_scx_dsq_next, KF_ITER_NEXT | KF_RET_NULL)
7400BTF_ID_FLAGS(func, bpf_iter_scx_dsq_destroy, KF_ITER_DESTROY)
7401BTF_ID_FLAGS(func, scx_bpf_exit_bstr, KF_TRUSTED_ARGS)
7402BTF_ID_FLAGS(func, scx_bpf_error_bstr, KF_TRUSTED_ARGS)
7403BTF_ID_FLAGS(func, scx_bpf_dump_bstr, KF_TRUSTED_ARGS)
7404BTF_ID_FLAGS(func, scx_bpf_reenqueue_local___v2)
7405BTF_ID_FLAGS(func, scx_bpf_cpuperf_cap)
7406BTF_ID_FLAGS(func, scx_bpf_cpuperf_cur)
7407BTF_ID_FLAGS(func, scx_bpf_cpuperf_set)
7408BTF_ID_FLAGS(func, scx_bpf_nr_node_ids)
7409BTF_ID_FLAGS(func, scx_bpf_nr_cpu_ids)
7410BTF_ID_FLAGS(func, scx_bpf_get_possible_cpumask, KF_ACQUIRE)
7411BTF_ID_FLAGS(func, scx_bpf_get_online_cpumask, KF_ACQUIRE)
7412BTF_ID_FLAGS(func, scx_bpf_put_cpumask, KF_RELEASE)
7413BTF_ID_FLAGS(func, scx_bpf_task_running, KF_RCU)
7414BTF_ID_FLAGS(func, scx_bpf_task_cpu, KF_RCU)
7415BTF_ID_FLAGS(func, scx_bpf_cpu_rq)
7416BTF_ID_FLAGS(func, scx_bpf_locked_rq, KF_RET_NULL)
7417BTF_ID_FLAGS(func, scx_bpf_cpu_curr, KF_RET_NULL | KF_RCU_PROTECTED)
7418#ifdef CONFIG_CGROUP_SCHED
7419BTF_ID_FLAGS(func, scx_bpf_task_cgroup, KF_RCU | KF_ACQUIRE)
7420#endif
7421BTF_ID_FLAGS(func, scx_bpf_now)
7422BTF_ID_FLAGS(func, scx_bpf_events, KF_TRUSTED_ARGS)
7423BTF_KFUNCS_END(scx_kfunc_ids_any)
7424
7425static const struct btf_kfunc_id_set scx_kfunc_set_any = {
7426	.owner			= THIS_MODULE,
7427	.set			= &scx_kfunc_ids_any,
7428};
7429
7430static int __init scx_init(void)
7431{
7432	int ret;
7433
7434	/*
7435	 * kfunc registration can't be done from init_sched_ext_class() as
7436	 * register_btf_kfunc_id_set() needs most of the system to be up.
7437	 *
7438	 * Some kfuncs are context-sensitive and can only be called from
7439	 * specific SCX ops. They are grouped into BTF sets accordingly.
7440	 * Unfortunately, BPF currently doesn't have a way of enforcing such
7441	 * restrictions. Eventually, the verifier should be able to enforce
7442	 * them. For now, register them the same and make each kfunc explicitly
7443	 * check using scx_kf_allowed().
7444	 */
7445	if ((ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
7446					     &scx_kfunc_set_enqueue_dispatch)) ||
7447	    (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
7448					     &scx_kfunc_set_dispatch)) ||
7449	    (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
7450					     &scx_kfunc_set_cpu_release)) ||
7451	    (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
7452					     &scx_kfunc_set_unlocked)) ||
7453	    (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL,
7454					     &scx_kfunc_set_unlocked)) ||
7455	    (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_STRUCT_OPS,
7456					     &scx_kfunc_set_any)) ||
7457	    (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_TRACING,
7458					     &scx_kfunc_set_any)) ||
7459	    (ret = register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL,
7460					     &scx_kfunc_set_any))) {
7461		pr_err("sched_ext: Failed to register kfunc sets (%d)\n", ret);
7462		return ret;
7463	}
7464
7465	ret = scx_idle_init();
7466	if (ret) {
7467		pr_err("sched_ext: Failed to initialize idle tracking (%d)\n", ret);
7468		return ret;
7469	}
7470
7471	ret = register_bpf_struct_ops(&bpf_sched_ext_ops, sched_ext_ops);
7472	if (ret) {
7473		pr_err("sched_ext: Failed to register struct_ops (%d)\n", ret);
7474		return ret;
7475	}
7476
7477	ret = register_pm_notifier(&scx_pm_notifier);
7478	if (ret) {
7479		pr_err("sched_ext: Failed to register PM notifier (%d)\n", ret);
7480		return ret;
7481	}
7482
7483	scx_kset = kset_create_and_add("sched_ext", &scx_uevent_ops, kernel_kobj);
7484	if (!scx_kset) {
7485		pr_err("sched_ext: Failed to create /sys/kernel/sched_ext\n");
7486		return -ENOMEM;
7487	}
7488
7489	ret = sysfs_create_group(&scx_kset->kobj, &scx_global_attr_group);
7490	if (ret < 0) {
7491		pr_err("sched_ext: Failed to add global attributes\n");
7492		return ret;
7493	}
7494
7495	if (!alloc_cpumask_var(&scx_bypass_lb_donee_cpumask, GFP_KERNEL) ||
7496	    !alloc_cpumask_var(&scx_bypass_lb_resched_cpumask, GFP_KERNEL)) {
7497		pr_err("sched_ext: Failed to allocate cpumasks\n");
7498		return -ENOMEM;
7499	}
7500
7501	return 0;
7502}
7503__initcall(scx_init);