3 files changed, 41 insertions, 35 deletions

diff --git a/include/linux/kthread.h b/include/linux/kthread.h
index c1961761311d..2803264c512f 100644
--- a/include/linux/kthread.h
+++ b/include/linux/kthread.h
@@ -62,6 +62,7 @@ void *kthread_probe_data(struct task_struct *k);
 int kthread_park(struct task_struct *k);
 void kthread_unpark(struct task_struct *k);
 void kthread_parkme(void);
+void kthread_park_complete(struct task_struct *k);
 
 int kthreadd(void *unused);
 extern struct task_struct *kthreadd_task;
diff --git a/kernel/kthread.c b/kernel/kthread.c
index cbee858e5815..2017a39ab490 100644
--- a/kernel/kthread.c
+++ b/kernel/kthread.c
@@ -55,7 +55,6 @@ enum KTHREAD_BITS {
         KTHREAD_IS_PER_CPU = 0,
         KTHREAD_SHOULD_STOP,
         KTHREAD_SHOULD_PARK,
-        KTHREAD_IS_PARKED,
 };
 
 static inline void set_kthread_struct(void *kthread)
@@ -181,11 +180,8 @@ static void __kthread_parkme(struct kthread *self)
                 set_current_state(TASK_PARKED);
                 if (!test_bit(KTHREAD_SHOULD_PARK, &self->flags))
                         break;
-                if (!test_and_set_bit(KTHREAD_IS_PARKED, &self->flags))
-                        complete(&self->parked);
                 schedule();
         }
-        clear_bit(KTHREAD_IS_PARKED, &self->flags);
         __set_current_state(TASK_RUNNING);
 }
 
@@ -195,6 +191,11 @@ void kthread_parkme(void)
 }
 EXPORT_SYMBOL_GPL(kthread_parkme);
 
+void kthread_park_complete(struct task_struct *k)
+{
+        complete(&to_kthread(k)->parked);
+}
+
 static int kthread(void *_create)
 {
         /* Copy data: it's on kthread's stack */
@@ -451,22 +452,15 @@ void kthread_unpark(struct task_struct *k)
 {
         struct kthread *kthread = to_kthread(k);
 
-        clear_bit(KTHREAD_SHOULD_PARK, &kthread->flags);
         /*
-         * We clear the IS_PARKED bit here as we don't wait
-         * until the task has left the park code. So if we'd
-         * park before that happens we'd see the IS_PARKED bit
-         * which might be about to be cleared.
+         * Newly created kthread was parked when the CPU was offline.
+         * The binding was lost and we need to set it again.
          */
-        if (test_and_clear_bit(KTHREAD_IS_PARKED, &kthread->flags)) {
-                /*
-                 * Newly created kthread was parked when the CPU was offline.
-                 * The binding was lost and we need to set it again.
-                 */
-                if (test_bit(KTHREAD_IS_PER_CPU, &kthread->flags))
-                        __kthread_bind(k, kthread->cpu, TASK_PARKED);
-                wake_up_state(k, TASK_PARKED);
-        }
+        if (test_bit(KTHREAD_IS_PER_CPU, &kthread->flags))
+                __kthread_bind(k, kthread->cpu, TASK_PARKED);
+
+        clear_bit(KTHREAD_SHOULD_PARK, &kthread->flags);
+        wake_up_state(k, TASK_PARKED);
 }
 EXPORT_SYMBOL_GPL(kthread_unpark);
 
@@ -489,12 +483,13 @@ int kthread_park(struct task_struct *k)
         if (WARN_ON(k->flags & PF_EXITING))
                 return -ENOSYS;
 
-        if (!test_bit(KTHREAD_IS_PARKED, &kthread->flags)) {
-                set_bit(KTHREAD_SHOULD_PARK, &kthread->flags);
-                if (k != current) {
-                        wake_up_process(k);
-                        wait_for_completion(&kthread->parked);
-                }
+        if (WARN_ON_ONCE(test_bit(KTHREAD_SHOULD_PARK, &kthread->flags)))
+                return -EBUSY;
+
+        set_bit(KTHREAD_SHOULD_PARK, &kthread->flags);
+        if (k != current) {
+                wake_up_process(k);
+                wait_for_completion(&kthread->parked);
         }
 
         return 0;
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 5e10aaeebfcc..7ad60e00a6a8 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -7,6 +7,8 @@
  */
 #include "sched.h"
 
+#include <linux/kthread.h>
+
 #include <asm/switch_to.h>
 #include <asm/tlb.h>
 
@@ -2718,20 +2720,28 @@ static struct rq *finish_task_switch(struct task_struct *prev)
                 membarrier_mm_sync_core_before_usermode(mm);
                 mmdrop(mm);
         }
-        if (unlikely(prev_state == TASK_DEAD)) {
-                if (prev->sched_class->task_dead)
-                        prev->sched_class->task_dead(prev);
+        if (unlikely(prev_state & (TASK_DEAD|TASK_PARKED))) {
+                switch (prev_state) {
+                case TASK_DEAD:
+                        if (prev->sched_class->task_dead)
+                                prev->sched_class->task_dead(prev);
 
-                /*
-                 * Remove function-return probe instances associated with this
-                 * task and put them back on the free list.
-                 */
-                kprobe_flush_task(prev);
+                        /*
+                         * Remove function-return probe instances associated with this
+                         * task and put them back on the free list.
+                         */
+                        kprobe_flush_task(prev);
 
-                /* Task is done with its stack. */
-                put_task_stack(prev);
+                        /* Task is done with its stack. */
+                        put_task_stack(prev);
 
-                put_task_struct(prev);
+                        put_task_struct(prev);
+                        break;
+
+                case TASK_PARKED:
+                        kthread_park_complete(prev);
+                        break;
+                }
         }
 
         tick_nohz_task_switch();