[PATCH sched_ext/for-7.2-fixes] 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
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.

Fixes: 7fb39e4eb4c3 ("sched_ext: Save and restore scx_locked_rq across SCX_CALL_OP")
Cc: [email protected] # 7.1+
Signed-off-by: Andrea Righi <[email protected]>
---
 kernel/sched/ext/ext.c | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index b6a635ba269cc..b8a61f7803d5c 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);
 }
 
 /**
-- 
2.55.0