sched_ext: Check remote rq eligibility under task's rq lock

task_can_run_on_remote_rq() operates under the assumption that
p->migration_disabled is stable, i.e. if the kernel observed
is_migration_disabled(p) == true, then the BPF scheduler must have also
been able to see this when dispatching the task, and it's the BPF
scheduler's fault that it tried to dispatch a task with migration
disabled to a CPU other than the task's current CPU.

This assumption does not always hold. It's possible that the BPF
scheduler saw is_migration_disabled(p) == false, while the kernel
observes is_migration_disabled(p) == true in dispatch_to_local_dsq()
-> task_can_run_on_remote_rq().

The crucial thing here is that with CONFIG_PREEMPT_RCU, migration is
disabled while a task is executing a BPF program. So, if there's a
situation where the BPF scheduler checks a task while it's not executing
a BPF program, while the kernel checks it while it is executing one,
the BPF scheduler will be killed through no fault of its own.

Consider the following scenario:

1. SCX task @p is executing on CPU A and CPU A gets preempted by a
   higher-priority scheduling class. On entry to __schedule(),
   p->migration_disabled == 0.

2. In put_prev_task_scx() @p is enqueued on the BPF scheduler's internal
   data structures, making it available for other CPUs to dispatch.

3. CPU B enters ops.dispatch(), pops @p from the BPF scheduler's data
   structures, checks is_migration_disabled(p) which returns false,
   and dispatches @p to CPU B's local DSQ.

4. On CPU A, @p hasn't been switched out yet. Execution reaches
   trace_sched_switch() which enters a BPF program, as the BPF scheduler
   hooks into the sched_switch tracepoint to detect idle->fair
   transitions. On entry into the BPF program, @p disables migration.

5. CPU B enters finish_dispatch() -> dispatch_to_local_dsq() ->
   task_can_run_on_remote_rq() which observes
   is_migration_disabled(p) == true, triggering scx_error().
   This all happens while holding CPU B's rq lock, so it's not
   synchronized with @p switching out.

This patch fixes this by moving the call to task_can_run_on_remote_rq()
after @p's rq lock is acquired in dispatch_to_local_dsq(). This way, we
synchronize with @p switching out, since @p holds its rq lock all
the way until it's switched out. Thus, any BPF programs that are called
between put_prev_task_scx() and the end of the context switch are
guaranteed to have finished and cannot influence p->migration_disabled.

Also add a lockdep assertion in task_can_run_on_remote_rq() which
ensures the task rq lock is held if enforce == true.

Signed-off-by: Kuba Piecuch <jpiecuch@google.com>
Reviewed-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 5efc97437c0a0e3776fb8c08dd4450dbdd443d0b..5418e6357249d7819b4a6473954d63a63282d7ed 100644 (file)
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -2307,6 +2307,7 @@ static void move_remote_task_to_local_dsq(struct task_struct *p, u64 enq_flags,
  *   no to the BPF scheduler initiated migrations while offline.
  *
  * The caller must ensure that @p and @rq are on different CPUs.
+ * If enforce == true, caller must hold @p's rq lock.
  */
 static bool task_can_run_on_remote_rq(struct scx_sched *sch,
                                      struct task_struct *p, struct rq *rq,
@@ -2314,6 +2315,14 @@ static bool task_can_run_on_remote_rq(struct scx_sched *sch,
 {
        s32 cpu = cpu_of(rq);
 
+       /*
+        * To prevent races with @p still running on its old CPU while switching
+        * out, make sure we're holding @p's rq lock so as not to risk
+        * erroneously killing the BPF scheduler.
+        */
+       if (enforce)
+               lockdep_assert_rq_held(task_rq(p));
+
        WARN_ON_ONCE(task_cpu(p) == cpu);
 
        /*
@@ -2581,13 +2590,6 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
                return;
        }
 
-       if (src_rq != dst_rq &&
-           unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq, true))) {
-               dispatch_enqueue(sch, rq, find_global_dsq(sch, task_cpu(p)), p,
-                                enq_flags | SCX_ENQ_CLEAR_OPSS | SCX_ENQ_GDSQ_FALLBACK);
-               return;
-       }
-
        /*
         * @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
@@ -2614,6 +2616,7 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
        /* task_rq couldn't have changed if we're still the holding cpu */
        if (likely(p->scx.holding_cpu == raw_smp_processor_id()) &&
            !WARN_ON_ONCE(src_rq != task_rq(p))) {
+               bool fallback = false;
                /*
                 * If @p is staying on the same rq, there's no need to go
                 * through the full deactivate/activate cycle. Optimize by
@@ -2623,6 +2626,11 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
                        p->scx.holding_cpu = -1;
                        dispatch_enqueue(sch, dst_rq, &dst_rq->scx.local_dsq, p,
                                         enq_flags);
+               } else if (unlikely(!task_can_run_on_remote_rq(sch, p, dst_rq, true))) {
+                       p->scx.holding_cpu = -1;
+                       fallback = true;
+                       dispatch_enqueue(sch, src_rq, find_global_dsq(sch, task_cpu(p)),
+                                        p, enq_flags | SCX_ENQ_GDSQ_FALLBACK);
                } else {
                        move_remote_task_to_local_dsq(p, enq_flags,
                                                      src_rq, dst_rq);
@@ -2631,7 +2639,7 @@ static void dispatch_to_local_dsq(struct scx_sched *sch, struct rq *rq,
                }
 
                /* if the destination CPU is idle, wake it up */
-               if (sched_class_above(p->sched_class, dst_rq->curr->sched_class))
+               if (!fallback && sched_class_above(p->sched_class, dst_rq->curr->sched_class))
                        resched_curr(dst_rq);
        }
 

diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 145272cb4d8a711a8634bbbad7cd3cc34042dfb2..673059fa9d728b12333c96453e58ef450b902ed5 100644 (file)
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -1469,21 +1469,24 @@ static const char *scx_enable_state_str[] = {
  * The sched_ext core uses a "lock dancing" protocol coordinated by
  * p->scx.holding_cpu. When moving a task to a different rq:
  *
- *   1. Verify task can be moved (CPU affinity, migration_disabled, etc.)
- *   2. Set p->scx.holding_cpu to the current CPU
- *   3. Set task state to %SCX_OPSS_NONE; dequeue waits while DISPATCHING
+ *   1. Set p->scx.holding_cpu to the current CPU
+ *   2. Set task state to %SCX_OPSS_NONE; dequeue waits while DISPATCHING
  *      is set, so clearing DISPATCHING first prevents the circular wait
  *      (safe to lock the rq we need)
- *   4. Unlock the current CPU's rq
- *   5. Lock src_rq (where the task currently lives)
- *   6. Verify p->scx.holding_cpu == current CPU, if not, dequeue won the
+ *   3. Unlock the current CPU's rq
+ *   4. Lock src_rq (where the task currently lives)
+ *   5. Verify p->scx.holding_cpu == current CPU, if not, dequeue won the
  *      race (dequeue clears holding_cpu to -1 when it takes the task), in
  *      this case migration is aborted
- *   7. If src_rq == dst_rq: clear holding_cpu and enqueue directly
+ *   6. If src_rq == dst_rq: clear holding_cpu and enqueue directly
  *      into dst_rq's local DSQ (no lock swap needed)
- *   8. Otherwise: call move_remote_task_to_local_dsq(), which releases
- *      src_rq, locks dst_rq, and performs the deactivate/activate
- *      migration cycle (dst_rq is held on return)
+ *   7. Otherwise, verify under src_rq lock that the task can be moved to dst_rq
+ *      (CPU affinity, migration_disabled, etc.). If not, clear holding_cpu,
+ *      leave the task on src_rq, and enqueue it on the fallback DSQ.
+ *   8. Otherwise (i.e. if the task can be moved to dst_rq), call
+ *      move_remote_task_to_local_dsq(), which releases src_rq, locks dst_rq,
+ *      and performs the deactivate/activate migration cycle
+ *      (dst_rq is held on return)
  *   9. Unlock dst_rq and re-lock the current CPU's rq to restore
  *      the lock state expected by the caller
  *