Quelle  ext.c   Sprache: C

/* SPDX-License-Identifier: GPL-2.0 */
/*
 * BPF extensible scheduler class: Documentation/scheduler/sched-ext.rst
 *
 * Copyright (c) 2022 Meta Platforms, Inc. and affiliates.
 * Copyright (c) 2022 Tejun Heo <tj@kernel.org>
 * Copyright (c) 2022 David Vernet <dvernet@meta.com>
 */
#include <linux/btf_ids.h>
#include "ext_idle.h"

/*
 * NOTE: sched_ext is in the process of growing multiple scheduler support and
 * scx_root usage is in a transitional state. Naked dereferences are safe if the
 * caller is one of the tasks attached to SCX and explicit RCU dereference is
 * necessary otherwise. Naked scx_root dereferences trigger sparse warnings but
 * are used as temporary markers to indicate that the dereferences need to be
 * updated to point to the associated scheduler instances rather than scx_root.
 */
static struct scx_sched __rcu *scx_root;

/*
 * During exit, a task may schedule after losing its PIDs. When disabling the
 * BPF scheduler, we need to be able to iterate tasks in every state to
 * guarantee system safety. Maintain a dedicated task list which contains every
 * task between its fork and eventual free.
 */
static DEFINE_SPINLOCK(scx_tasks_lock);
static LIST_HEAD(scx_tasks);

/* ops enable/disable */
static DEFINE_MUTEX(scx_enable_mutex);
DEFINE_STATIC_KEY_FALSE(__scx_enabled);
DEFINE_STATIC_PERCPU_RWSEM(scx_fork_rwsem);
static atomic_t scx_enable_state_var = ATOMIC_INIT(SCX_DISABLED);
static unsigned long scx_in_softlockup;
static atomic_t scx_breather_depth = ATOMIC_INIT(0);
static int scx_bypass_depth;
static bool scx_init_task_enabled;
static bool scx_switching_all;
DEFINE_STATIC_KEY_FALSE(__scx_switched_all);

static atomic_long_t scx_nr_rejected = ATOMIC_LONG_INIT(0);
static atomic_long_t scx_hotplug_seq = ATOMIC_LONG_INIT(0);

/*
 * A monotically increasing sequence number that is incremented every time a
 * scheduler is enabled. This can be used by to check if any custom sched_ext
 * scheduler has ever been used in the system.
 */
static atomic_long_t scx_enable_seq = ATOMIC_LONG_INIT(0);

/*
 * The maximum amount of time in jiffies that a task may be runnable without
 * being scheduled on a CPU. If this timeout is exceeded, it will trigger
 * scx_error().
 */
static unsigned long scx_watchdog_timeout;

/*
 * The last time the delayed work was run. This delayed work relies on
 * ksoftirqd being able to run to service timer interrupts, so it's possible
 * that this work itself could get wedged. To account for this, we check that
 * it's not stalled in the timer tick, and trigger an error if it is.
 */
static unsigned long scx_watchdog_timestamp = INITIAL_JIFFIES;

static struct delayed_work scx_watchdog_work;

/* for %SCX_KICK_WAIT */
static unsigned long __percpu *scx_kick_cpus_pnt_seqs;

/*
 * Direct dispatch marker.
 *
 * Non-NULL values are used for direct dispatch from enqueue path. A valid
 * pointer points to the task currently being enqueued. An ERR_PTR value is used
 * to indicate that direct dispatch has already happened.
 */
static DEFINE_PER_CPU(struct task_struct *, direct_dispatch_task);

static const struct rhashtable_params dsq_hash_params = {
 .key_len  = sizeof_field(struct scx_dispatch_q, id),
 .key_offset  = offsetof(struct scx_dispatch_q, id),
 .head_offset  = offsetof(struct scx_dispatch_q, hash_node),
};

static LLIST_HEAD(dsqs_to_free);

/* dispatch buf */
struct scx_dsp_buf_ent {
 struct task_struct *task;
 unsigned long  qseq;
 u64   dsq_id;
 u64   enq_flags;
};

static u32 scx_dsp_max_batch;

struct scx_dsp_ctx {
 struct rq  *rq;
 u32   cursor;
 u32   nr_tasks;
 struct scx_dsp_buf_ent buf[];
};

static struct scx_dsp_ctx __percpu *scx_dsp_ctx;

/* string formatting from BPF */
struct scx_bstr_buf {
 u64   data[MAX_BPRINTF_VARARGS];
 char   line[SCX_EXIT_MSG_LEN];
};

static DEFINE_RAW_SPINLOCK(scx_exit_bstr_buf_lock);
static struct scx_bstr_buf scx_exit_bstr_buf;

/* ops debug dump */
struct scx_dump_data {
 s32   cpu;
 bool   first;
 s32   cursor;
 struct seq_buf  *s;
 const char  *prefix;
 struct scx_bstr_buf buf;
};

static struct scx_dump_data scx_dump_data = {
 .cpu   = -1,
};

/* /sys/kernel/sched_ext interface */
static struct kset *scx_kset;

#define CREATE_TRACE_POINTS
#include <trace/events/sched_ext.h>

static void process_ddsp_deferred_locals(struct rq *rq);
static void scx_bpf_kick_cpu(s32 cpu, u64 flags);
static void scx_vexit(struct scx_sched *sch, enum scx_exit_kind kind,
        s64 exit_code, const char *fmt, va_list args);

static __printf(4, 5) void scx_exit(struct scx_sched *sch,
        enum scx_exit_kind kind, s64 exit_code,
        const char *fmt, ...)
{
 va_list args;

 va_start(args, fmt);
 scx_vexit(sch, kind, exit_code, fmt, args);
 va_end(args);
}

static __printf(3, 4) void scx_kf_exit(enum scx_exit_kind kind, s64 exit_code,
           const char *fmt, ...)
{
 struct scx_sched *sch;
 va_list args;

 rcu_read_lock();
 sch = rcu_dereference(scx_root);
 if (sch) {
  va_start(args, fmt);
  scx_vexit(sch, kind, exit_code, fmt, args);
  va_end(args);
 }
 rcu_read_unlock();
}

#define scx_error(sch, fmt, args...) scx_exit((sch), SCX_EXIT_ERROR, 0, fmt, ##args)
#define scx_kf_error(fmt, args...) scx_kf_exit(SCX_EXIT_ERROR, 0, fmt, ##args)

#define SCX_HAS_OP(sch, op) test_bit(SCX_OP_IDX(op), (sch)->has_op)

static long jiffies_delta_msecs(unsigned long at, unsigned long now)
{
 if (time_after(at, now))
  return jiffies_to_msecs(at - now);
 else
  return -(long)jiffies_to_msecs(now - at);
}

/* if the highest set bit is N, return a mask with bits [N+1, 31] set */
static u32 higher_bits(u32 flags)
{
 return ~((1 << fls(flags)) - 1);
}

/* return the mask with only the highest bit set */
static u32 highest_bit(u32 flags)
{
 int bit = fls(flags);
 return ((u64)1 << bit) >> 1;
}

static bool u32_before(u32 a, u32 b)
{
 return (s32)(a - b) < 0;
}

static struct scx_dispatch_q *find_global_dsq(struct task_struct *p)
{
 struct scx_sched *sch = scx_root;

 return sch->global_dsqs[cpu_to_node(task_cpu(p))];
}

static struct scx_dispatch_q *find_user_dsq(struct scx_sched *sch, u64 dsq_id)
{
 return rhashtable_lookup_fast(&sch->dsq_hash, &dsq_id, dsq_hash_params);
}

/*
 * scx_kf_mask enforcement. Some kfuncs can only be called from specific SCX
 * ops. When invoking SCX ops, SCX_CALL_OP[_RET]() should be used to indicate
 * the allowed kfuncs and those kfuncs should use scx_kf_allowed() to check
 * whether it's running from an allowed context.
 *
 * @mask is constant, always inline to cull the mask calculations.
 */
static __always_inline void scx_kf_allow(u32 mask)
{
 /* nesting is allowed only in increasing scx_kf_mask order */
 WARN_ONCE((mask | higher_bits(mask)) & current->scx.kf_mask,
    "invalid nesting current->scx.kf_mask=0x%x mask=0x%x\n",
    current->scx.kf_mask, mask);
 current->scx.kf_mask |= mask;
 barrier();
}

static void scx_kf_disallow(u32 mask)
{
 barrier();
 current->scx.kf_mask &= ~mask;
}

/*
 * Track the rq currently locked.
 *
 * This allows kfuncs to safely operate on rq from any scx ops callback,
 * knowing which rq is already locked.
 */
DEFINE_PER_CPU(struct rq *, scx_locked_rq_state);

static inline void update_locked_rq(struct rq *rq)
{
 /*
 * Check whether @rq is actually locked. This can help expose bugs
 * or incorrect assumptions about the context in which a kfunc or
 * callback is executed.
 */
 if (rq)
  lockdep_assert_rq_held(rq);
 __this_cpu_write(scx_locked_rq_state, rq);
}

#define SCX_CALL_OP(sch, mask, op, rq, args...)     \
do {          \
 if (rq)         \
  update_locked_rq(rq);      \
 if (mask) {        \
  scx_kf_allow(mask);      \
  (sch)->ops.op(args);      \
  scx_kf_disallow(mask);      \
 } else {        \
  (sch)->ops.op(args);      \
 }         \
 if (rq)         \
  update_locked_rq(NULL);      \
} while (0)

#define SCX_CALL_OP_RET(sch, mask, op, rq, args...)    \
({          \
 __typeof__((sch)->ops.op(args)) __ret;     \
          \
 if (rq)         \
  update_locked_rq(rq);      \
 if (mask) {        \
  scx_kf_allow(mask);      \
  __ret = (sch)->ops.op(args);     \
  scx_kf_disallow(mask);      \
 } else {        \
  __ret = (sch)->ops.op(args);     \
 }         \
 if (rq)         \
  update_locked_rq(NULL);      \
 __ret;         \
})

/*
 * Some kfuncs are allowed only on the tasks that are subjects of the
 * in-progress scx_ops operation for, e.g., locking guarantees. To enforce such
 * restrictions, the following SCX_CALL_OP_*() variants should be used when
 * invoking scx_ops operations that take task arguments. These can only be used
 * for non-nesting operations due to the way the tasks are tracked.
 *
 * kfuncs which can only operate on such tasks can in turn use
 * scx_kf_allowed_on_arg_tasks() to test whether the invocation is allowed on
 * the specific task.
 */
#define SCX_CALL_OP_TASK(sch, mask, op, rq, task, args...)   \
do {          \
 BUILD_BUG_ON((mask) & ~__SCX_KF_TERMINAL);    \
 current->scx.kf_tasks[0] = task;     \
 SCX_CALL_OP((sch), mask, op, rq, task, ##args);    \
 current->scx.kf_tasks[0] = NULL;     \
} while (0)

#define SCX_CALL_OP_TASK_RET(sch, mask, op, rq, task, args...)   \
({          \
 __typeof__((sch)->ops.op(task, ##args)) __ret;    \
 BUILD_BUG_ON((mask) & ~__SCX_KF_TERMINAL);    \
 current->scx.kf_tasks[0] = task;     \
 __ret = SCX_CALL_OP_RET((sch), mask, op, rq, task, ##args);  \
 current->scx.kf_tasks[0] = NULL;     \
 __ret;         \
})

#define SCX_CALL_OP_2TASKS_RET(sch, mask, op, rq, task0, task1, args...) \
({          \
 __typeof__((sch)->ops.op(task0, task1, ##args)) __ret;   \
 BUILD_BUG_ON((mask) & ~__SCX_KF_TERMINAL);    \
 current->scx.kf_tasks[0] = task0;     \
 current->scx.kf_tasks[1] = task1;     \
 __ret = SCX_CALL_OP_RET((sch), mask, op, rq, task0, task1, ##args); \
 current->scx.kf_tasks[0] = NULL;     \
 current->scx.kf_tasks[1] = NULL;     \
 __ret;         \
})

/* @mask is constant, always inline to cull unnecessary branches */
static __always_inline bool scx_kf_allowed(u32 mask)
{
 if (unlikely(!(current->scx.kf_mask & mask))) {
  scx_kf_error("kfunc with mask 0x%x called from an operation only allowing 0x%x",
        mask, current->scx.kf_mask);
  return false;
 }

 /*
 * Enforce nesting boundaries. e.g. A kfunc which can be called from
 * DISPATCH must not be called if we're running DEQUEUE which is nested
 * inside ops.dispatch(). We don't need to check boundaries for any
 * blocking kfuncs as the verifier ensures they're only called from
 * sleepable progs.
 */
 if (unlikely(highest_bit(mask) == SCX_KF_CPU_RELEASE &&
       (current->scx.kf_mask & higher_bits(SCX_KF_CPU_RELEASE)))) {
  scx_kf_error("cpu_release kfunc called from a nested operation");
  return false;
 }

 if (unlikely(highest_bit(mask) == SCX_KF_DISPATCH &&
       (current->scx.kf_mask & higher_bits(SCX_KF_DISPATCH)))) {
  scx_kf_error("dispatch kfunc called from a nested operation");
  return false;
 }

 return true;
}

/* see SCX_CALL_OP_TASK() */
static __always_inline bool scx_kf_allowed_on_arg_tasks(u32 mask,
       struct task_struct *p)
{
 if (!scx_kf_allowed(mask))
  return false;

 if (unlikely((p != current->scx.kf_tasks[0] &&
        p != current->scx.kf_tasks[1]))) {
  scx_kf_error("called on a task not being operated on");
  return false;
 }

 return true;
}

/**
 * nldsq_next_task - Iterate to the next task in a non-local DSQ
 * @dsq: user dsq being iterated
 * @cur: current position, %NULL to start iteration
 * @rev: walk backwards
 *
 * Returns %NULL when iteration is finished.
 */
static struct task_struct *nldsq_next_task(struct scx_dispatch_q *dsq,
        struct task_struct *cur, bool rev)
{
 struct list_head *list_node;
 struct scx_dsq_list_node *dsq_lnode;

 lockdep_assert_held(&dsq->lock);

 if (cur)
  list_node = &cur->scx.dsq_list.node;
 else
  list_node = &dsq->list;

 /* find the next task, need to skip BPF iteration cursors */
 do {
  if (rev)
   list_node = list_node->prev;
  else
   list_node = list_node->next;

  if (list_node == &dsq->list)
   return NULL;

  dsq_lnode = container_of(list_node, struct scx_dsq_list_node,
      node);
 } while (dsq_lnode->flags & SCX_DSQ_LNODE_ITER_CURSOR);

 return container_of(dsq_lnode, struct task_struct, scx.dsq_list);
}

#define nldsq_for_each_task(p, dsq)      \
 for ((p) = nldsq_next_task((dsq), NULL, false); (p);   \
      (p) = nldsq_next_task((dsq), (p), false))


/*
 * BPF DSQ iterator. Tasks in a non-local DSQ can be iterated in [reverse]
 * dispatch order. BPF-visible iterator is opaque and larger to allow future
 * changes without breaking backward compatibility. Can be used with
 * bpf_for_each(). See bpf_iter_scx_dsq_*().
 */
enum scx_dsq_iter_flags {
 /* iterate in the reverse dispatch order */
 SCX_DSQ_ITER_REV  = 1U << 16,

 __SCX_DSQ_ITER_HAS_SLICE = 1U << 30,
 __SCX_DSQ_ITER_HAS_VTIME = 1U << 31,

 __SCX_DSQ_ITER_USER_FLAGS = SCX_DSQ_ITER_REV,
 __SCX_DSQ_ITER_ALL_FLAGS = __SCX_DSQ_ITER_USER_FLAGS |
       __SCX_DSQ_ITER_HAS_SLICE |
       __SCX_DSQ_ITER_HAS_VTIME,
};

struct bpf_iter_scx_dsq_kern {
 struct scx_dsq_list_node cursor;
 struct scx_dispatch_q  *dsq;
 u64    slice;
 u64    vtime;
} __attribute__((aligned(8)));

struct bpf_iter_scx_dsq {
 u64    __opaque[6];
} __attribute__((aligned(8)));


/*
 * SCX task iterator.
 */
struct scx_task_iter {
 struct sched_ext_entity  cursor;
 struct task_struct  *locked;
 struct rq   *rq;
 struct rq_flags   rf;
 u32    cnt;
};

/**
 * scx_task_iter_start - Lock scx_tasks_lock and start a task iteration
 * @iter: iterator to init
 *
 * Initialize @iter and return with scx_tasks_lock held. Once initialized, @iter
 * must eventually be stopped with scx_task_iter_stop().
 *
 * scx_tasks_lock and the rq lock may be released using scx_task_iter_unlock()
 * between this and the first next() call or between any two next() calls. If
 * the locks are released between two next() calls, the caller is responsible
 * for ensuring that the task being iterated remains accessible either through
 * RCU read lock or obtaining a reference count.
 *
 * All tasks which existed when the iteration started are guaranteed to be
 * visited as long as they still exist.
 */
static void scx_task_iter_start(struct scx_task_iter *iter)
{
 BUILD_BUG_ON(__SCX_DSQ_ITER_ALL_FLAGS &
       ((1U << __SCX_DSQ_LNODE_PRIV_SHIFT) - 1));

 spin_lock_irq(&scx_tasks_lock);

 iter->cursor = (struct sched_ext_entity){ .flags = SCX_TASK_CURSOR };
 list_add(&iter->cursor.tasks_node, &scx_tasks);
 iter->locked = NULL;
 iter->cnt = 0;
}

static void __scx_task_iter_rq_unlock(struct scx_task_iter *iter)
{
 if (iter->locked) {
  task_rq_unlock(iter->rq, iter->locked, &iter->rf);
  iter->locked = NULL;
 }
}

/**
 * scx_task_iter_unlock - Unlock rq and scx_tasks_lock held by a task iterator
 * @iter: iterator to unlock
 *
 * If @iter is in the middle of a locked iteration, it may be locking the rq of
 * the task currently being visited in addition to scx_tasks_lock. Unlock both.
 * This function can be safely called anytime during an iteration.
 */
static void scx_task_iter_unlock(struct scx_task_iter *iter)
{
 __scx_task_iter_rq_unlock(iter);
 spin_unlock_irq(&scx_tasks_lock);
}

/**
 * scx_task_iter_relock - Lock scx_tasks_lock released by scx_task_iter_unlock()
 * @iter: iterator to re-lock
 *
 * Re-lock scx_tasks_lock unlocked by scx_task_iter_unlock(). Note that it
 * doesn't re-lock the rq lock. Must be called before other iterator operations.
 */
static void scx_task_iter_relock(struct scx_task_iter *iter)
{
 spin_lock_irq(&scx_tasks_lock);
}

/**
 * scx_task_iter_stop - Stop a task iteration and unlock scx_tasks_lock
 * @iter: iterator to exit
 *
 * Exit a previously initialized @iter. Must be called with scx_tasks_lock held
 * which is released on return. If the iterator holds a task's rq lock, that rq
 * lock is also released. See scx_task_iter_start() for details.
 */
static void scx_task_iter_stop(struct scx_task_iter *iter)
{
 list_del_init(&iter->cursor.tasks_node);
 scx_task_iter_unlock(iter);
}

/**
 * scx_task_iter_next - Next task
 * @iter: iterator to walk
 *
 * Visit the next task. See scx_task_iter_start() for details. Locks are dropped
 * and re-acquired every %SCX_TASK_ITER_BATCH iterations to avoid causing stalls
 * by holding scx_tasks_lock for too long.
 */
static struct task_struct *scx_task_iter_next(struct scx_task_iter *iter)
{
 struct list_head *cursor = &iter->cursor.tasks_node;
 struct sched_ext_entity *pos;

 if (!(++iter->cnt % SCX_TASK_ITER_BATCH)) {
  scx_task_iter_unlock(iter);
  cond_resched();
  scx_task_iter_relock(iter);
 }

 list_for_each_entry(pos, cursor, tasks_node) {
  if (&pos->tasks_node == &scx_tasks)
   return NULL;
  if (!(pos->flags & SCX_TASK_CURSOR)) {
   list_move(cursor, &pos->tasks_node);
   return container_of(pos, struct task_struct, scx);
  }
 }

 /* can't happen, should always terminate at scx_tasks above */
 BUG();
}

/**
 * scx_task_iter_next_locked - Next non-idle task with its rq locked
 * @iter: iterator to walk
 *
 * Visit the non-idle task with its rq lock held. Allows callers to specify
 * whether they would like to filter out dead tasks. See scx_task_iter_start()
 * for details.
 */
static struct task_struct *scx_task_iter_next_locked(struct scx_task_iter *iter)
{
 struct task_struct *p;

 __scx_task_iter_rq_unlock(iter);

 while ((p = scx_task_iter_next(iter))) {
  /*
 * scx_task_iter is used to prepare and move tasks into SCX
 * while loading the BPF scheduler and vice-versa while
 * unloading. The init_tasks ("swappers") should be excluded
 * from the iteration because:
 *
 * - It's unsafe to use __setschduler_prio() on an init_task to
 *   determine the sched_class to use as it won't preserve its
 *   idle_sched_class.
 *
 * - ops.init/exit_task() can easily be confused if called with
 *   init_tasks as they, e.g., share PID 0.
 *
 * As init_tasks are never scheduled through SCX, they can be
 * skipped safely. Note that is_idle_task() which tests %PF_IDLE
 * doesn't work here:
 *
 * - %PF_IDLE may not be set for an init_task whose CPU hasn't
 *   yet been onlined.
 *
 * - %PF_IDLE can be set on tasks that are not init_tasks. See
 *   play_idle_precise() used by CONFIG_IDLE_INJECT.
 *
 * Test for idle_sched_class as only init_tasks are on it.
 */
  if (p->sched_class != &idle_sched_class)
   break;
 }
 if (!p)
  return NULL;

 iter->rq = task_rq_lock(p, &iter->rf);
 iter->locked = p;

 return p;
}

/**
 * scx_add_event - Increase an event counter for 'name' by 'cnt'
 * @sch: scx_sched to account events for
 * @name: an event name defined in struct scx_event_stats
 * @cnt: the number of the event occurred
 *
 * This can be used when preemption is not disabled.
 */
#define scx_add_event(sch, name, cnt) do {     \
 this_cpu_add((sch)->pcpu->event_stats.name, (cnt));   \
 trace_sched_ext_event(#name, (cnt));     \
} while(0)

/**
 * __scx_add_event - Increase an event counter for 'name' by 'cnt'
 * @sch: scx_sched to account events for
 * @name: an event name defined in struct scx_event_stats
 * @cnt: the number of the event occurred
 *
 * This should be used only when preemption is disabled.
 */
#define __scx_add_event(sch, name, cnt) do {     \
 __this_cpu_add((sch)->pcpu->event_stats.name, (cnt));   \
 trace_sched_ext_event(#name, cnt);     \
} while(0)

/**
 * scx_agg_event - Aggregate an event counter 'kind' from 'src_e' to 'dst_e'
 * @dst_e: destination event stats
 * @src_e: source event stats
 * @kind: a kind of event to be aggregated
 */
#define scx_agg_event(dst_e, src_e, kind) do {     \
 (dst_e)->kind += READ_ONCE((src_e)->kind);    \
} while(0)

/**
 * scx_dump_event - Dump an event 'kind' in 'events' to 's'
 * @s: output seq_buf
 * @events: event stats
 * @kind: a kind of event to dump
 */
#define scx_dump_event(s, events, kind) do {     \
 dump_line(&(s), "%40s: %16lld", #kind, (events)->kind);   \
} while (0)


static void scx_read_events(struct scx_sched *sch,
       struct scx_event_stats *events);

static enum scx_enable_state scx_enable_state(void)
{
 return atomic_read(&scx_enable_state_var);
}

static enum scx_enable_state scx_set_enable_state(enum scx_enable_state to)
{
 return atomic_xchg(&scx_enable_state_var, to);
}

static bool scx_tryset_enable_state(enum scx_enable_state to,
        enum scx_enable_state from)
{
 int from_v = from;

 return atomic_try_cmpxchg(&scx_enable_state_var, &from_v, to);
}

/**
 * wait_ops_state - Busy-wait the specified ops state to end
 * @p: target task
 * @opss: state to wait the end of
 *
 * Busy-wait for @p to transition out of @opss. This can only be used when the
 * state part of @opss is %SCX_QUEUEING or %SCX_DISPATCHING. This function also
 * has load_acquire semantics to ensure that the caller can see the updates made
 * in the enqueueing and dispatching paths.
 */
static void wait_ops_state(struct task_struct *p, unsigned long opss)
{
 do {
  cpu_relax();
 } while (atomic_long_read_acquire(&p->scx.ops_state) == opss);
}

static inline bool __cpu_valid(s32 cpu)
{
 return likely(cpu >= 0 && cpu < nr_cpu_ids && cpu_possible(cpu));
}

/**
 * ops_cpu_valid - Verify a cpu number, to be used on ops input args
 * @sch: scx_sched to abort on error
 * @cpu: cpu number which came from a BPF ops
 * @where: extra information reported on error
 *
 * @cpu is a cpu number which came from the BPF scheduler and can be any value.
 * Verify that it is in range and one of the possible cpus. If invalid, trigger
 * an ops error.
 */
static bool ops_cpu_valid(struct scx_sched *sch, s32 cpu, const char *where)
{
 if (__cpu_valid(cpu)) {
  return true;
 } else {
  scx_error(sch, "invalid CPU %d%s%s", cpu, where ? " " : "", where ?: "");
  return false;
 }
}

/**
 * kf_cpu_valid - Verify a CPU number, to be used on kfunc input args
 * @cpu: cpu number which came from a BPF ops
 * @where: extra information reported on error
 *
 * The same as ops_cpu_valid() but @sch is implicit.
 */
static bool kf_cpu_valid(u32 cpu, const char *where)
{
 if (__cpu_valid(cpu)) {
  return true;
 } else {
  scx_kf_error("invalid CPU %d%s%s", cpu, where ? " " : "", where ?: "");
  return false;
 }
}

/**
 * ops_sanitize_err - Sanitize a -errno value
 * @sch: scx_sched to error out on error
 * @ops_name: operation to blame on failure
 * @err: -errno value to sanitize
 *
 * Verify @err is a valid -errno. If not, trigger scx_error() and return
 * -%EPROTO. This is necessary because returning a rogue -errno up the chain can
 * cause misbehaviors. For an example, a large negative return from
 * ops.init_task() triggers an oops when passed up the call chain because the
 * value fails IS_ERR() test after being encoded with ERR_PTR() and then is
 * handled as a pointer.
 */
static int ops_sanitize_err(struct scx_sched *sch, const char *ops_name, s32 err)
{
 if (err < 0 && err >= -MAX_ERRNO)
  return err;

 scx_error(sch, "ops.%s() returned an invalid errno %d", ops_name, err);
 return -EPROTO;
}

static void run_deferred(struct rq *rq)
{
 process_ddsp_deferred_locals(rq);
}

static void deferred_bal_cb_workfn(struct rq *rq)
{
 run_deferred(rq);
}

static void deferred_irq_workfn(struct irq_work *irq_work)
{
 struct rq *rq = container_of(irq_work, struct rq, scx.deferred_irq_work);

 raw_spin_rq_lock(rq);
 run_deferred(rq);
 raw_spin_rq_unlock(rq);
}

/**
 * schedule_deferred - Schedule execution of deferred actions on an rq
 * @rq: target rq
 *
 * Schedule execution of deferred actions on @rq. Must be called with @rq
 * locked. Deferred actions are executed with @rq locked but unpinned, and thus
 * can unlock @rq to e.g. migrate tasks to other rqs.
 */
static void schedule_deferred(struct rq *rq)
{
 lockdep_assert_rq_held(rq);

 /*
 * If in the middle of waking up a task, task_woken_scx() will be called
 * afterwards which will then run the deferred actions, no need to
 * schedule anything.
 */
 if (rq->scx.flags & SCX_RQ_IN_WAKEUP)
  return;

 /*
 * If in balance, the balance callbacks will be called before rq lock is
 * released. Schedule one.
 */
 if (rq->scx.flags & SCX_RQ_IN_BALANCE) {
  queue_balance_callback(rq, &rq->scx.deferred_bal_cb,
           deferred_bal_cb_workfn);
  return;
 }

 /*
 * No scheduler hooks available. Queue an irq work. They are executed on
 * IRQ re-enable which may take a bit longer than the scheduler hooks.
 * The above WAKEUP and BALANCE paths should cover most of the cases and
 * the time to IRQ re-enable shouldn't be long.
 */
 irq_work_queue(&rq->scx.deferred_irq_work);
}

/**
 * touch_core_sched - Update timestamp used for core-sched task ordering
 * @rq: rq to read clock from, must be locked
 * @p: task to update the timestamp for
 *
 * Update @p->scx.core_sched_at timestamp. This is used by scx_prio_less() to
 * implement global or local-DSQ FIFO ordering for core-sched. Should be called
 * when a task becomes runnable and its turn on the CPU ends (e.g. slice
 * exhaustion).
 */
static void touch_core_sched(struct rq *rq, struct task_struct *p)
{
 lockdep_assert_rq_held(rq);

#ifdef CONFIG_SCHED_CORE
 /*
 * It's okay to update the timestamp spuriously. Use
 * sched_core_disabled() which is cheaper than enabled().
 *
 * As this is used to determine ordering between tasks of sibling CPUs,
 * it may be better to use per-core dispatch sequence instead.
 */
 if (!sched_core_disabled())
  p->scx.core_sched_at = sched_clock_cpu(cpu_of(rq));
#endif
}

/**
 * touch_core_sched_dispatch - Update core-sched timestamp on dispatch
 * @rq: rq to read clock from, must be locked
 * @p: task being dispatched
 *
 * If the BPF scheduler implements custom core-sched ordering via
 * ops.core_sched_before(), @p->scx.core_sched_at is used to implement FIFO
 * ordering within each local DSQ. This function is called from dispatch paths
 * and updates @p->scx.core_sched_at if custom core-sched ordering is in effect.
 */
static void touch_core_sched_dispatch(struct rq *rq, struct task_struct *p)
{
 lockdep_assert_rq_held(rq);

#ifdef CONFIG_SCHED_CORE
 if (unlikely(SCX_HAS_OP(scx_root, core_sched_before)))
  touch_core_sched(rq, p);
#endif
}

static void update_curr_scx(struct rq *rq)
{
 struct task_struct *curr = rq->curr;
 s64 delta_exec;

 delta_exec = update_curr_common(rq);
 if (unlikely(delta_exec <= 0))
  return;

 if (curr->scx.slice != SCX_SLICE_INF) {
  curr->scx.slice -= min_t(u64, curr->scx.slice, delta_exec);
  if (!curr->scx.slice)
   touch_core_sched(rq, curr);
 }
}

static bool scx_dsq_priq_less(struct rb_node *node_a,
         const struct rb_node *node_b)
{
 const struct task_struct *a =
  container_of(node_a, struct task_struct, scx.dsq_priq);
 const struct task_struct *b =
  container_of(node_b, struct task_struct, scx.dsq_priq);

 return time_before64(a->scx.dsq_vtime, b->scx.dsq_vtime);
}

static void dsq_mod_nr(struct scx_dispatch_q *dsq, s32 delta)
{
 /* scx_bpf_dsq_nr_queued() reads ->nr without locking, use WRITE_ONCE() */
 WRITE_ONCE(dsq->nr, dsq->nr + delta);
}

static void refill_task_slice_dfl(struct task_struct *p)
{
 p->scx.slice = SCX_SLICE_DFL;
 __scx_add_event(scx_root, SCX_EV_REFILL_SLICE_DFL, 1);
}

static void dispatch_enqueue(struct scx_sched *sch, struct scx_dispatch_q *dsq,
        struct task_struct *p, u64 enq_flags)
{
 bool is_local = dsq->id == SCX_DSQ_LOCAL;

 WARN_ON_ONCE(p->scx.dsq || !list_empty(&p->scx.dsq_list.node));
 WARN_ON_ONCE((p->scx.dsq_flags & SCX_TASK_DSQ_ON_PRIQ) ||
       !RB_EMPTY_NODE(&p->scx.dsq_priq));

 if (!is_local) {
  raw_spin_lock(&dsq->lock);
  if (unlikely(dsq->id == SCX_DSQ_INVALID)) {
   scx_error(sch, "attempting to dispatch to a destroyed dsq");
   /* fall back to the global dsq */
   raw_spin_unlock(&dsq->lock);
   dsq = find_global_dsq(p);
   raw_spin_lock(&dsq->lock);
  }
 }

 if (unlikely((dsq->id & SCX_DSQ_FLAG_BUILTIN) &&
       (enq_flags & SCX_ENQ_DSQ_PRIQ))) {
  /*
 * SCX_DSQ_LOCAL and SCX_DSQ_GLOBAL DSQs always consume from
 * their FIFO queues. To avoid confusion and accidentally
 * starving vtime-dispatched tasks by FIFO-dispatched tasks, we
 * disallow any internal DSQ from doing vtime ordering of
 * tasks.
 */
  scx_error(sch, "cannot use vtime ordering for built-in DSQs");
  enq_flags &= ~SCX_ENQ_DSQ_PRIQ;
 }

 if (enq_flags & SCX_ENQ_DSQ_PRIQ) {
  struct rb_node *rbp;

  /*
 * A PRIQ DSQ shouldn't be using FIFO enqueueing. As tasks are
 * linked to both the rbtree and list on PRIQs, this can only be
 * tested easily when adding the first task.
 */
  if (unlikely(RB_EMPTY_ROOT(&dsq->priq) &&
        nldsq_next_task(dsq, NULL, false)))
   scx_error(sch, "DSQ ID 0x%016llx already had FIFO-enqueued tasks",
      dsq->id);

  p->scx.dsq_flags |= SCX_TASK_DSQ_ON_PRIQ;
  rb_add(&p->scx.dsq_priq, &dsq->priq, scx_dsq_priq_less);

  /*
 * Find the previous task and insert after it on the list so
 * that @dsq->list is vtime ordered.
 */
  rbp = rb_prev(&p->scx.dsq_priq);
  if (rbp) {
   struct task_struct *prev =
    container_of(rbp, struct task_struct,
          scx.dsq_priq);
   list_add(&p->scx.dsq_list.node, &prev->scx.dsq_list.node);
  } else {
   list_add(&p->scx.dsq_list.node, &dsq->list);
  }
 } else {
  /* a FIFO DSQ shouldn't be using PRIQ enqueuing */
  if (unlikely(!RB_EMPTY_ROOT(&dsq->priq)))
   scx_error(sch, "DSQ ID 0x%016llx already had PRIQ-enqueued tasks",
      dsq->id);

  if (enq_flags & (SCX_ENQ_HEAD | SCX_ENQ_PREEMPT))
   list_add(&p->scx.dsq_list.node, &dsq->list);
  else
   list_add_tail(&p->scx.dsq_list.node, &dsq->list);
 }

 /* seq records the order tasks are queued, used by BPF DSQ iterator */
 dsq->seq++;
 p->scx.dsq_seq = dsq->seq;

 dsq_mod_nr(dsq, 1);
 p->scx.dsq = dsq;

 /*
 * scx.ddsp_dsq_id and scx.ddsp_enq_flags are only relevant on the
 * direct dispatch path, but we clear them here because the direct
 * dispatch verdict may be overridden on the enqueue path during e.g.
 * bypass.
 */
 p->scx.ddsp_dsq_id = SCX_DSQ_INVALID;
 p->scx.ddsp_enq_flags = 0;

 /*
 * We're transitioning out of QUEUEING or DISPATCHING. store_release to
 * match waiters' load_acquire.
 */
 if (enq_flags & SCX_ENQ_CLEAR_OPSS)
  atomic_long_set_release(&p->scx.ops_state, SCX_OPSS_NONE);

 if (is_local) {
  struct rq *rq = container_of(dsq, struct rq, scx.local_dsq);
  bool preempt = false;

  if ((enq_flags & SCX_ENQ_PREEMPT) && p != rq->curr &&
      rq->curr->sched_class == &ext_sched_class) {
   rq->curr->scx.slice = 0;
   preempt = true;
  }

  if (preempt || sched_class_above(&ext_sched_class,
       rq->curr->sched_class))
   resched_curr(rq);
 } else {
  raw_spin_unlock(&dsq->lock);
 }
}

static void task_unlink_from_dsq(struct task_struct *p,
     struct scx_dispatch_q *dsq)
{
 WARN_ON_ONCE(list_empty(&p->scx.dsq_list.node));

 if (p->scx.dsq_flags & SCX_TASK_DSQ_ON_PRIQ) {
  rb_erase(&p->scx.dsq_priq, &dsq->priq);
  RB_CLEAR_NODE(&p->scx.dsq_priq);
  p->scx.dsq_flags &= ~SCX_TASK_DSQ_ON_PRIQ;
 }

 list_del_init(&p->scx.dsq_list.node);
 dsq_mod_nr(dsq, -1);
}

static void dispatch_dequeue(struct rq *rq, struct task_struct *p)
{
 struct scx_dispatch_q *dsq = p->scx.dsq;
 bool is_local = dsq == &rq->scx.local_dsq;

 if (!dsq) {
  /*
 * If !dsq && on-list, @p is on @rq's ddsp_deferred_locals.
 * Unlinking is all that's needed to cancel.
 */
  if (unlikely(!list_empty(&p->scx.dsq_list.node)))
   list_del_init(&p->scx.dsq_list.node);

  /*
 * When dispatching directly from the BPF scheduler to a local
 * DSQ, the task isn't associated with any DSQ but
 * @p->scx.holding_cpu may be set under the protection of
 * %SCX_OPSS_DISPATCHING.
 */
  if (p->scx.holding_cpu >= 0)
   p->scx.holding_cpu = -1;

  return;
 }

 if (!is_local)
  raw_spin_lock(&dsq->lock);

 /*
 * Now that we hold @dsq->lock, @p->holding_cpu and @p->scx.dsq_* can't
 * change underneath us.
*/
 if (p->scx.holding_cpu < 0) {
  /* @p must still be on @dsq, dequeue */
  task_unlink_from_dsq(p, dsq);
 } else {
  /*
 * We're racing against dispatch_to_local_dsq() which already
 * removed @p from @dsq and set @p->scx.holding_cpu. Clear the
 * holding_cpu which tells dispatch_to_local_dsq() that it lost
 * the race.
 */
  WARN_ON_ONCE(!list_empty(&p->scx.dsq_list.node));
  p->scx.holding_cpu = -1;
 }
 p->scx.dsq = NULL;

 if (!is_local)
  raw_spin_unlock(&dsq->lock);
}

static struct scx_dispatch_q *find_dsq_for_dispatch(struct scx_sched *sch,
          struct rq *rq, u64 dsq_id,
          struct task_struct *p)
{
 struct scx_dispatch_q *dsq;

 if (dsq_id == SCX_DSQ_LOCAL)
  return &rq->scx.local_dsq;

 if ((dsq_id & SCX_DSQ_LOCAL_ON) == SCX_DSQ_LOCAL_ON) {
  s32 cpu = dsq_id & SCX_DSQ_LOCAL_CPU_MASK;

  if (!ops_cpu_valid(sch, cpu, "in SCX_DSQ_LOCAL_ON dispatch verdict"))
   return find_global_dsq(p);

  return &cpu_rq(cpu)->scx.local_dsq;
 }

 if (dsq_id == SCX_DSQ_GLOBAL)
  dsq = find_global_dsq(p);
 else
  dsq = find_user_dsq(sch, dsq_id);

 if (unlikely(!dsq)) {
  scx_error(sch, "non-existent DSQ 0x%llx for %s[%d]",
     dsq_id, p->comm, p->pid);
  return find_global_dsq(p);
 }

 return dsq;
}

static void mark_direct_dispatch(struct task_struct *ddsp_task,
     struct task_struct *p, u64 dsq_id,
     u64 enq_flags)
{
 /*
 * Mark that dispatch already happened from ops.select_cpu() or
 * ops.enqueue() by spoiling direct_dispatch_task with a non-NULL value
 * which can never match a valid task pointer.
 */
 __this_cpu_write(direct_dispatch_task, ERR_PTR(-ESRCH));

 /* @p must match the task on the enqueue path */
 if (unlikely(p != ddsp_task)) {
  if (IS_ERR(ddsp_task))
   scx_kf_error("%s[%d] already direct-dispatched",
      p->comm, p->pid);
  else
   scx_kf_error("scheduling for %s[%d] but trying to direct-dispatch %s[%d]",
      ddsp_task->comm, ddsp_task->pid,
      p->comm, p->pid);
  return;
 }

 WARN_ON_ONCE(p->scx.ddsp_dsq_id != SCX_DSQ_INVALID);
 WARN_ON_ONCE(p->scx.ddsp_enq_flags);

 p->scx.ddsp_dsq_id = dsq_id;
 p->scx.ddsp_enq_flags = enq_flags;
}

static void direct_dispatch(struct scx_sched *sch, struct task_struct *p,
       u64 enq_flags)
{
 struct rq *rq = task_rq(p);
 struct scx_dispatch_q *dsq =
  find_dsq_for_dispatch(sch, rq, p->scx.ddsp_dsq_id, p);

 touch_core_sched_dispatch(rq, p);

 p->scx.ddsp_enq_flags |= enq_flags;

 /*
 * We are in the enqueue path with @rq locked and pinned, and thus can't
 * double lock a remote rq and enqueue to its local DSQ. For
 * DSQ_LOCAL_ON verdicts targeting the local DSQ of a remote CPU, defer
 * the enqueue so that it's executed when @rq can be unlocked.
 */
 if (dsq->id == SCX_DSQ_LOCAL && dsq != &rq->scx.local_dsq) {
  unsigned long opss;

  opss = atomic_long_read(&p->scx.ops_state) & SCX_OPSS_STATE_MASK;

  switch (opss & SCX_OPSS_STATE_MASK) {
  case SCX_OPSS_NONE:
   break;
  case SCX_OPSS_QUEUEING:
   /*
 * As @p was never passed to the BPF side, _release is
 * not strictly necessary. Still do it for consistency.
 */
   atomic_long_set_release(&p->scx.ops_state, SCX_OPSS_NONE);
   break;
  default:
   WARN_ONCE(true, "sched_ext: %s[%d] has invalid ops state 0x%lx in direct_dispatch()",
      p->comm, p->pid, opss);
   atomic_long_set_release(&p->scx.ops_state, SCX_OPSS_NONE);
   break;
  }

  WARN_ON_ONCE(p->scx.dsq || !list_empty(&p->scx.dsq_list.node));
  list_add_tail(&p->scx.dsq_list.node,
         &rq->scx.ddsp_deferred_locals);
  schedule_deferred(rq);
  return;
 }

 dispatch_enqueue(sch, dsq, p,
    p->scx.ddsp_enq_flags | SCX_ENQ_CLEAR_OPSS);
}

static bool scx_rq_online(struct rq *rq)
{
 /*
 * Test both cpu_active() and %SCX_RQ_ONLINE. %SCX_RQ_ONLINE indicates
 * the online state as seen from the BPF scheduler. cpu_active() test
 * guarantees that, if this function returns %true, %SCX_RQ_ONLINE will
 * stay set until the current scheduling operation is complete even if
 * we aren't locking @rq.
 */
 return likely((rq->scx.flags & SCX_RQ_ONLINE) && cpu_active(cpu_of(rq)));
}

static void do_enqueue_task(struct rq *rq, struct task_struct *p, u64 enq_flags,
       int sticky_cpu)
{
 struct scx_sched *sch = scx_root;
 struct task_struct **ddsp_taskp;
 unsigned long qseq;

 WARN_ON_ONCE(!(p->scx.flags & SCX_TASK_QUEUED));

 /* rq migration */
 if (sticky_cpu == cpu_of(rq))
  goto local_norefill;

 /*
 * If !scx_rq_online(), we already told the BPF scheduler that the CPU
 * is offline and are just running the hotplug path. Don't bother the
 * BPF scheduler.
 */
 if (!scx_rq_online(rq))
  goto local;

 if (scx_rq_bypassing(rq)) {
  __scx_add_event(sch, SCX_EV_BYPASS_DISPATCH, 1);
  goto global;
 }

 if (p->scx.ddsp_dsq_id != SCX_DSQ_INVALID)
  goto direct;

 /* see %SCX_OPS_ENQ_EXITING */
 if (!(sch->ops.flags & SCX_OPS_ENQ_EXITING) &&
     unlikely(p->flags & PF_EXITING)) {
  __scx_add_event(sch, SCX_EV_ENQ_SKIP_EXITING, 1);
  goto local;
 }

 /* see %SCX_OPS_ENQ_MIGRATION_DISABLED */
 if (!(sch->ops.flags & SCX_OPS_ENQ_MIGRATION_DISABLED) &&
     is_migration_disabled(p)) {
  __scx_add_event(sch, SCX_EV_ENQ_SKIP_MIGRATION_DISABLED, 1);
  goto local;
 }

 if (unlikely(!SCX_HAS_OP(sch, enqueue)))
  goto global;

 /* DSQ bypass didn't trigger, enqueue on the BPF scheduler */
 qseq = rq->scx.ops_qseq++ << SCX_OPSS_QSEQ_SHIFT;

 WARN_ON_ONCE(atomic_long_read(&p->scx.ops_state) != SCX_OPSS_NONE);
 atomic_long_set(&p->scx.ops_state, SCX_OPSS_QUEUEING | qseq);

 ddsp_taskp = this_cpu_ptr(&direct_dispatch_task);
 WARN_ON_ONCE(*ddsp_taskp);
 *ddsp_taskp = p;

 SCX_CALL_OP_TASK(sch, SCX_KF_ENQUEUE, enqueue, rq, p, enq_flags);

 *ddsp_taskp = NULL;
 if (p->scx.ddsp_dsq_id != SCX_DSQ_INVALID)
  goto direct;

 /*
 * If not directly dispatched, QUEUEING isn't clear yet and dispatch or
 * dequeue may be waiting. The store_release matches their load_acquire.
 */
 atomic_long_set_release(&p->scx.ops_state, SCX_OPSS_QUEUED | qseq);
 return;

direct:
 direct_dispatch(sch, p, enq_flags);
 return;

local:
 /*
 * For task-ordering, slice refill must be treated as implying the end
 * of the current slice. Otherwise, the longer @p stays on the CPU, the
 * higher priority it becomes from scx_prio_less()'s POV.
 */
 touch_core_sched(rq, p);
 refill_task_slice_dfl(p);
local_norefill:
 dispatch_enqueue(sch, &rq->scx.local_dsq, p, enq_flags);
 return;

global:
 touch_core_sched(rq, p); /* see the comment in local: */
 refill_task_slice_dfl(p);
 dispatch_enqueue(sch, find_global_dsq(p), p, enq_flags);
}

static bool task_runnable(const struct task_struct *p)
{
 return !list_empty(&p->scx.runnable_node);
}

static void set_task_runnable(struct rq *rq, struct task_struct *p)
{
 lockdep_assert_rq_held(rq);

 if (p->scx.flags & SCX_TASK_RESET_RUNNABLE_AT) {
  p->scx.runnable_at = jiffies;
  p->scx.flags &= ~SCX_TASK_RESET_RUNNABLE_AT;
 }

 /*
 * list_add_tail() must be used. scx_bypass() depends on tasks being
 * appended to the runnable_list.
 */
 list_add_tail(&p->scx.runnable_node, &rq->scx.runnable_list);
}

static void clr_task_runnable(struct task_struct *p, bool reset_runnable_at)
{
 list_del_init(&p->scx.runnable_node);
 if (reset_runnable_at)
  p->scx.flags |= SCX_TASK_RESET_RUNNABLE_AT;
}

static void enqueue_task_scx(struct rq *rq, struct task_struct *p, int enq_flags)
{
 struct scx_sched *sch = scx_root;
 int sticky_cpu = p->scx.sticky_cpu;

 if (enq_flags & ENQUEUE_WAKEUP)
  rq->scx.flags |= SCX_RQ_IN_WAKEUP;

 enq_flags |= rq->scx.extra_enq_flags;

 if (sticky_cpu >= 0)
  p->scx.sticky_cpu = -1;

 /*
 * Restoring a running task will be immediately followed by
 * set_next_task_scx() which expects the task to not be on the BPF
 * scheduler as tasks can only start running through local DSQs. Force
 * direct-dispatch into the local DSQ by setting the sticky_cpu.
 */
 if (unlikely(enq_flags & ENQUEUE_RESTORE) && task_current(rq, p))
  sticky_cpu = cpu_of(rq);

 if (p->scx.flags & SCX_TASK_QUEUED) {
  WARN_ON_ONCE(!task_runnable(p));
  goto out;
 }

 set_task_runnable(rq, p);
 p->scx.flags |= SCX_TASK_QUEUED;
 rq->scx.nr_running++;
 add_nr_running(rq, 1);

 if (SCX_HAS_OP(sch, runnable) && !task_on_rq_migrating(p))
  SCX_CALL_OP_TASK(sch, SCX_KF_REST, runnable, rq, p, enq_flags);

 if (enq_flags & SCX_ENQ_WAKEUP)
  touch_core_sched(rq, p);

 do_enqueue_task(rq, p, enq_flags, sticky_cpu);
out:
 rq->scx.flags &= ~SCX_RQ_IN_WAKEUP;

 if ((enq_flags & SCX_ENQ_CPU_SELECTED) &&
     unlikely(cpu_of(rq) != p->scx.selected_cpu))
  __scx_add_event(sch, SCX_EV_SELECT_CPU_FALLBACK, 1);
}

static void ops_dequeue(struct rq *rq, struct task_struct *p, u64 deq_flags)
{
 struct scx_sched *sch = scx_root;
 unsigned long opss;

 /* dequeue is always temporary, don't reset runnable_at */
 clr_task_runnable(p, false);

 /* acquire ensures that we see the preceding updates on QUEUED */
 opss = atomic_long_read_acquire(&p->scx.ops_state);

 switch (opss & SCX_OPSS_STATE_MASK) {
 case SCX_OPSS_NONE:
  break;
 case SCX_OPSS_QUEUEING:
  /*
 * QUEUEING is started and finished while holding @p's rq lock.
 * As we're holding the rq lock now, we shouldn't see QUEUEING.
 */
  BUG();
 case SCX_OPSS_QUEUED:
  if (SCX_HAS_OP(sch, dequeue))
   SCX_CALL_OP_TASK(sch, SCX_KF_REST, dequeue, rq,
      p, deq_flags);

  if (atomic_long_try_cmpxchg(&p->scx.ops_state, &opss,
         SCX_OPSS_NONE))
   break;
  fallthrough;
 case SCX_OPSS_DISPATCHING:
  /*
 * If @p is being dispatched from the BPF scheduler to a DSQ,
 * wait for the transfer to complete so that @p doesn't get
 * added to its DSQ after dequeueing is complete.
 *
 * As we're waiting on DISPATCHING with the rq locked, the
 * dispatching side shouldn't try to lock the rq while
 * DISPATCHING is set. See dispatch_to_local_dsq().
 *
 * DISPATCHING shouldn't have qseq set and control can reach
 * here with NONE @opss from the above QUEUED case block.
 * Explicitly wait on %SCX_OPSS_DISPATCHING instead of @opss.
 */
  wait_ops_state(p, SCX_OPSS_DISPATCHING);
  BUG_ON(atomic_long_read(&p->scx.ops_state) != SCX_OPSS_NONE);
  break;
 }
}

static bool dequeue_task_scx(struct rq *rq, struct task_struct *p, int deq_flags)
{
 struct scx_sched *sch = scx_root;

 if (!(p->scx.flags & SCX_TASK_QUEUED)) {
  WARN_ON_ONCE(task_runnable(p));
  return true;
 }

 ops_dequeue(rq, p, deq_flags);

 /*
 * A currently running task which is going off @rq first gets dequeued
 * and then stops running. As we want running <-> stopping transitions
 * to be contained within runnable <-> quiescent transitions, trigger
 * ->stopping() early here instead of in put_prev_task_scx().
 *
 * @p may go through multiple stopping <-> running transitions between
 * here and put_prev_task_scx() if task attribute changes occur while
 * balance_scx() leaves @rq unlocked. However, they don't contain any
 * information meaningful to the BPF scheduler and can be suppressed by
 * skipping the callbacks if the task is !QUEUED.
 */
 if (SCX_HAS_OP(sch, stopping) && task_current(rq, p)) {
  update_curr_scx(rq);
  SCX_CALL_OP_TASK(sch, SCX_KF_REST, stopping, rq, p, false);
 }

 if (SCX_HAS_OP(sch, quiescent) && !task_on_rq_migrating(p))
  SCX_CALL_OP_TASK(sch, SCX_KF_REST, quiescent, rq, p, deq_flags);

 if (deq_flags & SCX_DEQ_SLEEP)
  p->scx.flags |= SCX_TASK_DEQD_FOR_SLEEP;
 else
  p->scx.flags &= ~SCX_TASK_DEQD_FOR_SLEEP;

 p->scx.flags &= ~SCX_TASK_QUEUED;
 rq->scx.nr_running--;
 sub_nr_running(rq, 1);

 dispatch_dequeue(rq, p);
 return true;
}

static void yield_task_scx(struct rq *rq)
{
 struct scx_sched *sch = scx_root;
 struct task_struct *p = rq->curr;

 if (SCX_HAS_OP(sch, yield))
  SCX_CALL_OP_2TASKS_RET(sch, SCX_KF_REST, yield, rq, p, NULL);
 else
  p->scx.slice = 0;
}

static bool yield_to_task_scx(struct rq *rq, struct task_struct *to)
{
 struct scx_sched *sch = scx_root;
 struct task_struct *from = rq->curr;

 if (SCX_HAS_OP(sch, yield))
  return SCX_CALL_OP_2TASKS_RET(sch, SCX_KF_REST, yield, rq,
           from, to);
 else
  return false;
}

static void move_local_task_to_local_dsq(struct task_struct *p, u64 enq_flags,
      struct scx_dispatch_q *src_dsq,
      struct rq *dst_rq)
{
 struct scx_dispatch_q *dst_dsq = &dst_rq->scx.local_dsq;

 /* @dsq is locked and @p is on @dst_rq */
 lockdep_assert_held(&src_dsq->lock);
 lockdep_assert_rq_held(dst_rq);

 WARN_ON_ONCE(p->scx.holding_cpu >= 0);

 if (enq_flags & (SCX_ENQ_HEAD | SCX_ENQ_PREEMPT))
  list_add(&p->scx.dsq_list.node, &dst_dsq->list);
 else
  list_add_tail(&p->scx.dsq_list.node, &dst_dsq->list);

 dsq_mod_nr(dst_dsq, 1);
 p->scx.dsq = dst_dsq;
}

/**
 * move_remote_task_to_local_dsq - Move a task from a foreign rq to a local DSQ
 * @p: task to move
 * @enq_flags: %SCX_ENQ_*
 * @src_rq: rq to move the task from, locked on entry, released on return
 * @dst_rq: rq to move the task into, locked on return
 *
 * Move @p which is currently on @src_rq to @dst_rq's local DSQ.
 */
static void move_remote_task_to_local_dsq(struct task_struct *p, u64 enq_flags,
       struct rq *src_rq, struct rq *dst_rq)
{
 lockdep_assert_rq_held(src_rq);

 /* the following marks @p MIGRATING which excludes dequeue */
 deactivate_task(src_rq, p, 0);
 set_task_cpu(p, cpu_of(dst_rq));
 p->scx.sticky_cpu = cpu_of(dst_rq);

 raw_spin_rq_unlock(src_rq);
 raw_spin_rq_lock(dst_rq);

 /*
 * We want to pass scx-specific enq_flags but activate_task() will
 * truncate the upper 32 bit. As we own @rq, we can pass them through
 * @rq->scx.extra_enq_flags instead.
 */
 WARN_ON_ONCE(!cpumask_test_cpu(cpu_of(dst_rq), p->cpus_ptr));
 WARN_ON_ONCE(dst_rq->scx.extra_enq_flags);
 dst_rq->scx.extra_enq_flags = enq_flags;
 activate_task(dst_rq, p, 0);
 dst_rq->scx.extra_enq_flags = 0;
}

/*
 * Similar to kernel/sched/core.c::is_cpu_allowed(). However, there are two
 * differences:
 *
 * - is_cpu_allowed() asks "Can this task run on this CPU?" while
 *   task_can_run_on_remote_rq() asks "Can the BPF scheduler migrate the task to
 *   this CPU?".
 *
 *   While migration is disabled, is_cpu_allowed() has to say "yes" as the task
 *   must be allowed to finish on the CPU that it's currently on regardless of
 *   the CPU state. However, task_can_run_on_remote_rq() must say "no" as the
 *   BPF scheduler shouldn't attempt to migrate a task which has migration
 *   disabled.
 *
 * - The BPF scheduler is bypassed while the rq is offline and we can always say
 *   no to the BPF scheduler initiated migrations while offline.
 *
 * The caller must ensure that @p and @rq are on different CPUs.
 */
static bool task_can_run_on_remote_rq(struct scx_sched *sch,
          struct task_struct *p, struct rq *rq,
          bool enforce)
{
 int cpu = cpu_of(rq);

 WARN_ON_ONCE(task_cpu(p) == cpu);

 /*
 * If @p has migration disabled, @p->cpus_ptr is updated to contain only
 * the pinned CPU in migrate_disable_switch() while @p is being switched
 * out. However, put_prev_task_scx() is called before @p->cpus_ptr is
 * updated and thus another CPU may see @p on a DSQ inbetween leading to
 * @p passing the below task_allowed_on_cpu() check while migration is
 * disabled.
 *
 * Test the migration disabled state first as the race window is narrow
 * and the BPF scheduler failing to check migration disabled state can
 * easily be masked if task_allowed_on_cpu() is done first.
 */
 if (unlikely(is_migration_disabled(p))) {
  if (enforce)
   scx_error(sch, "SCX_DSQ_LOCAL[_ON] cannot move migration disabled %s[%d] from CPU %d to %d",
      p->comm, p->pid, task_cpu(p), cpu);
  return false;
 }

 /*
 * We don't require the BPF scheduler to avoid dispatching to offline
 * CPUs mostly for convenience but also because CPUs can go offline
 * between scx_bpf_dsq_insert() calls and here. Trigger error iff the
 * picked CPU is outside the allowed mask.
 */
 if (!task_allowed_on_cpu(p, cpu)) {
  if (enforce)
   scx_error(sch, "SCX_DSQ_LOCAL[_ON] target CPU %d not allowed for %s[%d]",
      cpu, p->comm, p->pid);
  return false;
 }

 if (!scx_rq_online(rq)) {
  if (enforce)
   __scx_add_event(scx_root,
     SCX_EV_DISPATCH_LOCAL_DSQ_OFFLINE, 1);
  return false;
 }

 return true;
}

/**
 * unlink_dsq_and_lock_src_rq() - Unlink task from its DSQ and lock its task_rq
 * @p: target task
 * @dsq: locked DSQ @p is currently on
 * @src_rq: rq @p is currently on, stable with @dsq locked
 *
 * Called with @dsq locked but no rq's locked. We want to move @p to a different
 * DSQ, including any local DSQ, but are not locking @src_rq. Locking @src_rq is
 * required when transferring into a local DSQ. Even when transferring into a
 * non-local DSQ, it's better to use the same mechanism to protect against
 * dequeues and maintain the invariant that @p->scx.dsq can only change while
 * @src_rq is locked, which e.g. scx_dump_task() depends on.
 *
 * We want to grab @src_rq but that can deadlock if we try while locking @dsq,
 * so we want to unlink @p from @dsq, drop its lock and then lock @src_rq. As
 * this may race with dequeue, which can't drop the rq lock or fail, do a little
 * dancing from our side.
 *
 * @p->scx.holding_cpu is set to this CPU before @dsq is unlocked. If @p gets
 * dequeued after we unlock @dsq but before locking @src_rq, the holding_cpu
 * would be cleared to -1. While other cpus may have updated it to different
 * values afterwards, as this operation can't be preempted or recurse, the
 * holding_cpu can never become this CPU again before we're done. Thus, we can
 * tell whether we lost to dequeue by testing whether the holding_cpu still
 * points to this CPU. See dispatch_dequeue() for the counterpart.
 *
 * On return, @dsq is unlocked and @src_rq is locked. Returns %true if @p is
 * still valid. %false if lost to dequeue.
 */
static bool unlink_dsq_and_lock_src_rq(struct task_struct *p,
           struct scx_dispatch_q *dsq,
           struct rq *src_rq)
{
 s32 cpu = raw_smp_processor_id();

 lockdep_assert_held(&dsq->lock);

 WARN_ON_ONCE(p->scx.holding_cpu >= 0);
 task_unlink_from_dsq(p, dsq);
 p->scx.holding_cpu = cpu;

 raw_spin_unlock(&dsq->lock);
 raw_spin_rq_lock(src_rq);

 /* task_rq couldn't have changed if we're still the holding cpu */
 return likely(p->scx.holding_cpu == cpu) &&
  !WARN_ON_ONCE(src_rq != task_rq(p));
}

static bool consume_remote_task(struct rq *this_rq, struct task_struct *p,
    struct scx_dispatch_q *dsq, struct rq *src_rq)
{
 raw_spin_rq_unlock(this_rq);

 if (unlink_dsq_and_lock_src_rq(p, dsq, src_rq)) {
  move_remote_task_to_local_dsq(p, 0, src_rq, this_rq);
  return true;
 } else {
  raw_spin_rq_unlock(src_rq);
  raw_spin_rq_lock(this_rq);
  return false;
 }
}

/**
 * move_task_between_dsqs() - Move a task from one DSQ to another
 * @sch: scx_sched being operated on
 * @p: target task
 * @enq_flags: %SCX_ENQ_*
 * @src_dsq: DSQ @p is currently on, must not be a local DSQ
 * @dst_dsq: DSQ @p is being moved to, can be any DSQ
 *
 * Must be called with @p's task_rq and @src_dsq locked. If @dst_dsq is a local
 * DSQ and @p is on a different CPU, @p will be migrated and thus its task_rq
 * will change. As @p's task_rq is locked, this function doesn't need to use the
 * holding_cpu mechanism.
 *
 * On return, @src_dsq is unlocked and only @p's new task_rq, which is the
 * return value, is locked.
 */
static struct rq *move_task_between_dsqs(struct scx_sched *sch,
      struct task_struct *p, u64 enq_flags,
      struct scx_dispatch_q *src_dsq,
      struct scx_dispatch_q *dst_dsq)
{
 struct rq *src_rq = task_rq(p), *dst_rq;

 BUG_ON(src_dsq->id == SCX_DSQ_LOCAL);
 lockdep_assert_held(&src_dsq->lock);
 lockdep_assert_rq_held(src_rq);

 if (dst_dsq->id == SCX_DSQ_LOCAL) {
  dst_rq = container_of(dst_dsq, struct rq, scx.local_dsq);
  if (src_rq != dst_rq &&
      unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq, true))) {
   dst_dsq = find_global_dsq(p);
   dst_rq = src_rq;
  }
 } else {
  /* no need to migrate if destination is a non-local DSQ */
  dst_rq = src_rq;
 }

 /*
 * Move @p into $dst_dsq. If $dst_dsq is the local DSQ of a different
 * CPU, @p will be migrated.
 */
 if (dst_dsq->id == SCX_DSQ_LOCAL) {
  /* @p is going from a non-local DSQ to a local DSQ */
  if (src_rq == dst_rq) {
   task_unlink_from_dsq(p, src_dsq);
   move_local_task_to_local_dsq(p, enq_flags,
           src_dsq, dst_rq);
   raw_spin_unlock(&src_dsq->lock);
  } else {
   raw_spin_unlock(&src_dsq->lock);
   move_remote_task_to_local_dsq(p, enq_flags,
            src_rq, dst_rq);
  }
 } else {
  /*
 * @p is going from a non-local DSQ to a non-local DSQ. As
 * $src_dsq is already locked, do an abbreviated dequeue.
 */
  task_unlink_from_dsq(p, src_dsq);
  p->scx.dsq = NULL;
  raw_spin_unlock(&src_dsq->lock);

  dispatch_enqueue(sch, dst_dsq, p, enq_flags);
 }

 return dst_rq;
}

/*
 * A poorly behaving BPF scheduler can live-lock the system by e.g. incessantly
 * banging on the same DSQ on a large NUMA system to the point where switching
 * to the bypass mode can take a long time. Inject artificial delays while the
 * bypass mode is switching to guarantee timely completion.
 */
static void scx_breather(struct rq *rq)
{
 u64 until;

 lockdep_assert_rq_held(rq);

 if (likely(!atomic_read(&scx_breather_depth)))
  return;

 raw_spin_rq_unlock(rq);

 until = ktime_get_ns() + NSEC_PER_MSEC;

 do {
  int cnt = 1024;
  while (atomic_read(&scx_breather_depth) && --cnt)
   cpu_relax();
 } while (atomic_read(&scx_breather_depth) &&
   time_before64(ktime_get_ns(), until));

 raw_spin_rq_lock(rq);
}

static bool consume_dispatch_q(struct scx_sched *sch, struct rq *rq,
          struct scx_dispatch_q *dsq)
{
 struct task_struct *p;
retry:
 /*
 * This retry loop can repeatedly race against scx_bypass() dequeueing
 * tasks from @dsq trying to put the system into the bypass mode. On
 * some multi-socket machines (e.g. 2x Intel 8480c), this can live-lock
 * the machine into soft lockups. Give a breather.
 */
 scx_breather(rq);

 /*
 * The caller can't expect to successfully consume a task if the task's
 * addition to @dsq isn't guaranteed to be visible somehow. Test
 * @dsq->list without locking and skip if it seems empty.
 */
 if (list_empty(&dsq->list))
  return false;

 raw_spin_lock(&dsq->lock);

 nldsq_for_each_task(p, dsq) {
  struct rq *task_rq = task_rq(p);

  if (rq == task_rq) {
   task_unlink_from_dsq(p, dsq);
   move_local_task_to_local_dsq(p, 0, dsq, rq);
   raw_spin_unlock(&dsq->lock);
   return true;
  }

  if (task_can_run_on_remote_rq(sch, p, rq, false)) {
   if (likely(consume_remote_task(rq, p, dsq, task_rq)))
    return true;
   goto retry;
  }
 }

 raw_spin_unlock(&dsq->lock);
 return false;
}

static bool consume_global_dsq(struct scx_sched *sch, struct rq *rq)
{
 int node = cpu_to_node(cpu_of(rq));

 return consume_dispatch_q(sch, rq, sch->global_dsqs[node]);
}

/**
 * dispatch_to_local_dsq - Dispatch a task to a local dsq
 * @sch: scx_sched being operated on
 * @rq: current rq which is locked
 * @dst_dsq: destination DSQ
 * @p: task to dispatch
 * @enq_flags: %SCX_ENQ_*
 *
 * We're holding @rq lock and want to dispatch @p to @dst_dsq which is a local
 * DSQ. This function performs all the synchronization dancing needed because
 * local DSQs are protected with rq locks.
 *
 * The caller must have exclusive ownership of @p (e.g. through
 * %SCX_OPSS_DISPATCHING).
 */
static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
      struct scx_dispatch_q *dst_dsq,
      struct task_struct *p, u64 enq_flags)
{
 struct rq *src_rq = task_rq(p);
 struct rq *dst_rq = container_of(dst_dsq, struct rq, scx.local_dsq);
 struct rq *locked_rq = rq;

 /*
 * We're synchronized against dequeue through DISPATCHING. As @p can't
 * be dequeued, its task_rq and cpus_allowed are stable too.
 *
 * If dispatching to @rq that @p is already on, no lock dancing needed.
 */
 if (rq == src_rq && rq == dst_rq) {
  dispatch_enqueue(sch, dst_dsq, p,
     enq_flags | SCX_ENQ_CLEAR_OPSS);
  return;
 }

 if (src_rq != dst_rq &&
     unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq, true))) {
  dispatch_enqueue(sch, find_global_dsq(p), p,
     enq_flags | SCX_ENQ_CLEAR_OPSS);
  return;
 }

 /*
 * @p is on a possibly remote @src_rq which we need to lock to move the
 * task. If dequeue is in progress, it'd be locking @src_rq and waiting
 * on DISPATCHING, so we can't grab @src_rq lock while holding
 * DISPATCHING.
 *
 * As DISPATCHING guarantees that @p is wholly ours, we can pretend that
 * we're moving from a DSQ and use the same mechanism - mark the task
 * under transfer with holding_cpu, release DISPATCHING and then follow
 * the same protocol. See unlink_dsq_and_lock_src_rq().
 */
 p->scx.holding_cpu = raw_smp_processor_id();

 /* store_release ensures that dequeue sees the above */
 atomic_long_set_release(&p->scx.ops_state, SCX_OPSS_NONE);

 /* switch to @src_rq lock */
 if (locked_rq != src_rq) {
  raw_spin_rq_unlock(locked_rq);
  locked_rq = src_rq;
  raw_spin_rq_lock(src_rq);
 }

 /* task_rq couldn't have changed if we're still the holding cpu */
 if (likely(p->scx.holding_cpu == raw_smp_processor_id()) &&
     !WARN_ON_ONCE(src_rq != task_rq(p))) {
  /*
 * If @p is staying on the same rq, there's no need to go
 * through the full deactivate/activate cycle. Optimize by
 * abbreviating move_remote_task_to_local_dsq().
 */
  if (src_rq == dst_rq) {
   p->scx.holding_cpu = -1;
   dispatch_enqueue(sch, &dst_rq->scx.local_dsq, p,
      enq_flags);
  } else {
   move_remote_task_to_local_dsq(p, enq_flags,
            src_rq, dst_rq);
   /* task has been moved to dst_rq, which is now locked */
   locked_rq = dst_rq;
  }

  /* if the destination CPU is idle, wake it up */
  if (sched_class_above(p->sched_class, dst_rq->curr->sched_class))
   resched_curr(dst_rq);
 }

 /* switch back to @rq lock */
 if (locked_rq != rq) {
  raw_spin_rq_unlock(locked_rq);
  raw_spin_rq_lock(rq);
 }
}

/**
 * finish_dispatch - Asynchronously finish dispatching a task
 * @rq: current rq which is locked
 * @p: task to finish dispatching
 * @qseq_at_dispatch: qseq when @p started getting dispatched
 * @dsq_id: destination DSQ ID
 * @enq_flags: %SCX_ENQ_*
 *
 * Dispatching to local DSQs may need to wait for queueing to complete or
 * require rq lock dancing. As we don't wanna do either while inside
 * ops.dispatch() to avoid locking order inversion, we split dispatching into
 * two parts. scx_bpf_dsq_insert() which is called by ops.dispatch() records the
 * task and its qseq. Once ops.dispatch() returns, this function is called to
 * finish up.
 *
 * There is no guarantee that @p is still valid for dispatching or even that it
 * was valid in the first place. Make sure that the task is still owned by the
 * BPF scheduler and claim the ownership before dispatching.
 */
static void finish_dispatch(struct scx_sched *sch, struct rq *rq,
       struct task_struct *p,
       unsigned long qseq_at_dispatch,
       u64 dsq_id, u64 enq_flags)
{
 struct scx_dispatch_q *dsq;
 unsigned long opss;

 touch_core_sched_dispatch(rq, p);
retry:
 /*
 * No need for _acquire here. @p is accessed only after a successful
 * try_cmpxchg to DISPATCHING.
 */
 opss = atomic_long_read(&p->scx.ops_state);

 switch (opss & SCX_OPSS_STATE_MASK) {
 case SCX_OPSS_DISPATCHING:
 case SCX_OPSS_NONE:
  /* someone else already got to it */
  return;
 case SCX_OPSS_QUEUED:
  /*
 * If qseq doesn't match, @p has gone through at least one
 * dispatch/dequeue and re-enqueue cycle between
 * scx_bpf_dsq_insert() and here and we have no claim on it.
 */
  if ((opss & SCX_OPSS_QSEQ_MASK) != qseq_at_dispatch)
   return;

  /*
 * While we know @p is accessible, we don't yet have a claim on
 * it - the BPF scheduler is allowed to dispatch tasks
 * spuriously and there can be a racing dequeue attempt. Let's
 * claim @p by atomically transitioning it from QUEUED to
 * DISPATCHING.
 */
  if (likely(atomic_long_try_cmpxchg(&p->scx.ops_state, &opss,
         SCX_OPSS_DISPATCHING)))
   break;
  goto retry;
 case SCX_OPSS_QUEUEING:
  /*
 * do_enqueue_task() is in the process of transferring the task
 * to the BPF scheduler while holding @p's rq lock. As we aren't
 * holding any kernel or BPF resource that the enqueue path may
 * depend upon, it's safe to wait.
 */
  wait_ops_state(p, opss);
  goto retry;
 }

 BUG_ON(!(p->scx.flags & SCX_TASK_QUEUED));

 dsq = find_dsq_for_dispatch(sch, this_rq(), dsq_id, p);

 if (dsq->id == SCX_DSQ_LOCAL)
  dispatch_to_local_dsq(sch, rq, dsq, p, enq_flags);
 else
  dispatch_enqueue(sch, dsq, p, enq_flags | SCX_ENQ_CLEAR_OPSS);
}

static void flush_dispatch_buf(struct scx_sched *sch, struct rq *rq)
{
 struct scx_dsp_ctx *dspc = this_cpu_ptr(scx_dsp_ctx);
 u32 u;

 for (u = 0; u < dspc->cursor; u++) {
  struct scx_dsp_buf_ent *ent = &dspc->buf[u];

  finish_dispatch(sch, rq, ent->task, ent->qseq, ent->dsq_id,
    ent->enq_flags);
 }

 dspc->nr_tasks += dspc->cursor;
 dspc->cursor = 0;
}

static int balance_one(struct rq *rq, struct task_struct *prev)
{
 struct scx_sched *sch = scx_root;
 struct scx_dsp_ctx *dspc = this_cpu_ptr(scx_dsp_ctx);
 bool prev_on_scx = prev->sched_class == &ext_sched_class;
 bool prev_on_rq = prev->scx.flags & SCX_TASK_QUEUED;
 int nr_loops = SCX_DSP_MAX_LOOPS;

 lockdep_assert_rq_held(rq);
 rq->scx.flags |= SCX_RQ_IN_BALANCE;
 rq->scx.flags &= ~(SCX_RQ_BAL_PENDING | SCX_RQ_BAL_KEEP);

 if ((sch->ops.flags & SCX_OPS_HAS_CPU_PREEMPT) &&
     unlikely(rq->scx.cpu_released)) {
  /*
 * If the previous sched_class for the current CPU was not SCX,
 * notify the BPF scheduler that it again has control of the
 * core. This callback complements ->cpu_release(), which is
 * emitted in switch_class().
 */
  if (SCX_HAS_OP(sch, cpu_acquire))
   SCX_CALL_OP(sch, SCX_KF_REST, cpu_acquire, rq,
        cpu_of(rq), NULL);
  rq->scx.cpu_released = false;
 }

 if (prev_on_scx) {
  update_curr_scx(rq);

  /*
 * If @prev is runnable & has slice left, it has priority and
 * fetching more just increases latency for the fetched tasks.
 * Tell pick_task_scx() to keep running @prev. If the BPF
 * scheduler wants to handle this explicitly, it should
 * implement ->cpu_release().
 *
 * See scx_disable_workfn() for the explanation on the bypassing
 * test.
 */
  if (prev_on_rq && prev->scx.slice && !scx_rq_bypassing(rq)) {
   rq->scx.flags |= SCX_RQ_BAL_KEEP;
   goto has_tasks;
  }
 }

 /* if there already are tasks to run, nothing to do */
 if (rq->scx.local_dsq.nr)
  goto has_tasks;

 if (consume_global_dsq(sch, rq))
  goto has_tasks;

 if (unlikely(!SCX_HAS_OP(sch, dispatch)) ||
     scx_rq_bypassing(rq) || !scx_rq_online(rq))
  goto no_tasks;

 dspc->rq = rq;

 /*
 * The dispatch loop. Because flush_dispatch_buf() may drop the rq lock,
 * the local DSQ might still end up empty after a successful
 * ops.dispatch(). If the local DSQ is empty even after ops.dispatch()
 * produced some tasks, retry. The BPF scheduler may depend on this
 * looping behavior to simplify its implementation.
 */
 do {
  dspc->nr_tasks = 0;

  SCX_CALL_OP(sch, SCX_KF_DISPATCH, dispatch, rq,
       cpu_of(rq), prev_on_scx ? prev : NULL);

  flush_dispatch_buf(sch, rq);

  if (prev_on_rq && prev->scx.slice) {
   rq->scx.flags |= SCX_RQ_BAL_KEEP;
   goto has_tasks;
  }
  if (rq->scx.local_dsq.nr)
   goto has_tasks;
  if (consume_global_dsq(sch, rq))
   goto has_tasks;

  /*
 * ops.dispatch() can trap us in this loop by repeatedly
 * dispatching ineligible tasks. Break out once in a while to
 * allow the watchdog to run. As IRQ can't be enabled in
 * balance(), we want to complete this scheduling cycle and then
 * start a new one. IOW, we want to call resched_curr() on the
 * next, most likely idle, task, not the current one. Use
 * scx_bpf_kick_cpu() for deferred kicking.
 */
  if (unlikely(!--nr_loops)) {
   scx_bpf_kick_cpu(cpu_of(rq), 0);
   break;
  }
 } while (dspc->nr_tasks);

no_tasks:
 /*
 * Didn't find another task to run. Keep running @prev unless
 * %SCX_OPS_ENQ_LAST is in effect.
 */
 if (prev_on_rq &&
     (!(sch->ops.flags & SCX_OPS_ENQ_LAST) || scx_rq_bypassing(rq))) {
  rq->scx.flags |= SCX_RQ_BAL_KEEP;
  __scx_add_event(sch, SCX_EV_DISPATCH_KEEP_LAST, 1);
  goto has_tasks;
 }
 rq->scx.flags &= ~SCX_RQ_IN_BALANCE;
 return false;

has_tasks:
 rq->scx.flags &= ~SCX_RQ_IN_BALANCE;
 return true;
}

static int balance_scx(struct rq *rq, struct task_struct *prev,
         struct rq_flags *rf)
{
 int ret;

 rq_unpin_lock(rq, rf);

 ret = balance_one(rq, prev);

#ifdef CONFIG_SCHED_SMT
 /*
 * When core-sched is enabled, this ops.balance() call will be followed
 * by pick_task_scx() on this CPU and the SMT siblings. Balance the
 * siblings too.
 */
 if (sched_core_enabled(rq)) {
  const struct cpumask *smt_mask = cpu_smt_mask(cpu_of(rq));
  int scpu;

  for_each_cpu_andnot(scpu, smt_mask, cpumask_of(cpu_of(rq))) {
   struct rq *srq = cpu_rq(scpu);
   struct task_struct *sprev = srq->curr;

   WARN_ON_ONCE(__rq_lockp(rq) != __rq_lockp(srq));
   update_rq_clock(srq);
   balance_one(srq, sprev);
  }
 }
#endif
 rq_repin_lock(rq, rf);

 return ret;
}

static void process_ddsp_deferred_locals(struct rq *rq)
{
 struct task_struct *p;

 lockdep_assert_rq_held(rq);

 /*
 * Now that @rq can be unlocked, execute the deferred enqueueing of
 * tasks directly dispatched to the local DSQs of other CPUs. See
 * direct_dispatch(). Keep popping from the head instead of using
 * list_for_each_entry_safe() as dispatch_local_dsq() may unlock @rq
 * temporarily.
 */
 while ((p = list_first_entry_or_null(&rq->scx.ddsp_deferred_locals,
    struct task_struct, scx.dsq_list.node))) {
  struct scx_sched *sch = scx_root;
  struct scx_dispatch_q *dsq;

  list_del_init(&p->scx.dsq_list.node);

  dsq = find_dsq_for_dispatch(sch, rq, p->scx.ddsp_dsq_id, p);
  if (!WARN_ON_ONCE(dsq->id != SCX_DSQ_LOCAL))
--> --------------------

--> maximum size reached

--> --------------------