2 files changed, 23 insertions, 0 deletions

diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index f808ddf2a868..3b0b75de1141 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -1514,6 +1514,21 @@ struct task_numa_env {
 static void task_numa_assign(struct task_numa_env *env,
                              struct task_struct *p, long imp)
 {
+        struct rq *rq = cpu_rq(env->dst_cpu);
+
+        /* Bail out if run-queue part of active NUMA balance. */
+        if (xchg(&rq->numa_migrate_on, 1))
+                return;
+
+        /*
+         * Clear previous best_cpu/rq numa-migrate flag, since task now
+         * found a better CPU to move/swap.
+         */
+        if (env->best_cpu != -1) {
+                rq = cpu_rq(env->best_cpu);
+                WRITE_ONCE(rq->numa_migrate_on, 0);
+        }
+
         if (env->best_task)
                 put_task_struct(env->best_task);
         if (p)
@@ -1569,6 +1584,9 @@ static void task_numa_compare(struct task_numa_env *env,
         long moveimp = imp;
         int dist = env->dist;
 
+        if (READ_ONCE(dst_rq->numa_migrate_on))
+                return;
+
         rcu_read_lock();
         cur = task_rcu_dereference(&dst_rq->curr);
         if (cur && ((cur->flags & PF_EXITING) || is_idle_task(cur)))
@@ -1710,6 +1728,7 @@ static int task_numa_migrate(struct task_struct *p)
                 .best_cpu = -1,
         };
         struct sched_domain *sd;
+        struct rq *best_rq;
         unsigned long taskweight, groupweight;
         int nid, ret, dist;
         long taskimp, groupimp;
@@ -1811,14 +1830,17 @@ static int task_numa_migrate(struct task_struct *p)
          */
         p->numa_scan_period = task_scan_start(p);
 
+        best_rq = cpu_rq(env.best_cpu);
         if (env.best_task == NULL) {
                 ret = migrate_task_to(p, env.best_cpu);
+                WRITE_ONCE(best_rq->numa_migrate_on, 0);
                 if (ret != 0)
                         trace_sched_stick_numa(p, env.src_cpu, env.best_cpu);
                 return ret;
         }
 
         ret = migrate_swap(p, env.best_task, env.best_cpu, env.src_cpu);
+        WRITE_ONCE(best_rq->numa_migrate_on, 0);
 
         if (ret != 0)
                 trace_sched_stick_numa(p, env.src_cpu, task_cpu(env.best_task));
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 4a2e8cae63c4..0b9161241bda 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -783,6 +783,7 @@ struct rq {
 #ifdef CONFIG_NUMA_BALANCING
         unsigned int            nr_numa_running;
         unsigned int            nr_preferred_running;
+        unsigned int            numa_migrate_on;
 #endif
         #define CPU_LOAD_IDX_MAX 5
         unsigned long           cpu_load[CPU_LOAD_IDX_MAX];