sched_ext: Preserve rq tracking across local DSQ dispatch

dispatch_to_local_dsq() can run from scx_bpf_dsq_move_to_local() while
ops.dispatch() has recorded the current rq. Moving a task to a local DSQ
may switch to the source or destination rq before synchronously invoking
ops.dequeue() through the following path:

  SCX_CALL_OP(dispatch, rq)
    ops.dispatch()
      scx_bpf_dsq_move_to_local()
        scx_flush_dispatch_buf()
          finish_dispatch()
            dispatch_to_local_dsq()
              scx_dispatch_enqueue()
                local_dsq_post_enq()
                  call_task_dequeue()
                    SCX_CALL_OP_TASK(dequeue, locked_rq, ...)

The nested callback saves the recorded rq and restores it on return. If
the rq tracking does not follow the lock switch, update_locked_rq() can
trigger the following lockdep assertion while restoring an rq which is
no longer held:

  WARNING: kernel/sched/sched.h:1641 at call_task_dequeue+0x160/0x170
  Call Trace:
    scx_dispatch_enqueue+0x2b0/0x460
    dispatch_to_local_dsq+0x138/0x230
    scx_flush_dispatch_buf+0x1af/0x220
    scx_bpf_dsq_move_to_local___v2+0xe2/0x1c0
    bpf__sched_ext_ops_dispatch+0x4b/0xa7
    do_pick_task_scx+0x3b6/0x910
    __pick_next_task+0x105/0x1f0
    __schedule+0x3e7/0x1980

Introduce switch_rq_lock() to update the tracking state together with
each rq lock handoff. Use it in dispatch_to_local_dsq(),
move_remote_task_to_local_dsq() and the in-balance paths of
scx_dsq_move(), ensuring that scx_locked_rq() consistently refers to the
rq whose lock is actually held throughout the lock dance.

Fixes: 7fb39e4eb4c3 ("sched_ext: Save and restore scx_locked_rq across SCX_CALL_OP")
Cc: stable@vger.kernel.org # 7.1+
Signed-off-by: Andrea Righi <arighi@nvidia.com>
Signed-off-by: Tejun Heo <tj@kernel.org>

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index e75e2fd5ab7e33cfbea9504a32e2081aea952ff6..ee66dacfc92c32dcbaa270a99c3f23fe5605356b 100644 (file)
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -479,6 +479,18 @@ static bool rq_is_open(struct rq *rq, u64 enq_flags)
  */
 DEFINE_PER_CPU(struct rq *, scx_locked_rq_state);
 
+static void switch_rq_lock(struct rq *from, struct rq *to)
+{
+       bool tracked = scx_locked_rq() == from;
+
+       if (tracked)
+               update_locked_rq(NULL);
+       raw_spin_rq_unlock(from);
+       raw_spin_rq_lock(to);
+       if (tracked)
+               update_locked_rq(to);
+}
+
 /*
  * Flipped on enable per sch->is_cid_type. Declared in internal.h so
  * subsystem inlines can read it.
@@ -2274,8 +2286,7 @@ static void move_remote_task_to_local_dsq(struct task_struct *p, u64 enq_flags,
        deactivate_task(src_rq, p, 0);
        set_task_cpu(p, cpu_of(dst_rq));
 
-       raw_spin_rq_unlock(src_rq);
-       raw_spin_rq_lock(dst_rq);
+       switch_rq_lock(src_rq, dst_rq);
 
        /*
         * We want to pass scx-specific enq_flags but activate_task() will
@@ -2608,9 +2619,8 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
 
        /* switch to @src_rq lock */
        if (locked_rq != src_rq) {
-               raw_spin_rq_unlock(locked_rq);
+               switch_rq_lock(locked_rq, src_rq);
                locked_rq = src_rq;
-               raw_spin_rq_lock(src_rq);
        }
 
        /* task_rq couldn't have changed if we're still the holding cpu */
@@ -2644,10 +2654,8 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
        }
 
        /* switch back to @rq lock */
-       if (locked_rq != rq) {
-               raw_spin_rq_unlock(locked_rq);
-               raw_spin_rq_lock(rq);
-       }
+       if (locked_rq != rq)
+               switch_rq_lock(locked_rq, rq);
 }
 
 /**
@@ -8825,10 +8833,8 @@ static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
        in_balance = this_rq->scx.flags & SCX_RQ_IN_BALANCE;
 
        if (in_balance) {
-               if (this_rq != src_rq) {
-                       raw_spin_rq_unlock(this_rq);
-                       raw_spin_rq_lock(src_rq);
-               }
+               if (this_rq != src_rq)
+                       switch_rq_lock(this_rq, src_rq);
        } else {
                raw_spin_rq_lock(src_rq);
        }
@@ -8860,10 +8866,8 @@ static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
        dispatched = true;
 out:
        if (in_balance) {
-               if (this_rq != locked_rq) {
-                       raw_spin_rq_unlock(locked_rq);
-                       raw_spin_rq_lock(this_rq);
-               }
+               if (this_rq != locked_rq)
+                       switch_rq_lock(locked_rq, this_rq);
        } else {
                raw_spin_rq_unlock_irqrestore(locked_rq, flags);
        }