4 files changed, 34 insertions, 6 deletions

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index a8a73b8897bf..cde573d3f12e 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -2586,24 +2586,28 @@ static inline void schedule_debug(struct task_struct *prev)
 static inline struct task_struct *
 pick_next_task(struct rq *rq, struct task_struct *prev)
 {
-        const struct sched_class *class;
+        const struct sched_class *class = &fair_sched_class;
         struct task_struct *p;
 
         /*
          * Optimization: we know that if all tasks are in
          * the fair class we can call that function directly:
          */
-        if (likely(prev->sched_class == &fair_sched_class &&
+        if (likely(prev->sched_class == class &&
                    rq->nr_running == rq->cfs.h_nr_running)) {
                 p = fair_sched_class.pick_next_task(rq, prev);
-                if (likely(p))
+                if (likely(p && p != RETRY_TASK))
                         return p;
         }
 
+again:
         for_each_class(class) {
                 p = class->pick_next_task(rq, prev);
-                if (p)
+                if (p) {
+                        if (unlikely(p == RETRY_TASK))
+                                goto again;
                         return p;
+                }
         }
 
         BUG(); /* the idle class will always have a runnable task */
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index be4f7d9eaf03..16042b58a32f 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -4686,6 +4686,7 @@ pick_next_task_fair(struct rq *rq, struct task_struct *prev)
         struct cfs_rq *cfs_rq = &rq->cfs;
         struct sched_entity *se;
         struct task_struct *p;
+        int new_tasks;
 
 again:
 #ifdef CONFIG_FAIR_GROUP_SCHED
@@ -4784,7 +4785,17 @@ simple:
         return p;
 
 idle:
-        if (idle_balance(rq)) /* drops rq->lock */
+        /*
+         * Because idle_balance() 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.
+         */
+        new_tasks = idle_balance(rq);
+
+        if (rq->nr_running != rq->cfs.h_nr_running)
+                return RETRY_TASK;
+
+        if (new_tasks)
                 goto again;
 
         return NULL;
diff --git a/kernel/sched/rt.c b/kernel/sched/rt.c
index 4d4b386598aa..398b3f990823 100644
--- a/kernel/sched/rt.c
+++ b/kernel/sched/rt.c
@@ -1360,8 +1360,16 @@ pick_next_task_rt(struct rq *rq, struct task_struct *prev)
         struct task_struct *p;
         struct rt_rq *rt_rq = &rq->rt;
 
-        if (need_pull_rt_task(rq, prev))
+        if (need_pull_rt_task(rq, prev)) {
                 pull_rt_task(rq);
+                /*
+                 * pull_rt_task() can drop (and re-acquire) rq->lock; this
+                 * means a dl task can slip in, in which case we need to
+                 * re-start task selection.
+                 */
+                if (unlikely(rq->dl.dl_nr_running))
+                        return RETRY_TASK;
+        }
 
         if (!rt_rq->rt_nr_running)
                 return NULL;
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 046084ebb1fb..1929deb3f29d 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1091,6 +1091,8 @@ static const u32 prio_to_wmult[40] = {
 
 #define DEQUEUE_SLEEP           1
 
+#define RETRY_TASK              ((void *)-1UL)
+
 struct sched_class {
         const struct sched_class *next;
 
@@ -1105,6 +1107,9 @@ struct sched_class {
          * It is the responsibility of the pick_next_task() method that will
          * return the next task to call put_prev_task() on the @prev task or
          * something equivalent.
+         *
+         * May return RETRY_TASK when it finds a higher prio class has runnable
+         * tasks.
          */
         struct task_struct * (*pick_next_task) (struct rq *rq,
                                                 struct task_struct *prev);