[PATCH v2 sched_ext/for-7.2-fixes] sched_ext: Preserve rq tracking across local DSQ dispatch

[Andrea Righi <arighi@nvidia.com>]

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
dispatch_to_local_dsq() switched locks, that rq is no longer held and
update_locked_rq() can trigger the following lockdep assertion while
restoring it:

  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

Clear the tracked rq before switching locks and restore it only after
reacquiring the original rq.

scx_dsq_move() has the same problem when it drops this_rq while
ops.dispatch() is active, so apply the same tracking treatment there too.

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>
---
Changes in v2:
 - Apply the same rq tracking fix to scx_dsq_move() (sashiko AI)
 - Link to v1: https://lore.kernel.org/all/20260707135854.1379730-1-arighi@nvidia.com/

 kernel/sched/ext/ext.c | 29 ++++++++++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index e75e2fd5ab7e3..872d1e1f83d14 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -2577,6 +2577,7 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
 	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;
+	struct rq *tracked_rq = scx_locked_rq();
 
 	/*
 	 * We're synchronized against dequeue through DISPATCHING. As @p can't
@@ -2590,6 +2591,16 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
 		return;
 	}
 
+	/*
+	 * This may be called from an SCX op with @rq recorded as the currently
+	 * locked rq. Clear the tracking while switching rq locks so nested
+	 * callbacks don't restore an rq which isn't locked anymore.
+	 */
+	if (tracked_rq) {
+		WARN_ON_ONCE(tracked_rq != rq);
+		update_locked_rq(NULL);
+	}
+
 	/*
 	 * @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
@@ -2648,6 +2659,8 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
 		raw_spin_rq_unlock(locked_rq);
 		raw_spin_rq_lock(rq);
 	}
+	if (tracked_rq)
+		update_locked_rq(tracked_rq);
 }
 
 /**
@@ -8782,7 +8795,7 @@ static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
 {
 	struct scx_dispatch_q *src_dsq = kit->dsq, *dst_dsq;
 	struct scx_sched *sch;
-	struct rq *this_rq, *src_rq, *locked_rq;
+	struct rq *this_rq, *src_rq, *locked_rq, *tracked_rq;
 	bool dispatched = false;
 	bool in_balance;
 	unsigned long flags;
@@ -8824,6 +8837,18 @@ static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
 	this_rq = this_rq();
 	in_balance = this_rq->scx.flags & SCX_RQ_IN_BALANCE;
 
+	/*
+	 * When called from ops.dispatch(), @this_rq is recorded as the
+	 * currently locked rq. Clear the tracking while switching rq locks so
+	 * nested callbacks don't restore an rq which is no longer locked.
+	 */
+	tracked_rq = scx_locked_rq();
+	if (tracked_rq) {
+		WARN_ON_ONCE(!in_balance);
+		WARN_ON_ONCE(tracked_rq != this_rq);
+		update_locked_rq(NULL);
+	}
+
 	if (in_balance) {
 		if (this_rq != src_rq) {
 			raw_spin_rq_unlock(this_rq);
@@ -8864,6 +8889,8 @@ static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
 			raw_spin_rq_unlock(locked_rq);
 			raw_spin_rq_lock(this_rq);
 		}
+		if (tracked_rq)
+			update_locked_rq(tracked_rq);
 	} else {
 		raw_spin_rq_unlock_irqrestore(locked_rq, flags);
 	}
-- 
2.55.0