sched: Collate affine_move_task() stoppers

The SCA_MIGRATE_ENABLE and task_running() cases are almost identical,
collapse them to avoid further duplication.

Fixes: 6d337eab041d ("sched: Fix migrate_disable() vs set_cpus_allowed_ptr()")
Cc: stable@kernel.org
Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
Signed-off-by: Ingo Molnar <mingo@kernel.org>
Reviewed-by: Valentin Schneider <valentin.schneider@arm.com>
Link: https://lkml.kernel.org/r/20210224131355.500108964@infradead.org

diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 088e8f492271b08b8d5941a97f2499d0416e5d52..84b657f056258fd2dfae749466edc775468a125b 100644 (file)
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -2239,30 +2239,23 @@ static int affine_move_task(struct rq *rq, struct task_struct *p, struct rq_flag
                return -EINVAL;
        }
 
-       if (flags & SCA_MIGRATE_ENABLE) {
-
-               refcount_inc(&pending->refs); /* pending->{arg,stop_work} */
-               p->migration_flags &= ~MDF_PUSH;
-               task_rq_unlock(rq, p, rf);
-
-               stop_one_cpu_nowait(cpu_of(rq), migration_cpu_stop,
-                                   &pending->arg, &pending->stop_work);
-
-               return 0;
-       }
-
        if (task_running(rq, p) || p->state == TASK_WAKING) {
                /*
-                * Lessen races (and headaches) by delegating
-                * is_migration_disabled(p) checks to the stopper, which will
-                * run on the same CPU as said p.
+                * MIGRATE_ENABLE gets here because 'p == current', but for
+                * anything else we cannot do is_migration_disabled(), punt
+                * and have the stopper function handle it all race-free.
                 */
+
                refcount_inc(&pending->refs); /* pending->{arg,stop_work} */
+               if (flags & SCA_MIGRATE_ENABLE)
+                       p->migration_flags &= ~MDF_PUSH;
                task_rq_unlock(rq, p, rf);
 
                stop_one_cpu_nowait(cpu_of(rq), migration_cpu_stop,
                                    &pending->arg, &pending->stop_work);
 
+               if (flags & SCA_MIGRATE_ENABLE)
+                       return 0;
        } else {
 
                if (!is_migration_disabled(p)) {