[RFC PATCH 1/5] sched_ext: Make legacy scx_bpf_reenqueue_local() target the locked rq

[Kuba Piecuch <jpiecuch@google.com>]

The legacy, synchronous scx_bpf_reenqueue_local() can only be called
from ops.cpu_release(). It assumes that ops.cpu_release() always runs on
the CPU being released and re-enqueues the tasks on the local DSQ of
smp_processor_id().

That assumption is about to stop holding. A following patch emits
ops.cpu_release() when the task running on a CPU leaves SCX through
sched_setscheduler() or a PI boost. That can happen from any CPU, with
the target CPU's rq locked. Picking the rq from smp_processor_id() would
then trip lockdep_assert_rq_held() and walk another CPU's local DSQ
without holding that CPU's rq lock.

SCX_CALL_OP() already records the rq the operation is invoked for in
scx_locked_rq(), and scx_bpf_dsq_reenq() /
scx_bpf_reenqueue_local___v2() already resolve SCX_DSQ_LOCAL against it.
Do the same in the legacy kfunc. Nothing else in reenq_local() depends
on running on the target CPU: the local DSQ is protected by the rq lock,
and direct dispatches to a third CPU's local DSQ are already deferred
through ddsp_deferred_locals.

No functional change for the existing callers, which all run on the CPU
whose rq is locked.

Also update the kerneldoc of scx_bpf_reenqueue_local___v2(), which still
says that it works on the local DSQ of the caller's CPU, to match
scx_bpf_dsq_reenq().

Signed-off-by: Kuba Piecuch <jpiecuch@google.com>
---
 kernel/sched/ext/ext.c | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index f1b9d52f624a..f64438290218 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -9501,9 +9501,10 @@ __bpf_kfunc_start_defs();
  * scx_bpf_reenqueue_local - Re-enqueue tasks on a local DSQ
  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
  *
- * Iterate over all of the tasks currently enqueued on the local DSQ of the
- * caller's CPU, and re-enqueue them in the BPF scheduler. Returns the number of
- * processed tasks. Can only be called from ops.cpu_release().
+ * Iterate over all of the tasks currently enqueued on the local DSQ of the CPU
+ * ops.cpu_release() is being invoked for, and re-enqueue them in the BPF
+ * scheduler. Returns the number of processed tasks. Can only be called from
+ * ops.cpu_release().
  */
 __bpf_kfunc u32 scx_bpf_reenqueue_local(const struct bpf_prog_aux *aux)
 {
@@ -9515,7 +9516,13 @@ __bpf_kfunc u32 scx_bpf_reenqueue_local(const struct bpf_prog_aux *aux)
 	if (unlikely(!sch))
 		return 0;
 
-	rq = cpu_rq(smp_processor_id());
+	/*
+	 * Use the rq ops.cpu_release() was invoked for rather than the local
+	 * one. They are not necessarily the same CPU.
+	 */
+	rq = scx_locked_rq();
+	if (WARN_ON_ONCE(!rq))
+		return 0;
 	lockdep_assert_rq_held(rq);
 
 	return reenq_local(sch, rq, SCX_REENQ_ANY);
@@ -10057,7 +10064,9 @@ __bpf_kfunc void scx_bpf_dsq_reenq(u64 dsq_id, u64 reenq_flags,
  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
  *
  * Iterate over all of the tasks currently enqueued on the local DSQ of the
- * caller's CPU, and re-enqueue them in the BPF scheduler. Can be called from
+ * rq the current scheduler operation is locked to - e.g. the CPU
+ * ops.cpu_release() is being invoked for - or the calling CPU's when no rq is
+ * locked, and re-enqueue them in the BPF scheduler. Can be called from
  * anywhere.
  *
  * This is now a special case of scx_bpf_dsq_reenq() and may be removed in the
-- 
2.56.0.176.g3575e65682-goog