sched: Rework prev_balance() to avoid stale prev references

Historically, the prev value from __schedule() was the rq->curr.
This prev value is passed down through numerous functions, and
used in the class scheduler implementations. The fact that
prev was on_cpu until the end of __schedule(), meant it was
stable across the rq lock drops that the class->balance()
implementations often do.

However, with proxy-exec, the prev passed to functions called
by __schedule() is rq->donor, which may not be the same as
rq->curr and may not be on_cpu, this makes the prev value
potentially unstable across rq lock drops.

A recently found issue with proxy-exec, is when we begin doing
return migration from try_to_wake_up(), its possible we may be
waking up the rq->donor.  When we do this, we proxy_resched_idle()
to put_prev_set_next() setting the rq->donor to rq->idle, allowing
the rq->donor to be return migrated and allowed to run.

This however runs into trouble, as on another cpu we might be in
the middle of calling __schedule(). Conceptually the rq lock is
held for the majority of the time, but in calling prev_balance()
its possible the class->balance() handler call may briefly drop the rq lock.
This opens a window for try_to_wake_up() to wake and return migrate the
rq->donor before the class logic reacquires the rq lock.

Unfortunately prev_balance() pass in a prev argument, to which we pass
rq->donor. However this prev value can now become stale and incorrect across a
rq lock drop.

So, to correct this, rework the prev_balance() call so that it does not take a
"prev" argument.

Signed-off-by: John Stultz <jstultz@google.com>
Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
Link: https://patch.msgid.link/20260512025635.2840817-2-jstultz@google.com

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 3c8bfd697e2c1877d7434a99495be36d321a53cb..a9c9b89260cd44a5ef103522901cc5479e970d16 100644 (file)
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -5986,10 +5986,9 @@ static inline void schedule_debug(struct task_struct *prev, bool preempt)
        schedstat_inc(this_rq()->sched_count);
 }
 
-static void prev_balance(struct rq *rq, struct task_struct *prev,
-                        struct rq_flags *rf)
+static void prev_balance(struct rq *rq, struct rq_flags *rf)
 {
-       const struct sched_class *start_class = prev->sched_class;
+       const struct sched_class *start_class = rq->donor->sched_class;
        const struct sched_class *class;
 
        /*
@@ -6001,7 +6000,7 @@ static void prev_balance(struct rq *rq, struct task_struct *prev,
         * a runnable task of @class priority or higher.
         */
        for_active_class_range(class, start_class, &idle_sched_class) {
-               if (class->balance && class->balance(rq, prev, rf))
+               if (class->balance && class->balance(rq, rf))
                        break;
        }
 }
@@ -6010,7 +6009,7 @@ static void prev_balance(struct rq *rq, struct task_struct *prev,
  * Pick up the highest-prio task:
  */
 static inline struct task_struct *
-__pick_next_task(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
+__pick_next_task(struct rq *rq, struct rq_flags *rf)
        __must_hold(__rq_lockp(rq))
 {
        const struct sched_class *class;
@@ -6027,7 +6026,7 @@ __pick_next_task(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
         * higher scheduling class, because otherwise those lose the
         * opportunity to pull in more work from other CPUs.
         */
-       if (likely(!sched_class_above(prev->sched_class, &fair_sched_class) &&
+       if (likely(!sched_class_above(rq->donor->sched_class, &fair_sched_class) &&
                   rq->nr_running == rq->cfs.h_nr_queued)) {
 
                p = pick_task_fair(rq, rf);
@@ -6038,19 +6037,19 @@ __pick_next_task(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
                if (!p)
                        p = pick_task_idle(rq, rf);
 
-               put_prev_set_next_task(rq, prev, p);
+               put_prev_set_next_task(rq, rq->donor, p);
                return p;
        }
 
 restart:
-       prev_balance(rq, prev, rf);
+       prev_balance(rq, rf);
 
        for_each_active_class(class) {
                p = class->pick_task(rq, rf);
                if (unlikely(p == RETRY_TASK))
                        goto restart;
                if (p) {
-                       put_prev_set_next_task(rq, prev, p);
+                       put_prev_set_next_task(rq, rq->donor, p);
                        return p;
                }
        }
@@ -6102,7 +6101,7 @@ extern void task_vruntime_update(struct rq *rq, struct task_struct *p, bool in_f
 static void queue_core_balance(struct rq *rq);
 
 static struct task_struct *
-pick_next_task(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
+pick_next_task(struct rq *rq, struct rq_flags *rf)
        __must_hold(__rq_lockp(rq))
 {
        struct task_struct *next, *p, *max;
@@ -6115,7 +6114,7 @@ pick_next_task(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
        bool need_sync;
 
        if (!sched_core_enabled(rq))
-               return __pick_next_task(rq, prev, rf);
+               return __pick_next_task(rq, rf);
 
        cpu = cpu_of(rq);
 
@@ -6128,7 +6127,7 @@ pick_next_task(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
                 */
                rq->core_pick = NULL;
                rq->core_dl_server = NULL;
-               return __pick_next_task(rq, prev, rf);
+               return __pick_next_task(rq, rf);
        }
 
        /*
@@ -6152,7 +6151,7 @@ pick_next_task(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
                goto out_set_next;
        }
 
-       prev_balance(rq, prev, rf);
+       prev_balance(rq, rf);
 
        smt_mask = cpu_smt_mask(cpu);
        need_sync = !!rq->core->core_cookie;
@@ -6334,7 +6333,7 @@ restart_multi:
        }
 
 out_set_next:
-       put_prev_set_next_task(rq, prev, next);
+       put_prev_set_next_task(rq, rq->donor, next);
        if (rq->core->core_forceidle_count && next == rq->idle)
                queue_core_balance(rq);
 
@@ -6557,10 +6556,10 @@ static inline void sched_core_cpu_deactivate(unsigned int cpu) {}
 static inline void sched_core_cpu_dying(unsigned int cpu) {}
 
 static struct task_struct *
-pick_next_task(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
+pick_next_task(struct rq *rq, struct rq_flags *rf)
        __must_hold(__rq_lockp(rq))
 {
-       return __pick_next_task(rq, prev, rf);
+       return __pick_next_task(rq, rf);
 }
 
 #endif /* !CONFIG_SCHED_CORE */
@@ -7108,7 +7107,7 @@ static void __sched notrace __schedule(int sched_mode)
 
 pick_again:
        assert_balance_callbacks_empty(rq);
-       next = pick_next_task(rq, rq->donor, &rf);
+       next = pick_next_task(rq, &rf);
        rq->next_class = next->sched_class;
        if (sched_proxy_exec()) {
                struct task_struct *prev_donor = rq->donor;

diff --git a/kernel/sched/deadline.c b/kernel/sched/deadline.c
index f9e62ed08d775d3ce75139f6992d0f40737ada2a..6ef5a808e13e09c9fad6cb639467a3bb9a9da6b4 100644 (file)
--- a/kernel/sched/deadline.c
+++ b/kernel/sched/deadline.c
@@ -2698,8 +2698,14 @@ static void check_preempt_equal_dl(struct rq *rq, struct task_struct *p)
        resched_curr(rq);
 }
 
-static int balance_dl(struct rq *rq, struct task_struct *p, struct rq_flags *rf)
+static int balance_dl(struct rq *rq, struct rq_flags *rf)
 {
+       /*
+        * Note, rq->donor may change during rq lock drops,
+        * so don't re-use prev across lock drops
+        */
+       struct task_struct *p = rq->donor;
+
        if (!on_dl_rq(&p->dl) && need_pull_dl_task(rq, p)) {
                /*
                 * This is OK, because current is on_cpu, which avoids it being

diff --git a/kernel/sched/idle.c b/kernel/sched/idle.c
index a83be0c834ddb5ec8a236c525d8a5edd84b3175a..ff39120d723a91f6a9469d9875cb373bf467f9af 100644 (file)
--- a/kernel/sched/idle.c
+++ b/kernel/sched/idle.c
@@ -462,7 +462,7 @@ select_task_rq_idle(struct task_struct *p, int cpu, int flags)
 }
 
 static int
-balance_idle(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
+balance_idle(struct rq *rq, struct rq_flags *rf)
 {
        return WARN_ON_ONCE(1);
 }

diff --git a/kernel/sched/rt.c b/kernel/sched/rt.c
index e6ea728f519ee866918b414dc4d29e392f4d36dd..e474c31d8fe66b7f50b251b44c15afd16d245c83 100644 (file)
--- a/kernel/sched/rt.c
+++ b/kernel/sched/rt.c
@@ -1596,8 +1596,14 @@ static void check_preempt_equal_prio(struct rq *rq, struct task_struct *p)
        resched_curr(rq);
 }
 
-static int balance_rt(struct rq *rq, struct task_struct *p, struct rq_flags *rf)
+static int balance_rt(struct rq *rq, struct rq_flags *rf)
 {
+       /*
+        * Note, rq->donor may change during rq lock drops,
+        * so don't re-use p across lock drops
+        */
+       struct task_struct *p = rq->donor;
+
        if (!on_rt_rq(&p->rt) && need_pull_rt_task(rq, p)) {
                /*
                 * This is OK, because current is on_cpu, which avoids it being

diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 332ecf8930b46347afceb6f0c1584ceb10fbd4ff..ef715f2acbaa20d89750683367e69c989fc00c9d 100644 (file)
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -2587,7 +2587,7 @@ struct sched_class {
        /*
         * schedule/pick_next_task/prev_balance: rq->lock
         */
-       int (*balance)(struct rq *rq, struct task_struct *prev, struct rq_flags *rf);
+       int (*balance)(struct rq *rq, struct rq_flags *rf);
 
        /*
         * schedule/pick_next_task: rq->lock

diff --git a/kernel/sched/stop_task.c b/kernel/sched/stop_task.c
index f95798baddebbd9b2b721b3f828a69da3ea615f6..c909ca0d8c87c1055ad5ffe58ae9f044194ee358 100644 (file)
--- a/kernel/sched/stop_task.c
+++ b/kernel/sched/stop_task.c
@@ -16,7 +16,7 @@ select_task_rq_stop(struct task_struct *p, int cpu, int flags)
 }
 
 static int
-balance_stop(struct rq *rq, struct task_struct *prev, struct rq_flags *rf)
+balance_stop(struct rq *rq, struct rq_flags *rf)
 {
        return sched_stop_runnable(rq);
 }