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

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

[Andrea Righi]

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

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

[Tejun Heo]

Hello, Andrea.

Instead of clearing and restoring the tracking at the function boundaries,
can you wrap the lock switch itself so scx_locked_rq() follows it? Something
like:

	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);
	}

Then use it at the unlock/lock pairs in dispatch_to_local_dsq(),
move_remote_task_to_local_dsq(), and scx_dsq_move() (the in_balance ones;
the !in_balance fresh lock stays). That keeps scx_locked_rq() naming the
actually-held rq the whole time instead of going NULL across the dance, and
drops the tracked_rq bookkeeping, the WARN_ON_ONCE()s, and the
in_balance-vs-scx_locked_rq() coupling.

The == from guard makes it a no-op for the consume path (there it's this_rq,
not the rq being released), so that stays as-is for now - it's harmless
today because the deactivate runs under the migration guards. Would be nice
to bring consume under the same helper on for-7.3 so the tracking is
faithful everywhere.

Thanks.

--
tejun

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

[Andrea Righi]

Hi Tejun,

On Tue, Jul 07, 2026 at 12:36:34PM -1000, Tejun Heo wrote:
> Hello, Andrea.
> 
> Instead of clearing and restoring the tracking at the function boundaries,
> can you wrap the lock switch itself so scx_locked_rq() follows it? Something
> like:
> 
> 	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);
> 	}
> 
> Then use it at the unlock/lock pairs in dispatch_to_local_dsq(),
> move_remote_task_to_local_dsq(), and scx_dsq_move() (the in_balance ones;
> the !in_balance fresh lock stays). That keeps scx_locked_rq() naming the
> actually-held rq the whole time instead of going NULL across the dance, and
> drops the tracked_rq bookkeeping, the WARN_ON_ONCE()s, and the
> in_balance-vs-scx_locked_rq() coupling.

Makes sense, I'll send a v3 shortly with these changes for for-7.2-fixes.

> 
> The == from guard makes it a no-op for the consume path (there it's this_rq,
> not the rq being released), so that stays as-is for now - it's harmless
> today because the deactivate runs under the migration guards. Would be nice
> to bring consume under the same helper on for-7.3 so the tracking is
> faithful everywhere.

Ack, I'll follow up with a separate for-7.3 patch to keep the rq tracking
accurate in the consume path as well.

Thanks,
-Andrea