sched/fair: Add newidle balance to pick_task_fair()

With commit 50653216e4ff ("sched: Add support to pick functions to
take rf") removing the balance callback, the pick_task() callback is
in charge of newidle balancing.

This means pick_task_fair() should do so too. This hasn't been a
problem in practise because pick_next_task_fair() is used. However,
since we'll be removing that one shortly, make sure pick_next_task()
is up to scratch.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
Reviewed-by: Vincent Guittot <vincent.guittot@linaro.org>
Link: https://patch.msgid.link/20260511120627.944705718@infradead.org

diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 8e858ca6bcd0e4c714d7c4831974300b0c799e15..5f48af700fd44880fbf2bcc67df7e34215c783a1 100644 (file)
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -9863,16 +9863,18 @@ preempt:
 }
 
 static struct task_struct *pick_task_fair(struct rq *rq, struct rq_flags *rf)
+       __must_hold(__rq_lockp(rq))
 {
        struct sched_entity *se;
        struct cfs_rq *cfs_rq;
        struct task_struct *p;
        bool throttled;
+       int new_tasks;
 
 again:
        cfs_rq = &rq->cfs;
        if (!cfs_rq->nr_queued)
-               return NULL;
+               goto idle;
 
        throttled = false;
 
@@ -9893,6 +9895,14 @@ again:
        if (unlikely(throttled))
                task_throttle_setup_work(p);
        return p;
+
+idle:
+       new_tasks = sched_balance_newidle(rq, rf);
+       if (new_tasks < 0)
+               return RETRY_TASK;
+       if (new_tasks > 0)
+               goto again;
+       return NULL;
 }
 
 static void __set_next_task_fair(struct rq *rq, struct task_struct *p, bool first);
@@ -9904,12 +9914,12 @@ pick_next_task_fair(struct rq *rq, struct task_struct *prev, struct rq_flags *rf
 {
        struct sched_entity *se;
        struct task_struct *p;
-       int new_tasks;
 
-again:
        p = pick_task_fair(rq, rf);
+       if (unlikely(p == RETRY_TASK))
+               return p;
        if (!p)
-               goto idle;
+               return p;
        se = &p->se;
 
 #ifdef CONFIG_FAIR_GROUP_SCHED
@@ -9959,29 +9969,11 @@ simple:
 #endif /* CONFIG_FAIR_GROUP_SCHED */
        put_prev_set_next_task(rq, prev, p);
        return p;
-
-idle:
-       if (rf) {
-               new_tasks = sched_balance_newidle(rq, rf);
-
-               /*
-                * Because sched_balance_newidle() releases (and re-acquires)
-                * rq->lock, it is possible for any higher priority task to
-                * appear. In that case we must re-start the pick_next_entity()
-                * loop.
-                */
-               if (new_tasks < 0)
-                       return RETRY_TASK;
-
-               if (new_tasks > 0)
-                       goto again;
-       }
-
-       return NULL;
 }
 
 static struct task_struct *
 fair_server_pick_task(struct sched_dl_entity *dl_se, struct rq_flags *rf)
+       __must_hold(__rq_lockp(dl_se->rq))
 {
        return pick_task_fair(dl_se->rq, rf);
 }